def test_write_messages_unsupported_features(self, mock_pubsub):
data = b'data'
attributes = {'key': 'value'}
payloads = [PubsubMessage(data, attributes)]
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
with self.assertRaisesRegex(NotImplementedError,
r'id_label is not supported'):
with TestPipeline(options=options) as p:
_ = (
p
| Create(payloads)
| WriteToPubSub(
'projects/fakeprj/topics/a_topic', id_label='a_label'))
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
with self.assertRaisesRegex(NotImplementedError,
r'timestamp_attribute is not supported'):
with TestPipeline(options=options) as p:
_ = (
p
| Create(payloads)
| WriteToPubSub(
'projects/fakeprj/topics/a_topic',
timestamp_attribute='timestamp'))
def test_write_messages_unsupported_features(self, mock_pubsub):
data = 'data'
attributes = {'key': 'value'}
payloads = [PubsubMessage(data, attributes)]
expected_payloads = [[data, attributes]]
mock_pubsub.Client = functools.partial(FakePubsubClient,
messages_write=expected_payloads)
p = TestPipeline()
p.options.view_as(StandardOptions).streaming = True
_ = (p
| Create(payloads)
| WriteToPubSub('projects/fakeprj/topics/a_topic',
id_label='a_label'))
with self.assertRaisesRegexp(NotImplementedError,
r'id_label is not supported'):
p.run()
p = TestPipeline()
p.options.view_as(StandardOptions).streaming = True
_ = (p
| Create(payloads)
| WriteToPubSub('projects/fakeprj/topics/a_topic',
timestamp_attribute='timestamp'))
with self.assertRaisesRegexp(NotImplementedError,
r'timestamp_attribute is not supported'):
p.run()
def test_expand(self):
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
p = TestPipeline(options=options)
pcoll = (
p
| ReadFromPubSub('projects/fakeprj/topics/baz')
| WriteToPubSub(
'projects/fakeprj/topics/a_topic', with_attributes=True)
| beam.Map(lambda x: x))
# Apply the necessary PTransformOverrides.
overrides = _get_transform_overrides(options)
p.replace_all(overrides)
# Note that the direct output of ReadFromPubSub will be replaced
# by a PTransformOverride, so we use a no-op Map.
write_transform = pcoll.producer.inputs[0].producer.transform
# Ensure that the properties passed through correctly
self.assertEqual('a_topic', write_transform.dofn.short_topic_name)
self.assertEqual(True, write_transform.dofn.with_attributes)
# TODO(BEAM-4275): These properties aren't supported yet in direct runner.
self.assertEqual(None, write_transform.dofn.id_label)
self.assertEqual(None, write_transform.dofn.timestamp_attribute)
def test_write_messages_success(self, mock_pubsub):
data = 'data'
payloads = [data]
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
with TestPipeline(options=options) as p:
_ = (p
| Create(payloads)
| WriteToPubSub('projects/fakeprj/topics/a_topic',
with_attributes=False))
mock_pubsub.return_value.publish.assert_has_calls(
[mock.call(mock.ANY, data)])
def test_display_data(self):
sink = WriteToPubSub('projects/fakeprj/topics/a_topic',
id_label='id',
timestamp_attribute='time')
dd = DisplayData.create_from(sink)
expected_items = [
DisplayDataItemMatcher('topic', 'projects/fakeprj/topics/a_topic'),
DisplayDataItemMatcher('id_label', 'id'),
DisplayDataItemMatcher('with_attributes', True),
DisplayDataItemMatcher('timestamp_attribute', 'time'),
]
hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
def test_write_messages_with_attributes_error(self, mock_pubsub):
data = 'data'
# Sending raw data when WriteToPubSub expects a PubsubMessage object.
payloads = [data]
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
with self.assertRaisesRegex(AttributeError,
r'str.*has no attribute.*data'):
with TestPipeline(options=options) as p:
_ = (p
| Create(payloads)
| WriteToPubSub('projects/fakeprj/topics/a_topic',
with_attributes=True))
def test_write_messages_with_attributes_success(self, mock_pubsub):
data = 'data'
attributes = {'key': 'value'}
payloads = [PubsubMessage(data, attributes)]
p = TestPipeline()
p.options.view_as(StandardOptions).streaming = True
_ = (p
| Create(payloads)
| WriteToPubSub('projects/fakeprj/topics/a_topic',
with_attributes=True))
p.run()
mock_pubsub.return_value.publish.assert_has_calls(
[mock.call(mock.ANY, data, **attributes)])
def test_write_messages_success(self, mock_pubsub):
data = 'data'
payloads = [data]
expected_payloads = [[data, {}]]
mock_pubsub.Client = functools.partial(FakePubsubClient,
messages_write=expected_payloads)
p = TestPipeline()
p.options.view_as(StandardOptions).streaming = True
_ = (p
| Create(payloads)
| WriteToPubSub('projects/fakeprj/topics/a_topic',
with_attributes=False))
p.run()
def test_write_messages_with_attributes_error(self, mock_pubsub):
data = 'data'
# Sending raw data when WriteToPubSub expects a PubsubMessage object.
payloads = [data]
mock_pubsub.Client = functools.partial(FakePubsubClient)
p = TestPipeline()
p.options.view_as(StandardOptions).streaming = True
_ = (p
| Create(payloads)
| WriteToPubSub('projects/fakeprj/topics/a_topic',
with_attributes=True))
with self.assertRaisesRegexp(AttributeError,
r'str.*has no attribute.*data'):
p.run()
