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()