def test_read_messages_timestamp_attribute_rfc3339_success(
mocker,
patch_sub_client,
patch_msg_manager,
):
exp_entity_id = "entity_id"
kmsg = klio_pb2.KlioMessage()
kmsg.data.element = bytes(exp_entity_id, "utf-8")
data = kmsg.SerializeToString()
attributes = {"time": "2018-03-12T13:37:01.234567Z"}
publish_time_secs = 1337000000
publish_time_nanos = 133700000
ack_id = "ack_id"
pull_response = beam_test_utils.create_pull_response([
beam_test_utils.PullResponseMessage(data, attributes,
publish_time_secs,
publish_time_nanos, ack_id)
])
pmsg = b_pubsub.PubsubMessage(data, attributes)
expected_elements = [
beam_testing_util.TestWindowedValue(
pmsg,
beam_utils.timestamp.Timestamp.from_rfc3339(attributes["time"]),
[beam_transforms.window.GlobalWindow()],
),
]
patch_sub_client.pull.return_value = pull_response
options = pipeline_options.PipelineOptions([])
options.view_as(pipeline_options.StandardOptions).streaming = True
with beam_test_pipeline.TestPipeline(options=options) as p:
pcoll = p | b_pubsub.ReadFromPubSub(
"projects/fakeprj/topics/a_topic",
None,
None,
with_attributes=True,
timestamp_attribute="time",
)
# Check original functionality that was kept the same
beam_testing_util.assert_that(
pcoll,
beam_testing_util.equal_to(expected_elements),
reify_windows=True,
)
# Check overridden functionality:
# 1. Check that auto-acking is skipped
patch_sub_client.acknowledge.assert_not_called()
# 2. Check that MessageManager daemon threads were started
patch_msg_manager.assert_called_once_with(
patch_sub_client.subscription_path())
# 3. Check that messages were added to the MessageManager
patch_msg_manager.return_value.add.assert_called_once_with(ack_id, pmsg)
# 4. Check that one message is handled at a time, instead of the
# original 10
patch_sub_client.pull.assert_called_once_with(mocker.ANY,
max_messages=1,
return_immediately=True)
patch_sub_client.api.transport.channel.close.assert_called_once_with()
def expand(self, pvalue):
if self.with_attributes:
pcoll = pvalue | 'ToProto' >> Map(
pubsub.WriteToPubSub.to_proto_str)
else:
pcoll = pvalue | 'ToProto' >> Map(
lambda x: pubsub.PubsubMessage(x, {})._to_proto_str())
pcoll.element_type = bytes
return pcoll.apply(
ExternalTransform(self.URN,
NamedTupleBasedPayloadBuilder(self.params),
self.expansion_service))
def test_convert_raw_pubsub_message(mocker, monkeypatch, msg_manager):
mock_event = mocker.Mock()
monkeypatch.setattr(pmm.threading, "Event", mock_event)
exp_message = pmm.PubSubKlioMessage("ack_id1", "kmsg_id1")
kmsg = klio_pb2.KlioMessage()
kmsg.data.element = b"kmsg_id1"
kmsg_bytes = kmsg.SerializeToString()
pmsg = beam_pubsub.PubsubMessage(data=kmsg_bytes, attributes={})
act_message = msg_manager._convert_raw_pubsub_message("ack_id1", pmsg)
# comparing class attributes (via __dict__) since we'd need to implement
# __eq__ on the PubSubKlioMessage class, but doing so would make it un-
# hashable. Which can be addressed, but this just seems easier for now.
assert _compare_objects_dicts(exp_message, act_message)
def _get_pubsub_message(klio_element):
kmsg = klio_pb2.KlioMessage()
kmsg.data.element = bytes(klio_element.encode("utf-8"))
kmsg_bytes = kmsg.SerializeToString()
pmsg = beam_pubsub.PubsubMessage(data=kmsg_bytes, attributes={})
return pmsg