def test_runner_api_transformation_properties_none(self,
unused_mock_pubsub):
# Confirming that properties stay None after a runner API transformation.
sink = _PubSubSink(
topic='projects/fakeprj/topics/a_topic',
id_label=None,
with_attributes=True,
# We expect encoded PubSub write transform to always return attributes.
timestamp_attribute=None)
transform = Write(sink)
context = pipeline_context.PipelineContext()
proto_transform_spec = transform.to_runner_api(context)
self.assertEqual(common_urns.composites.PUBSUB_WRITE.urn,
proto_transform_spec.urn)
pubsub_write_payload = (proto_utils.parse_Bytes(
proto_transform_spec.payload,
beam_runner_api_pb2.PubSubWritePayload))
proto_transform = beam_runner_api_pb2.PTransform(
unique_name="dummy_label", spec=proto_transform_spec)
transform_from_proto = Write.from_runner_api_parameter(
proto_transform, pubsub_write_payload, None)
self.assertTrue(isinstance(transform_from_proto, Write))
self.assertTrue(isinstance(transform_from_proto.sink, _PubSubSink))
self.assertTrue(transform_from_proto.sink.with_attributes)
self.assertIsNone(transform_from_proto.sink.id_label)
self.assertIsNone(transform_from_proto.sink.timestamp_attribute)
def test_display_data(self):
sink = _PubSubSink('projects/fakeprj/topics/a_topic',
id_label='id', with_attributes=False,
timestamp_attribute='time')
dd = DisplayData.create_from(sink)
expected_items = [
DisplayDataItemMatcher('topic', 'projects/fakeprj/topics/a_topic'),
DisplayDataItemMatcher('id_label', 'id'),
DisplayDataItemMatcher('with_attributes', False),
DisplayDataItemMatcher('timestamp_attribute', 'time'),
]
hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
def test_display_data(self):
sink = _PubSubSink('projects/fakeprj/topics/a_topic',
id_label='id', with_attributes=False,
timestamp_attribute='time')
dd = DisplayData.create_from(sink)
expected_items = [
DisplayDataItemMatcher('topic', 'projects/fakeprj/topics/a_topic'),
DisplayDataItemMatcher('id_label', 'id'),
DisplayDataItemMatcher('with_attributes', False),
DisplayDataItemMatcher('timestamp_attribute', 'time'),
]
hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
