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