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)