def test_runner_api_transformation_with_subscription( self, unused_mock_pubsub): source = _PubSubSource( topic=None, subscription='projects/fakeprj/subscriptions/a_subscription', id_label='a_label', timestamp_attribute='b_label', with_attributes=True) transform = Read(source) context = pipeline_context.PipelineContext() proto_transform_spec = transform.to_runner_api(context) self.assertEqual(common_urns.composites.PUBSUB_READ.urn, proto_transform_spec.urn) pubsub_read_payload = (proto_utils.parse_Bytes( proto_transform_spec.payload, beam_runner_api_pb2.PubSubReadPayload)) self.assertEqual('projects/fakeprj/subscriptions/a_subscription', pubsub_read_payload.subscription) self.assertEqual('a_label', pubsub_read_payload.id_attribute) self.assertEqual('b_label', pubsub_read_payload.timestamp_attribute) self.assertEqual('', pubsub_read_payload.topic) self.assertTrue(pubsub_read_payload.with_attributes) proto_transform = beam_runner_api_pb2.PTransform( unique_name="dummy_label", spec=proto_transform_spec) transform_from_proto = Read.from_runner_api_parameter( proto_transform, pubsub_read_payload, None) self.assertTrue(isinstance(transform_from_proto, Read)) self.assertTrue(isinstance(transform_from_proto.source, _PubSubSource)) self.assertTrue(transform_from_proto.source.with_attributes) self.assertEqual('projects/fakeprj/subscriptions/a_subscription', transform_from_proto.source.full_subscription)
def test_display_data_no_subscription(self): source = _PubSubSource('projects/fakeprj/topics/a_topic') dd = DisplayData.create_from(source) expected_items = [ DisplayDataItemMatcher('topic', 'projects/fakeprj/topics/a_topic')] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
def test_display_data_no_subscription(self): source = _PubSubSource('projects/fakeprj/topics/a_topic') dd = DisplayData.create_from(source) expected_items = [ DisplayDataItemMatcher('topic', 'projects/fakeprj/topics/a_topic')] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
def test_display_data_topic(self): source = _PubSubSource('projects/fakeprj/topics/a_topic', None, 'a_label') dd = DisplayData.create_from(source) expected_items = [ DisplayDataItemMatcher('topic', 'projects/fakeprj/topics/a_topic'), DisplayDataItemMatcher('id_label', 'a_label'), DisplayDataItemMatcher('with_attributes', False), ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
def test_display_data_subscription(self): source = _PubSubSource( None, 'projects/fakeprj/subscriptions/a_subscription', 'a_label') dd = DisplayData.create_from(source) expected_items = [ DisplayDataItemMatcher( 'subscription', 'projects/fakeprj/subscriptions/a_subscription'), DisplayDataItemMatcher('id_label', 'a_label') ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
def test_display_data_subscription(self): source = _PubSubSource( None, 'projects/fakeprj/subscriptions/a_subscription', 'a_label') dd = DisplayData.create_from(source) expected_items = [ DisplayDataItemMatcher( 'subscription', 'projects/fakeprj/subscriptions/a_subscription'), DisplayDataItemMatcher('id_label', 'a_label')] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
def test_display_data_topic(self): source = _PubSubSource( 'projects/fakeprj/topics/a_topic', None, 'a_label') dd = DisplayData.create_from(source) expected_items = [ DisplayDataItemMatcher( 'topic', 'projects/fakeprj/topics/a_topic'), DisplayDataItemMatcher('id_label', 'a_label'), DisplayDataItemMatcher('with_attributes', False), ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
def test_source_to_cache_changed_when_source_is_altered(self, cell): with cell: # Cell 1 pipeline = _build_an_empty_stream_pipeline() transform = beam.io.ReadFromPubSub(subscription=_FOO_PUBSUB_SUB) read_foo = pipeline | 'Read' >> transform ib.watch({'read_foo': read_foo}) # Sets the signature for current pipeline state. ie.current_env().set_cached_source_signature( pipeline, bcj.extract_source_to_cache_signature(pipeline)) with cell: # Cell 2 from apache_beam.io.gcp.pubsub import _PubSubSource # Alter the transform. transform._source = _PubSubSource(subscription=_BAR_PUBSUB_SUB) self.assertTrue(bcj.is_source_to_cache_changed(pipeline))
def test_runner_api_transformation_properties_none(self, unused_mock_pubsub): # Confirming that properties stay None after a runner API transformation. source = _PubSubSource(topic='projects/fakeprj/topics/a_topic', with_attributes=True) transform = Read(source) context = pipeline_context.PipelineContext() proto_transform_spec = transform.to_runner_api(context) self.assertEqual(common_urns.composites.PUBSUB_READ.urn, proto_transform_spec.urn) pubsub_read_payload = (proto_utils.parse_Bytes( proto_transform_spec.payload, beam_runner_api_pb2.PubSubReadPayload)) proto_transform = beam_runner_api_pb2.PTransform( unique_name="dummy_label", spec=proto_transform_spec) transform_from_proto = Read.from_runner_api_parameter( proto_transform, pubsub_read_payload, None) self.assertIsNone(transform_from_proto.source.full_subscription) self.assertIsNone(transform_from_proto.source.id_label) self.assertIsNone(transform_from_proto.source.timestamp_attribute)