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)
