| # coding=utf-8 |
| # |
| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You under the Apache License, Version 2.0 |
| # (the "License"); you may not use this file except in compliance with |
| # the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| |
| """Unit tests for PubSub sources and sinks.""" |
| |
| from __future__ import absolute_import |
| |
| import functools |
| import logging |
| import unittest |
| from builtins import object |
| from builtins import range |
| from builtins import zip |
| |
| import hamcrest as hc |
| import mock |
| |
| import apache_beam as beam |
| from apache_beam.io.gcp.pubsub import PubsubMessage |
| from apache_beam.io.gcp.pubsub import ReadFromPubSub |
| from apache_beam.io.gcp.pubsub import ReadStringsFromPubSub |
| from apache_beam.io.gcp.pubsub import WriteStringsToPubSub |
| from apache_beam.io.gcp.pubsub import WriteToPubSub |
| from apache_beam.io.gcp.pubsub import _PubSubSink |
| from apache_beam.io.gcp.pubsub import _PubSubSource |
| from apache_beam.options.pipeline_options import StandardOptions |
| from apache_beam.runners.direct import transform_evaluator |
| from apache_beam.runners.direct.direct_runner import _DirectReadFromPubSub |
| from apache_beam.runners.direct.direct_runner import _get_transform_overrides |
| from apache_beam.runners.direct.transform_evaluator import _PubSubReadEvaluator |
| from apache_beam.testing.test_pipeline import TestPipeline |
| from apache_beam.testing.util import TestWindowedValue |
| from apache_beam.testing.util import assert_that |
| from apache_beam.testing.util import equal_to |
| from apache_beam.transforms import window |
| from apache_beam.transforms.core import Create |
| from apache_beam.transforms.display import DisplayData |
| from apache_beam.transforms.display_test import DisplayDataItemMatcher |
| from apache_beam.utils import timestamp |
| |
| # Protect against environments where the PubSub library is not available. |
| # pylint: disable=wrong-import-order, wrong-import-position |
| try: |
| from google.cloud import pubsub |
| except ImportError: |
| pubsub = None |
| # pylint: enable=wrong-import-order, wrong-import-position |
| |
| # The protobuf library is only used for running on Dataflow. |
| try: |
| from google.cloud.proto.pubsub.v1 import pubsub_pb2 |
| except ImportError: |
| pubsub_pb2 = None |
| |
| |
| class TestPubsubMessage(unittest.TestCase): |
| |
| def test_payload_valid(self): |
| _ = PubsubMessage('', None) |
| _ = PubsubMessage('data', None) |
| _ = PubsubMessage(None, {'k': 'v'}) |
| |
| def test_payload_invalid(self): |
| with self.assertRaisesRegexp(ValueError, r'data.*attributes.*must be set'): |
| _ = PubsubMessage(None, None) |
| with self.assertRaisesRegexp(ValueError, r'data.*attributes.*must be set'): |
| _ = PubsubMessage(None, {}) |
| |
| @unittest.skipIf(pubsub_pb2 is None, |
| 'PubSub proto dependencies are not installed') |
| def test_proto_conversion(self): |
| data = 'data' |
| attributes = {'k1': 'v1', 'k2': 'v2'} |
| m = PubsubMessage(data, attributes) |
| m_converted = PubsubMessage._from_proto_str(m._to_proto_str()) |
| self.assertEqual(m_converted.data, data) |
| self.assertEqual(m_converted.attributes, attributes) |
| |
| def test_eq(self): |
| a = PubsubMessage('abc', {1: 2, 3: 4}) |
| b = PubsubMessage('abc', {1: 2, 3: 4}) |
| c = PubsubMessage('abc', {1: 2}) |
| self.assertTrue(a == b) |
| self.assertTrue(a != c) |
| self.assertTrue(b != c) |
| |
| def test_hash(self): |
| a = PubsubMessage('abc', {1: 2, 3: 4}) |
| b = PubsubMessage('abc', {1: 2, 3: 4}) |
| c = PubsubMessage('abc', {1: 2}) |
| self.assertTrue(hash(a) == hash(b)) |
| self.assertTrue(hash(a) != hash(c)) |
| self.assertTrue(hash(b) != hash(c)) |
| |
| def test_repr(self): |
| a = PubsubMessage('abc', {1: 2, 3: 4}) |
| b = PubsubMessage('abc', {1: 2, 3: 4}) |
| c = PubsubMessage('abc', {1: 2}) |
| self.assertTrue(repr(a) == repr(b)) |
| self.assertTrue(repr(a) != repr(c)) |
| self.assertTrue(repr(b) != repr(c)) |
| |
| |
| @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') |
| class TestReadFromPubSubOverride(unittest.TestCase): |
| |
| 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, |
| timestamp_attribute=None) |
| | beam.Map(lambda x: x)) |
| self.assertEqual(bytes, pcoll.element_type) |
| |
| # Apply the necessary PTransformOverrides. |
| overrides = _get_transform_overrides(p.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. |
| 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_subscription(self): |
| p = TestPipeline() |
| p.options.view_as(StandardOptions).streaming = True |
| pcoll = (p |
| | ReadFromPubSub( |
| None, 'projects/fakeprj/subscriptions/a_subscription', |
| 'a_label', with_attributes=False, timestamp_attribute=None) |
| | beam.Map(lambda x: x)) |
| self.assertEqual(bytes, pcoll.element_type) |
| |
| # Apply the necessary PTransformOverrides. |
| overrides = _get_transform_overrides(p.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. |
| read_transform = pcoll.producer.inputs[0].producer.transform |
| |
| # Ensure that the properties passed through correctly |
| source = read_transform._source |
| self.assertEqual('a_subscription', source.subscription_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, |
| timestamp_attribute=None) |
| |
| 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, timestamp_attribute=None) |
| |
| def test_expand_with_other_options(self): |
| p = TestPipeline() |
| p.options.view_as(StandardOptions).streaming = True |
| pcoll = (p |
| | ReadFromPubSub('projects/fakeprj/topics/a_topic', |
| None, 'a_label', with_attributes=True, |
| timestamp_attribute='time') |
| | beam.Map(lambda x: x)) |
| self.assertEqual(PubsubMessage, pcoll.element_type) |
| |
| # Apply the necessary PTransformOverrides. |
| overrides = _get_transform_overrides(p.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. |
| read_transform = pcoll.producer.inputs[0].producer.transform |
| |
| # Ensure that the properties passed through correctly |
| source = read_transform._source |
| self.assertTrue(source.with_attributes) |
| self.assertEqual('time', source.timestamp_attribute) |
| |
| |
| @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') |
| class TestWriteStringsToPubSubOverride(unittest.TestCase): |
| def test_expand_deprecated(self): |
| p = TestPipeline() |
| p.options.view_as(StandardOptions).streaming = True |
| pcoll = (p |
| | ReadFromPubSub('projects/fakeprj/topics/baz') |
| | WriteStringsToPubSub('projects/fakeprj/topics/a_topic') |
| | beam.Map(lambda x: x)) |
| |
| # Apply the necessary PTransformOverrides. |
| overrides = _get_transform_overrides(p.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.topic_name) |
| |
| def test_expand(self): |
| p = TestPipeline() |
| p.options.view_as(StandardOptions).streaming = True |
| 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(p.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.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) |
| |
| |
| @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') |
| class TestPubSubSource(unittest.TestCase): |
| def test_display_data_topic(self): |
| source = _PubSubSource( |
| 'projects/fakeprj/topics/a_topic', |
| None, |
| 'a_label') |
| dd = DisplayData.create_from(source) |
| expected_items = [ |
| DisplayDataItemMatcher( |
| 'topic', 'projects/fakeprj/topics/a_topic'), |
| DisplayDataItemMatcher('id_label', 'a_label'), |
| DisplayDataItemMatcher('with_attributes', False), |
| ] |
| |
| hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) |
| |
| def test_display_data_subscription(self): |
| source = _PubSubSource( |
| None, |
| 'projects/fakeprj/subscriptions/a_subscription', |
| 'a_label') |
| dd = DisplayData.create_from(source) |
| expected_items = [ |
| DisplayDataItemMatcher( |
| 'subscription', 'projects/fakeprj/subscriptions/a_subscription'), |
| DisplayDataItemMatcher('id_label', 'a_label'), |
| DisplayDataItemMatcher('with_attributes', False), |
| ] |
| |
| hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) |
| |
| def test_display_data_no_subscription(self): |
| source = _PubSubSource('projects/fakeprj/topics/a_topic') |
| dd = DisplayData.create_from(source) |
| expected_items = [ |
| DisplayDataItemMatcher('topic', 'projects/fakeprj/topics/a_topic'), |
| DisplayDataItemMatcher('with_attributes', False), |
| ] |
| |
| hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) |
| |
| |
| @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') |
| class TestPubSubSink(unittest.TestCase): |
| def test_display_data(self): |
| sink = _PubSubSink('projects/fakeprj/topics/a_topic', |
| id_label='id', with_attributes=False, |
| timestamp_attribute='time') |
| dd = DisplayData.create_from(sink) |
| expected_items = [ |
| DisplayDataItemMatcher('topic', 'projects/fakeprj/topics/a_topic'), |
| DisplayDataItemMatcher('id_label', 'id'), |
| DisplayDataItemMatcher('with_attributes', False), |
| DisplayDataItemMatcher('timestamp_attribute', 'time'), |
| ] |
| |
| hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) |
| |
| |
| class TestPubSubReadEvaluator(object): |
| """Wrapper of _PubSubReadEvaluator that makes it bounded.""" |
| |
| _pubsub_read_evaluator = _PubSubReadEvaluator |
| |
| def __init__(self, *args, **kwargs): |
| self._evaluator = self._pubsub_read_evaluator(*args, **kwargs) |
| |
| def start_bundle(self): |
| return self._evaluator.start_bundle() |
| |
| def process_element(self, element): |
| return self._evaluator.process_element(element) |
| |
| def finish_bundle(self): |
| result = self._evaluator.finish_bundle() |
| result.unprocessed_bundles = [] |
| result.keyed_watermark_holds = {None: None} |
| return result |
| |
| |
| transform_evaluator.TransformEvaluatorRegistry._test_evaluators_overrides = { |
| _DirectReadFromPubSub: TestPubSubReadEvaluator, |
| } |
| |
| |
| class FakePubsubTopic(object): |
| |
| def __init__(self, name, client): |
| self.name = name |
| self.client = client |
| |
| def subscription(self, name): |
| return FakePubsubSubscription(name, self.name, self.client) |
| |
| def batch(self): |
| if self.client.batch is None: |
| self.client.batch = FakeBatch(self.client) |
| return self.client.batch |
| |
| |
| class FakePubsubSubscription(object): |
| |
| def __init__(self, name, topic, client): |
| self.name = name |
| self.topic = topic |
| self.client = client |
| |
| def create(self): |
| pass |
| |
| |
| class FakeAutoAck(object): |
| |
| def __init__(self, sub, **unused_kwargs): |
| self.sub = sub |
| |
| def __enter__(self): |
| messages = self.sub.client.messages_read |
| self.ack_id_to_msg = dict(zip(range(len(messages)), messages)) |
| return self.ack_id_to_msg |
| |
| def __exit__(self, exc_type, exc_val, exc_tb): |
| pass |
| |
| |
| class FakeBatch(object): |
| """Context manager that accept Pubsub client writes via publish(). |
| |
| Verifies writes on exit. |
| """ |
| |
| def __init__(self, client): |
| self.client = client |
| self.published = [] |
| |
| def __enter__(self): |
| return self |
| |
| def __exit__(self, exc_type, exc_val, exc_tb): |
| if exc_type is not None: |
| return # Exception will be raised. |
| hc.assert_that(self.published, |
| hc.only_contains(*self.client.messages_write)) |
| |
| def publish(self, message, **attrs): |
| self.published.append([message, attrs]) |
| |
| |
| class FakePubsubClient(object): |
| |
| def __init__(self, messages_read=None, messages_write=None, project=None, |
| **unused_kwargs): |
| """Creates a Pubsub client fake. |
| |
| Args: |
| messages_read: List of PubsubMessage objects to return. |
| messages_write: List of [data, attributes] pairs, corresponding to |
| messages expected to be written to the client. |
| project: Name of GCP project. |
| """ |
| self.messages_read = messages_read |
| self.messages_write = messages_write |
| self.project = project |
| self.batch = None |
| |
| def topic(self, name): |
| return FakePubsubTopic(name, self) |
| |
| |
| def create_client_message(data, message_id, attributes, publish_time): |
| """Returns a message as it would be returned from Cloud Pub/Sub client. |
| |
| This is what the reader sees. |
| """ |
| msg = pubsub.message.Message(data, message_id, attributes) |
| msg._service_timestamp = publish_time |
| return msg |
| |
| |
| @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') |
| class TestReadFromPubSub(unittest.TestCase): |
| |
| @mock.patch('google.cloud.pubsub') |
| def test_read_messages_success(self, mock_pubsub): |
| data = 'data' |
| message_id = 'message_id' |
| publish_time = '2018-03-12T13:37:01.234567Z' |
| attributes = {'key': 'value'} |
| payloads = [create_client_message( |
| data, message_id, attributes, publish_time)] |
| expected_elements = [ |
| TestWindowedValue(PubsubMessage(data, attributes), |
| timestamp.Timestamp(1520861821.234567), |
| [window.GlobalWindow()])] |
| |
| mock_pubsub.Client = functools.partial(FakePubsubClient, payloads) |
| mock_pubsub.subscription.AutoAck = FakeAutoAck |
| |
| p = TestPipeline() |
| p.options.view_as(StandardOptions).streaming = True |
| pcoll = (p |
| | ReadFromPubSub('projects/fakeprj/topics/a_topic', |
| None, None, with_attributes=True)) |
| assert_that(pcoll, equal_to(expected_elements), reify_windows=True) |
| p.run() |
| |
| @mock.patch('google.cloud.pubsub') |
| def test_read_strings_success(self, mock_pubsub): |
| data = u'🤷 ¯\\_(ツ)_/¯' |
| data_encoded = data.encode('utf-8') |
| publish_time = '2018-03-12T13:37:01.234567Z' |
| payloads = [create_client_message(data_encoded, None, None, publish_time)] |
| expected_elements = [data] |
| |
| mock_pubsub.Client = functools.partial(FakePubsubClient, payloads) |
| mock_pubsub.subscription.AutoAck = FakeAutoAck |
| |
| p = TestPipeline() |
| p.options.view_as(StandardOptions).streaming = True |
| pcoll = (p |
| | ReadStringsFromPubSub('projects/fakeprj/topics/a_topic', |
| None, None)) |
| assert_that(pcoll, equal_to(expected_elements)) |
| p.run() |
| |
| @mock.patch('google.cloud.pubsub') |
| def test_read_data_success(self, mock_pubsub): |
| data_encoded = u'🤷 ¯\\_(ツ)_/¯'.encode('utf-8') |
| publish_time = '2018-03-12T13:37:01.234567Z' |
| payloads = [create_client_message(data_encoded, None, None, publish_time)] |
| expected_elements = [data_encoded] |
| |
| mock_pubsub.Client = functools.partial(FakePubsubClient, payloads) |
| mock_pubsub.subscription.AutoAck = FakeAutoAck |
| |
| p = TestPipeline() |
| p.options.view_as(StandardOptions).streaming = True |
| pcoll = (p |
| | ReadFromPubSub('projects/fakeprj/topics/a_topic', None, None)) |
| assert_that(pcoll, equal_to(expected_elements)) |
| p.run() |
| |
| @mock.patch('google.cloud.pubsub') |
| def test_read_messages_timestamp_attribute_milli_success(self, mock_pubsub): |
| data = 'data' |
| message_id = 'message_id' |
| attributes = {'time': '1337'} |
| publish_time = '2018-03-12T13:37:01.234567Z' |
| payloads = [ |
| create_client_message(data, message_id, attributes, publish_time)] |
| expected_elements = [ |
| TestWindowedValue( |
| PubsubMessage(data, attributes), |
| timestamp.Timestamp(micros=int(attributes['time']) * 1000), |
| [window.GlobalWindow()]), |
| ] |
| |
| mock_pubsub.Client = functools.partial(FakePubsubClient, payloads) |
| mock_pubsub.subscription.AutoAck = FakeAutoAck |
| |
| p = TestPipeline() |
| p.options.view_as(StandardOptions).streaming = True |
| pcoll = (p |
| | ReadFromPubSub( |
| 'projects/fakeprj/topics/a_topic', None, None, |
| with_attributes=True, timestamp_attribute='time')) |
| assert_that(pcoll, equal_to(expected_elements), reify_windows=True) |
| p.run() |
| |
| @mock.patch('google.cloud.pubsub') |
| def test_read_messages_timestamp_attribute_rfc3339_success(self, mock_pubsub): |
| data = 'data' |
| message_id = 'message_id' |
| attributes = {'time': '2018-03-12T13:37:01.234567Z'} |
| publish_time = '2018-03-12T13:37:01.234567Z' |
| payloads = [ |
| create_client_message(data, message_id, attributes, publish_time)] |
| expected_elements = [ |
| TestWindowedValue( |
| PubsubMessage(data, attributes), |
| timestamp.Timestamp.from_rfc3339(attributes['time']), |
| [window.GlobalWindow()]), |
| ] |
| |
| mock_pubsub.Client = functools.partial(FakePubsubClient, payloads) |
| mock_pubsub.subscription.AutoAck = FakeAutoAck |
| |
| p = TestPipeline() |
| p.options.view_as(StandardOptions).streaming = True |
| pcoll = (p |
| | ReadFromPubSub( |
| 'projects/fakeprj/topics/a_topic', None, None, |
| with_attributes=True, timestamp_attribute='time')) |
| assert_that(pcoll, equal_to(expected_elements), reify_windows=True) |
| p.run() |
| |
| @mock.patch('google.cloud.pubsub') |
| def test_read_messages_timestamp_attribute_missing(self, mock_pubsub): |
| data = 'data' |
| message_id = 'message_id' |
| attributes = {} |
| publish_time = '2018-03-12T13:37:01.234567Z' |
| payloads = [ |
| create_client_message(data, message_id, attributes, publish_time)] |
| expected_elements = [ |
| TestWindowedValue( |
| PubsubMessage(data, attributes), |
| timestamp.Timestamp.from_rfc3339(publish_time), |
| [window.GlobalWindow()]), |
| ] |
| |
| mock_pubsub.Client = functools.partial(FakePubsubClient, payloads) |
| mock_pubsub.subscription.AutoAck = FakeAutoAck |
| |
| p = TestPipeline() |
| p.options.view_as(StandardOptions).streaming = True |
| pcoll = (p |
| | ReadFromPubSub( |
| 'projects/fakeprj/topics/a_topic', None, None, |
| with_attributes=True, timestamp_attribute='nonexistent')) |
| assert_that(pcoll, equal_to(expected_elements), reify_windows=True) |
| p.run() |
| |
| @mock.patch('google.cloud.pubsub') |
| def test_read_messages_timestamp_attribute_fail_parse(self, mock_pubsub): |
| data = 'data' |
| message_id = 'message_id' |
| attributes = {'time': '1337 unparseable'} |
| publish_time = '2018-03-12T13:37:01.234567Z' |
| payloads = [ |
| create_client_message(data, message_id, attributes, publish_time)] |
| |
| mock_pubsub.Client = functools.partial(FakePubsubClient, payloads) |
| mock_pubsub.subscription.AutoAck = FakeAutoAck |
| |
| p = TestPipeline() |
| p.options.view_as(StandardOptions).streaming = True |
| _ = (p |
| | ReadFromPubSub( |
| 'projects/fakeprj/topics/a_topic', None, None, |
| with_attributes=True, timestamp_attribute='time')) |
| with self.assertRaisesRegexp(ValueError, r'parse'): |
| p.run() |
| |
| @mock.patch('google.cloud.pubsub') |
| def test_read_message_id_label_unsupported(self, mock_pubsub): |
| # id_label is unsupported in DirectRunner. |
| data = 'data' |
| message_id = 'message_id' |
| attributes = {'time': '1337 unparseable'} |
| publish_time = '2018-03-12T13:37:01.234567Z' |
| payloads = [ |
| create_client_message(data, message_id, attributes, publish_time)] |
| |
| mock_pubsub.Client = functools.partial(FakePubsubClient, payloads) |
| mock_pubsub.subscription.AutoAck = FakeAutoAck |
| |
| p = TestPipeline() |
| p.options.view_as(StandardOptions).streaming = True |
| _ = (p | ReadFromPubSub('projects/fakeprj/topics/a_topic', None, 'a_label')) |
| with self.assertRaisesRegexp(NotImplementedError, |
| r'id_label is not supported'): |
| p.run() |
| |
| |
| @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') |
| class TestWriteToPubSub(unittest.TestCase): |
| |
| @mock.patch('google.cloud.pubsub') |
| 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() |
| |
| @mock.patch('google.cloud.pubsub') |
| def test_write_messages_deprecated(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) |
| | WriteStringsToPubSub('projects/fakeprj/topics/a_topic')) |
| p.run() |
| |
| @mock.patch('google.cloud.pubsub') |
| def test_write_messages_with_attributes_success(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', |
| with_attributes=True)) |
| p.run() |
| |
| @mock.patch('google.cloud.pubsub') |
| 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() |
| |
| @mock.patch('google.cloud.pubsub') |
| 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() |
| |
| |
| if __name__ == '__main__': |
| logging.getLogger().setLevel(logging.INFO) |
| unittest.main() |
