def test_expand_with_wrong_source(self):
with self.assertRaisesRegex(
ValueError, r'PubSub source descriptor must be in the form '
r'"projects/<project>/topics/<topic>"'
' or "projects/<project>/subscription/<subscription>".*'):
MultipleReadFromPubSub(
[PubSubSourceDescriptor('not_a_proper_source')])
def test_expand_with_multiple_sources(self):
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
p = TestPipeline(options=options)
topics = [
'projects/fakeprj/topics/a_topic',
'projects/fakeprj2/topics/b_topic'
]
subscriptions = ['projects/fakeprj/subscriptions/a_subscription']
pubsub_sources = [
PubSubSourceDescriptor(descriptor)
for descriptor in topics + subscriptions
]
pcoll = (p | MultipleReadFromPubSub(pubsub_sources)
| beam.Map(lambda x: x))
# Apply the necessary PTransformOverrides.
overrides = _get_transform_overrides(options)
p.replace_all(overrides)
self.assertEqual(bytes, pcoll.element_type)
# Ensure that the sources are passed through correctly
read_transforms = pcoll.producer.inputs[0].producer.inputs
topics_list = []
subscription_list = []
for read_transform in read_transforms:
source = read_transform.producer.transform._source
if source.full_topic:
topics_list.append(source.full_topic)
else:
subscription_list.append(source.full_subscription)
self.assertEqual(topics_list, topics)
self.assertEqual(subscription_list, subscriptions)
def test_expand_with_multiple_sources_and_other_options(self):
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
p = TestPipeline(options=options)
sources = [
'projects/fakeprj/topics/a_topic',
'projects/fakeprj2/topics/b_topic',
'projects/fakeprj/subscriptions/a_subscription'
]
id_labels = ['a_label_topic', 'b_label_topic', 'a_label_subscription']
timestamp_attributes = [
'a_ta_topic', 'b_ta_topic', 'a_ta_subscription'
]
pubsub_sources = [
PubSubSourceDescriptor(source=source,
id_label=id_label,
timestamp_attribute=timestamp_attribute)
for source, id_label, timestamp_attribute in zip(
sources, id_labels, timestamp_attributes)
]
pcoll = (p | MultipleReadFromPubSub(pubsub_sources)
| beam.Map(lambda x: x))
# Apply the necessary PTransformOverrides.
overrides = _get_transform_overrides(options)
p.replace_all(overrides)
self.assertEqual(bytes, pcoll.element_type)
# Ensure that the sources are passed through correctly
read_transforms = pcoll.producer.inputs[0].producer.inputs
for i, read_transform in enumerate(read_transforms):
id_label = id_labels[i]
timestamp_attribute = timestamp_attributes[i]
source = read_transform.producer.transform._source
self.assertEqual(source.id_label, id_label)
self.assertEqual(source.with_attributes, False)
self.assertEqual(source.timestamp_attribute, timestamp_attribute)