def test_expand_with_both_topic_and_subscription(self):
with self.assertRaisesRegexp(
ValueError,
"Only one of topic or subscription should be provided."):
_ReadFromPubSub('a_topic',
'a_subscription',
'a_label',
with_attributes=False)
def test_expand_with_topic(self):
p = TestPipeline()
p.options.view_as(StandardOptions).streaming = True
pcoll = (p
| _ReadFromPubSub('projects/fakeprj/topics/a_topic',
None,
'a_label',
with_attributes=False)
| beam.Map(lambda x: x))
self.assertEqual(str, pcoll.element_type)
# Apply the necessary PTransformOverrides.
overrides = _get_transform_overrides(p.options)
p.replace_all(overrides)
# Note that the direct output of ReadMessagesFromPubSub will be replaced
# by a PTransformOverride, so we use a no-op Map.
read_transform = pcoll.producer.inputs[0].producer.transform
# Ensure that the properties passed through correctly
source = read_transform._source
self.assertEqual('a_topic', source.topic_name)
self.assertEqual('a_label', source.id_label)
def test_expand_with_no_topic_or_subscription(self):
with self.assertRaisesRegexp(
ValueError,
"Either a topic or subscription must be provided."):
_ReadFromPubSub(None, None, 'a_label', with_attributes=False)