[BEAM-11524] Change WriteStringsToPubsub to use WriteToPubsub. (apach…

…e#13614)
twosom · Dec 28, 2020 · 8df07b6 · 8df07b6
1 parent b8f273e
commit 8df07b6
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 10 deletions.
diff --git a/sdks/python/apache_beam/io/gcp/pubsub.py b/sdks/python/apache_beam/io/gcp/pubsub.py
@@ -252,16 +252,12 @@ def __init__(self, topic):
       topic: Cloud Pub/Sub topic in the form "/topics/<project>/<topic>".
     """
     super(_WriteStringsToPubSub, self).__init__()
-    self.with_attributes = False
-    self.id_label = None
-    self.timestamp_attribute = None
-    self.project, self.topic_name = parse_topic(topic)
-    self._sink = _PubSubSink(topic, id_label=None, timestamp_attribute=None)
+    self.topic = topic
 
   def expand(self, pcoll):
     pcoll = pcoll | 'EncodeString' >> Map(lambda s: s.encode('utf-8'))
     pcoll.element_type = bytes
-    return pcoll | Write(self._sink)
+    return pcoll | WriteToPubSub(self.topic)
 
 
 class WriteToPubSub(PTransform):

diff --git a/sdks/python/apache_beam/io/gcp/pubsub_test.py b/sdks/python/apache_beam/io/gcp/pubsub_test.py
@@ -829,6 +829,7 @@ def test_write_messages_success(self, mock_pubsub):
 
   def test_write_messages_deprecated(self, mock_pubsub):
     data = 'data'
+    data_bytes = b'data'
     payloads = [data]
 
     options = PipelineOptions([])
@@ -839,7 +840,7 @@ def test_write_messages_deprecated(self, mock_pubsub):
           | Create(payloads)
           | WriteStringsToPubSub('projects/fakeprj/topics/a_topic'))
     mock_pubsub.return_value.publish.assert_has_calls(
-        [mock.call(mock.ANY, data)])
+        [mock.call(mock.ANY, data_bytes)])
 
   def test_write_messages_with_attributes_success(self, mock_pubsub):
     data = b'data'

diff --git a/sdks/python/apache_beam/runners/direct/direct_runner.py b/sdks/python/apache_beam/runners/direct/direct_runner.py
@@ -472,9 +472,7 @@ def get_replacement_transform_for_applied_ptransform(
 
   class WriteToPubSubOverride(PTransformOverride):
     def matches(self, applied_ptransform):
-      return isinstance(
-          applied_ptransform.transform,
-          (beam_pubsub.WriteToPubSub, beam_pubsub._WriteStringsToPubSub))
+      return isinstance(applied_ptransform.transform, beam_pubsub.WriteToPubSub)
 
     def get_replacement_transform_for_applied_ptransform(
         self, applied_ptransform):