Merge pull request #6564 from udim/pubsub-0-35-4

[BEAM-5513] Upgrade Python SDK to PubSub 0.35.4
diff --git a/sdks/python/apache_beam/examples/complete/game/game_stats_it_test.py b/sdks/python/apache_beam/examples/complete/game/game_stats_it_test.py
index 6dc60d0..2fc19da 100644
--- a/sdks/python/apache_beam/examples/complete/game/game_stats_it_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/game_stats_it_test.py
@@ -72,15 +72,15 @@
 
     # Set up PubSub environment.
     from google.cloud import pubsub
-    self.pubsub_client = pubsub.Client(project=self.project)
-    unique_topic_name = self.INPUT_TOPIC + _unique_id
-    unique_subscrition_name = self.INPUT_SUB + _unique_id
-    self.input_topic = self.pubsub_client.topic(unique_topic_name)
-    self.input_sub = self.input_topic.subscription(unique_subscrition_name)
+    self.pub_client = pubsub.PublisherClient()
+    self.input_topic = self.pub_client.create_topic(
+        self.pub_client.topic_path(self.project, self.INPUT_TOPIC + _unique_id))
 
-    self.input_topic.create()
-    test_utils.wait_for_topics_created([self.input_topic])
-    self.input_sub.create()
+    self.sub_client = pubsub.SubscriberClient()
+    self.input_sub = self.sub_client.create_subscription(
+        self.sub_client.subscription_path(self.project,
+                                          self.INPUT_SUB + _unique_id),
+        self.input_topic.name)
 
     # Set up BigQuery environment
     from google.cloud import bigquery
@@ -95,14 +95,15 @@
     """Inject game events as test data to PubSub."""
 
     logging.debug('Injecting %d game events to topic %s',
-                  message_count, topic.full_name)
+                  message_count, topic.name)
 
     for _ in range(message_count):
-      topic.publish(self.INPUT_EVENT % self._test_timestamp)
+      self.pub_client.publish(topic.name,
+                              self.INPUT_EVENT % self._test_timestamp)
 
   def _cleanup_pubsub(self):
-    test_utils.cleanup_subscriptions([self.input_sub])
-    test_utils.cleanup_topics([self.input_topic])
+    test_utils.cleanup_subscriptions(self.sub_client, [self.input_sub])
+    test_utils.cleanup_topics(self.pub_client, [self.input_topic])
 
   def _cleanup_dataset(self):
     self.dataset.delete()
@@ -123,9 +124,9 @@
 
     # TODO(mariagh): Add teams table verifier once game_stats.py is fixed.
 
-    extra_opts = {'subscription': self.input_sub.full_name,
+    extra_opts = {'subscription': self.input_sub.name,
                   'dataset': self.dataset.name,
-                  'topic': self.input_topic.full_name,
+                  'topic': self.input_topic.name,
                   'fixed_window_duration': 1,
                   'user_activity_window_duration': 1,
                   'wait_until_finish_duration':
@@ -143,8 +144,6 @@
                     self.dataset.name, self.OUTPUT_TABLE_TEAMS)
 
     # Generate input data and inject to PubSub.
-    test_utils.wait_for_subscriptions_created([self.input_topic,
-                                               self.input_sub])
     self._inject_pubsub_game_events(self.input_topic, self.DEFAULT_INPUT_COUNT)
 
     # Get pipeline options from command argument: --test-pipeline-options,
diff --git a/sdks/python/apache_beam/examples/complete/game/leader_board_it_test.py b/sdks/python/apache_beam/examples/complete/game/leader_board_it_test.py
index ab10942..e0e309b 100644
--- a/sdks/python/apache_beam/examples/complete/game/leader_board_it_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/leader_board_it_test.py
@@ -73,15 +73,16 @@
 
     # Set up PubSub environment.
     from google.cloud import pubsub
-    self.pubsub_client = pubsub.Client(project=self.project)
-    unique_topic_name = self.INPUT_TOPIC + _unique_id
-    unique_subscrition_name = self.INPUT_SUB + _unique_id
-    self.input_topic = self.pubsub_client.topic(unique_topic_name)
-    self.input_sub = self.input_topic.subscription(unique_subscrition_name)
 
-    self.input_topic.create()
-    test_utils.wait_for_topics_created([self.input_topic])
-    self.input_sub.create()
+    self.pub_client = pubsub.PublisherClient()
+    self.input_topic = self.pub_client.create_topic(
+        self.pub_client.topic_path(self.project, self.INPUT_TOPIC + _unique_id))
+
+    self.sub_client = pubsub.SubscriberClient()
+    self.input_sub = self.sub_client.create_subscription(
+        self.sub_client.subscription_path(self.project,
+                                          self.INPUT_SUB + _unique_id),
+        self.input_topic.name)
 
     # Set up BigQuery environment
     from google.cloud import bigquery
@@ -96,14 +97,15 @@
     """Inject game events as test data to PubSub."""
 
     logging.debug('Injecting %d game events to topic %s',
-                  message_count, topic.full_name)
+                  message_count, topic.name)
 
     for _ in range(message_count):
-      topic.publish(self.INPUT_EVENT % self._test_timestamp)
+      self.pub_client.publish(topic.name,
+                              self.INPUT_EVENT % self._test_timestamp)
 
   def _cleanup_pubsub(self):
-    test_utils.cleanup_subscriptions([self.input_sub])
-    test_utils.cleanup_topics([self.input_topic])
+    test_utils.cleanup_subscriptions(self.sub_client, [self.input_sub])
+    test_utils.cleanup_topics(self.pub_client, [self.input_topic])
 
   def _cleanup_dataset(self):
     self.dataset.delete()
@@ -131,9 +133,9 @@
                                         teams_query,
                                         self.DEFAULT_EXPECTED_CHECKSUM)
 
-    extra_opts = {'subscription': self.input_sub.full_name,
+    extra_opts = {'subscription': self.input_sub.name,
                   'dataset': self.dataset.name,
-                  'topic': self.input_topic.full_name,
+                  'topic': self.input_topic.name,
                   'team_window_duration': 1,
                   'wait_until_finish_duration':
                       self.WAIT_UNTIL_FINISH_DURATION,
@@ -151,8 +153,6 @@
                     self.dataset.name, self.OUTPUT_TABLE_TEAMS)
 
     # Generate input data and inject to PubSub.
-    test_utils.wait_for_subscriptions_created([self.input_topic,
-                                               self.input_sub])
     self._inject_pubsub_game_events(self.input_topic, self.DEFAULT_INPUT_COUNT)
 
     # Get pipeline options from command argument: --test-pipeline-options,
diff --git a/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py b/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py
index 3c0cfa9..78e89a1 100644
--- a/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py
+++ b/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py
@@ -52,31 +52,31 @@
 
     # Set up PubSub environment.
     from google.cloud import pubsub
-    self.pubsub_client = pubsub.Client(project=self.project)
-    self.input_topic = self.pubsub_client.topic(INPUT_TOPIC + self.uuid)
-    self.output_topic = self.pubsub_client.topic(OUTPUT_TOPIC + self.uuid)
-    self.input_sub = self.input_topic.subscription(INPUT_SUB + self.uuid)
-    self.output_sub = self.output_topic.subscription(OUTPUT_SUB + self.uuid)
+    self.pub_client = pubsub.PublisherClient()
+    self.input_topic = self.pub_client.create_topic(
+        self.pub_client.topic_path(self.project, INPUT_TOPIC + self.uuid))
+    self.output_topic = self.pub_client.create_topic(
+        self.pub_client.topic_path(self.project, OUTPUT_TOPIC + self.uuid))
 
-    self.input_topic.create()
-    self.output_topic.create()
-    test_utils.wait_for_topics_created([self.input_topic, self.output_topic])
-    self.input_sub.create()
-    self.output_sub.create()
+    self.sub_client = pubsub.SubscriberClient()
+    self.input_sub = self.sub_client.create_subscription(
+        self.sub_client.subscription_path(self.project, INPUT_SUB + self.uuid),
+        self.input_topic.name)
+    self.output_sub = self.sub_client.create_subscription(
+        self.sub_client.subscription_path(self.project, OUTPUT_SUB + self.uuid),
+        self.output_topic.name)
 
   def _inject_numbers(self, topic, num_messages):
     """Inject numbers as test data to PubSub."""
-    logging.debug('Injecting %d numbers to topic %s',
-                  num_messages, topic.full_name)
+    logging.debug('Injecting %d numbers to topic %s', num_messages, topic.name)
     for n in range(num_messages):
-      topic.publish(str(n))
-
-  def _cleanup_pubsub(self):
-    test_utils.cleanup_subscriptions([self.input_sub, self.output_sub])
-    test_utils.cleanup_topics([self.input_topic, self.output_topic])
+      self.pub_client.publish(self.input_topic.name, str(n))
 
   def tearDown(self):
-    self._cleanup_pubsub()
+    test_utils.cleanup_subscriptions(self.sub_client,
+                                     [self.input_sub, self.output_sub])
+    test_utils.cleanup_topics(self.pub_client,
+                              [self.input_topic, self.output_topic])
 
   @attr('IT')
   def test_streaming_wordcount_it(self):
@@ -86,17 +86,16 @@
     # Set extra options to the pipeline for test purpose
     state_verifier = PipelineStateMatcher(PipelineState.RUNNING)
     pubsub_msg_verifier = PubSubMessageMatcher(self.project,
-                                               OUTPUT_SUB + self.uuid,
+                                               self.output_sub.name,
                                                expected_msg,
                                                timeout=400)
-    extra_opts = {'input_subscription': self.input_sub.full_name,
-                  'output_topic': self.output_topic.full_name,
+    extra_opts = {'input_subscription': self.input_sub.name,
+                  'output_topic': self.output_topic.name,
                   'wait_until_finish_duration': WAIT_UNTIL_FINISH_DURATION,
                   'on_success_matcher': all_of(state_verifier,
                                                pubsub_msg_verifier)}
 
     # Generate input data and inject to PubSub.
-    test_utils.wait_for_subscriptions_created([self.input_sub])
     self._inject_numbers(self.input_topic, DEFAULT_INPUT_NUMBERS)
 
     # Get pipeline options from command argument: --test-pipeline-options,
diff --git a/sdks/python/apache_beam/io/gcp/pubsub.py b/sdks/python/apache_beam/io/gcp/pubsub.py
index 2414194..a1644ab 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub.py
@@ -38,11 +38,10 @@
 from apache_beam.transforms.display import DisplayDataItem
 from apache_beam.utils.annotations import deprecated
 
-# The protobuf library is only used for running on Dataflow.
 try:
-  from google.cloud.proto.pubsub.v1 import pubsub_pb2
+  from google.cloud import pubsub
 except ImportError:
-  pubsub_pb2 = None
+  pubsub = None
 
 __all__ = ['PubsubMessage', 'ReadFromPubSub', 'ReadStringsFromPubSub',
            'WriteStringsToPubSub', 'WriteToPubSub']
@@ -92,7 +91,7 @@
     Returns:
       A new PubsubMessage object.
     """
-    msg = pubsub_pb2.PubsubMessage()
+    msg = pubsub.types.pubsub_pb2.PubsubMessage()
     msg.ParseFromString(proto_msg)
     # Convert ScalarMapContainer to dict.
     attributes = dict((key, msg.attributes[key]) for key in msg.attributes)
@@ -109,7 +108,7 @@
       https://cloud.google.com/pubsub/docs/reference/rpc/google.pubsub.v1#google.pubsub.v1.PubsubMessage
       containing the payload of this object.
     """
-    msg = pubsub_pb2.PubsubMessage()
+    msg = pubsub.types.pubsub_pb2.PubsubMessage()
     msg.data = self.data
     for key, value in self.attributes.iteritems():
       msg.attributes[key] = value
@@ -117,9 +116,9 @@
 
   @staticmethod
   def _from_message(msg):
-    """Construct from ``google.cloud.pubsub.message.Message``.
+    """Construct from ``google.cloud.pubsub_v1.subscriber.message.Message``.
 
-    https://google-cloud-python.readthedocs.io/en/latest/pubsub/subscriber/api/message.html
+    https://googleapis.github.io/google-cloud-python/latest/pubsub/subscriber/api/message.html
     """
     # Convert ScalarMapContainer to dict.
     attributes = dict((key, msg.attributes[key]) for key in msg.attributes)
diff --git a/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py b/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py
index 9bb81fc..5b060e5 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py
@@ -100,21 +100,25 @@
 
     # Set up PubSub environment.
     from google.cloud import pubsub
-    self.pubsub_client = pubsub.Client(project=self.project)
-    self.input_topic = self.pubsub_client.topic(INPUT_TOPIC + self.uuid)
-    self.output_topic = self.pubsub_client.topic(OUTPUT_TOPIC + self.uuid)
-    self.input_sub = self.input_topic.subscription(INPUT_SUB + self.uuid)
-    self.output_sub = self.output_topic.subscription(OUTPUT_SUB + self.uuid)
+    self.pub_client = pubsub.PublisherClient()
+    self.input_topic = self.pub_client.create_topic(
+        self.pub_client.topic_path(self.project, INPUT_TOPIC + self.uuid))
+    self.output_topic = self.pub_client.create_topic(
+        self.pub_client.topic_path(self.project, OUTPUT_TOPIC + self.uuid))
 
-    self.input_topic.create()
-    self.output_topic.create()
-    test_utils.wait_for_topics_created([self.input_topic, self.output_topic])
-    self.input_sub.create()
-    self.output_sub.create()
+    self.sub_client = pubsub.SubscriberClient()
+    self.input_sub = self.sub_client.create_subscription(
+        self.sub_client.subscription_path(self.project, INPUT_SUB + self.uuid),
+        self.input_topic.name)
+    self.output_sub = self.sub_client.create_subscription(
+        self.sub_client.subscription_path(self.project, OUTPUT_SUB + self.uuid),
+        self.output_topic.name)
 
   def tearDown(self):
-    test_utils.cleanup_subscriptions([self.input_sub, self.output_sub])
-    test_utils.cleanup_topics([self.input_topic, self.output_topic])
+    test_utils.cleanup_subscriptions(self.sub_client,
+                                     [self.input_sub, self.output_sub])
+    test_utils.cleanup_topics(self.pub_client,
+                              [self.input_topic, self.output_topic])
 
   def _test_streaming(self, with_attributes):
     """Runs IT pipeline with message verifier.
@@ -139,21 +143,20 @@
       strip_attributes = [self.ID_LABEL, self.TIMESTAMP_ATTRIBUTE]
     pubsub_msg_verifier = PubSubMessageMatcher(
         self.project,
-        OUTPUT_SUB + self.uuid,
+        self.output_sub.name,
         expected_messages,
         timeout=MESSAGE_MATCHER_TIMEOUT_S,
         with_attributes=with_attributes,
         strip_attributes=strip_attributes)
-    extra_opts = {'input_subscription': self.input_sub.full_name,
-                  'output_topic': self.output_topic.full_name,
+    extra_opts = {'input_subscription': self.input_sub.name,
+                  'output_topic': self.output_topic.name,
                   'wait_until_finish_duration': TEST_PIPELINE_DURATION_MS,
                   'on_success_matcher': all_of(state_verifier,
                                                pubsub_msg_verifier)}
 
     # Generate input data and inject to PubSub.
-    test_utils.wait_for_subscriptions_created([self.input_sub])
     for msg in self.INPUT_MESSAGES[self.runner_name]:
-      self.input_topic.publish(msg.data, **msg.attributes)
+      self.pub_client.publish(self.input_topic.name, msg.data, **msg.attributes)
 
     # Get pipeline options from command argument: --test-pipeline-options,
     # and start pipeline job by calling pipeline main function.
diff --git a/sdks/python/apache_beam/io/gcp/pubsub_test.py b/sdks/python/apache_beam/io/gcp/pubsub_test.py
index 6e19950..a95ffc6 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub_test.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub_test.py
@@ -20,12 +20,9 @@
 
 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
@@ -43,6 +40,7 @@
 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 import test_utils
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.util import TestWindowedValue
 from apache_beam.testing.util import assert_that
@@ -54,18 +52,10 @@
 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):
@@ -81,8 +71,7 @@
     with self.assertRaisesRegexp(ValueError, r'data.*attributes.*must be set'):
       _ = PubsubMessage(None, {})
 
-  @unittest.skipIf(pubsub_pb2 is None,
-                   'PubSub proto dependencies are not installed')
+  @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed')
   def test_proto_conversion(self):
     data = 'data'
     attributes = {'k1': 'v1', 'k2': 'v2'}
@@ -220,7 +209,7 @@
     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('a_topic', write_transform.dofn.short_topic_name)
 
   def test_expand(self):
     p = TestPipeline()
@@ -240,7 +229,7 @@
     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('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)
@@ -333,118 +322,25 @@
 }
 
 
-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')
+@mock.patch('google.cloud.pubsub.SubscriberClient')
 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'
+    publish_time_secs = 1520861821
+    publish_time_nanos = 234567000
     attributes = {'key': 'value'}
-    payloads = [create_client_message(
-        data, message_id, attributes, publish_time)]
+    ack_id = 'ack_id'
+    pull_response = test_utils.create_pull_response([
+        test_utils.PullResponseMessage(
+            data, attributes, publish_time_secs, publish_time_nanos, ack_id)
+    ])
     expected_elements = [
         TestWindowedValue(PubsubMessage(data, attributes),
                           timestamp.Timestamp(1520861821.234567),
                           [window.GlobalWindow()])]
-
-    mock_pubsub.Client = functools.partial(FakePubsubClient, payloads)
-    mock_pubsub.subscription.AutoAck = FakeAutoAck
+    mock_pubsub.return_value.pull.return_value = pull_response
 
     p = TestPipeline()
     p.options.view_as(StandardOptions).streaming = True
@@ -453,17 +349,18 @@
                               None, None, with_attributes=True))
     assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
     p.run()
+    mock_pubsub.return_value.acknowledge.assert_has_calls([
+        mock.call(mock.ANY, [ack_id])])
 
-  @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)]
+    ack_id = 'ack_id'
+    pull_response = test_utils.create_pull_response([
+        test_utils.PullResponseMessage(data_encoded, ack_id=ack_id)
+    ])
     expected_elements = [data]
-
-    mock_pubsub.Client = functools.partial(FakePubsubClient, payloads)
-    mock_pubsub.subscription.AutoAck = FakeAutoAck
+    mock_pubsub.return_value.pull.return_value = pull_response
 
     p = TestPipeline()
     p.options.view_as(StandardOptions).streaming = True
@@ -472,16 +369,16 @@
                                      None, None))
     assert_that(pcoll, equal_to(expected_elements))
     p.run()
+    mock_pubsub.return_value.acknowledge.assert_has_calls([
+        mock.call(mock.ANY, [ack_id])])
 
-  @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)]
+    ack_id = 'ack_id'
+    pull_response = test_utils.create_pull_response([
+        test_utils.PullResponseMessage(data_encoded, ack_id=ack_id)])
     expected_elements = [data_encoded]
-
-    mock_pubsub.Client = functools.partial(FakePubsubClient, payloads)
-    mock_pubsub.subscription.AutoAck = FakeAutoAck
+    mock_pubsub.return_value.pull.return_value = pull_response
 
     p = TestPipeline()
     p.options.view_as(StandardOptions).streaming = True
@@ -489,24 +386,26 @@
              | ReadFromPubSub('projects/fakeprj/topics/a_topic', None, None))
     assert_that(pcoll, equal_to(expected_elements))
     p.run()
+    mock_pubsub.return_value.acknowledge.assert_has_calls([
+        mock.call(mock.ANY, [ack_id])])
 
-  @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)]
+    publish_time_secs = 1520861821
+    publish_time_nanos = 234567000
+    ack_id = 'ack_id'
+    pull_response = test_utils.create_pull_response([
+        test_utils.PullResponseMessage(
+            data, attributes, publish_time_secs, publish_time_nanos, ack_id)
+    ])
     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
+    mock_pubsub.return_value.pull.return_value = pull_response
 
     p = TestPipeline()
     p.options.view_as(StandardOptions).streaming = True
@@ -516,24 +415,26 @@
                  with_attributes=True, timestamp_attribute='time'))
     assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
     p.run()
+    mock_pubsub.return_value.acknowledge.assert_has_calls([
+        mock.call(mock.ANY, [ack_id])])
 
-  @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)]
+    publish_time_secs = 1337000000
+    publish_time_nanos = 133700000
+    ack_id = 'ack_id'
+    pull_response = test_utils.create_pull_response([
+        test_utils.PullResponseMessage(
+            data, attributes, publish_time_secs, publish_time_nanos, ack_id)
+    ])
     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
+    mock_pubsub.return_value.pull.return_value = pull_response
 
     p = TestPipeline()
     p.options.view_as(StandardOptions).streaming = True
@@ -543,24 +444,27 @@
                  with_attributes=True, timestamp_attribute='time'))
     assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
     p.run()
+    mock_pubsub.return_value.acknowledge.assert_has_calls([
+        mock.call(mock.ANY, [ack_id])])
 
-  @mock.patch('google.cloud.pubsub')
   def test_read_messages_timestamp_attribute_missing(self, mock_pubsub):
     data = 'data'
-    message_id = 'message_id'
     attributes = {}
+    publish_time_secs = 1520861821
+    publish_time_nanos = 234567000
     publish_time = '2018-03-12T13:37:01.234567Z'
-    payloads = [
-        create_client_message(data, message_id, attributes, publish_time)]
+    ack_id = 'ack_id'
+    pull_response = test_utils.create_pull_response([
+        test_utils.PullResponseMessage(
+            data, attributes, publish_time_secs, publish_time_nanos, ack_id)
+    ])
     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
+    mock_pubsub.return_value.pull.return_value = pull_response
 
     p = TestPipeline()
     p.options.view_as(StandardOptions).streaming = True
@@ -570,18 +474,20 @@
                  with_attributes=True, timestamp_attribute='nonexistent'))
     assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
     p.run()
+    mock_pubsub.return_value.acknowledge.assert_has_calls([
+        mock.call(mock.ANY, [ack_id])])
 
-  @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
+    publish_time_secs = 1520861821
+    publish_time_nanos = 234567000
+    ack_id = 'ack_id'
+    pull_response = test_utils.create_pull_response([
+        test_utils.PullResponseMessage(
+            data, attributes, publish_time_secs, publish_time_nanos, ack_id)
+    ])
+    mock_pubsub.return_value.pull.return_value = pull_response
 
     p = TestPipeline()
     p.options.view_as(StandardOptions).streaming = True
@@ -591,20 +497,10 @@
              with_attributes=True, timestamp_attribute='time'))
     with self.assertRaisesRegexp(ValueError, r'parse'):
       p.run()
+    mock_pubsub.return_value.acknowledge.assert_not_called()
 
-  @mock.patch('google.cloud.pubsub')
-  def test_read_message_id_label_unsupported(self, mock_pubsub):
+  def test_read_message_id_label_unsupported(self, unused_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'))
@@ -614,16 +510,12 @@
 
 
 @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed')
+@mock.patch('google.cloud.pubsub.PublisherClient')
 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
@@ -632,15 +524,12 @@
          | WriteToPubSub('projects/fakeprj/topics/a_topic',
                          with_attributes=False))
     p.run()
+    mock_pubsub.return_value.publish.assert_has_calls([
+        mock.call(mock.ANY, data)])
 
-  @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
@@ -648,16 +537,13 @@
          | Create(payloads)
          | WriteStringsToPubSub('projects/fakeprj/topics/a_topic'))
     p.run()
+    mock_pubsub.return_value.publish.assert_has_calls([
+        mock.call(mock.ANY, data)])
 
-  @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
@@ -666,15 +552,14 @@
          | 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)])
 
-  @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
@@ -685,15 +570,10 @@
                                  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
diff --git a/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher.py b/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher.py
index 6217faf..7a0b5c8 100644
--- a/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher.py
+++ b/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher.py
@@ -31,12 +31,10 @@
 
 
 # Protect against environments where 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
 
 DEFAULT_TIMEOUT = 5 * 60
 MAX_MESSAGES_IN_ONE_PULL = 50
@@ -49,8 +47,9 @@
   subscription until all expected messages are shown or timeout.
   """
 
-  def __init__(self, project, sub_name, expected_msg, timeout=DEFAULT_TIMEOUT,
-               with_attributes=False, strip_attributes=None):
+  def __init__(self, project, sub_name, expected_msg,
+               timeout=DEFAULT_TIMEOUT, with_attributes=False,
+               strip_attributes=None):
     """Initialize PubSubMessageMatcher object.
 
     Args:
@@ -59,8 +58,9 @@
       expected_msg: A string list that contains expected message data pulled
         from the subscription. See also: with_attributes.
       timeout: Timeout in seconds to wait for all expected messages appears.
-      with_attributes: Whether expected_msg is a list of
-        ``PubsubMessage`` objects.
+      with_attributes: If True, will match against both message data and
+        attributes. If True, expected_msg should be a list of ``PubsubMessage``
+        objects. Otherwise, it should be a list of ``bytes``.
       strip_attributes: List of strings. If with_attributes==True, strip the
         attributes keyed by these values from incoming messages.
         If a key is missing, will add an attribute with an error message as
@@ -86,28 +86,26 @@
 
   def _matches(self, _):
     if self.messages is None:
-      self.messages = self._wait_for_messages(self._get_subscription(),
-                                              len(self.expected_msg),
+      self.messages = self._wait_for_messages(len(self.expected_msg),
                                               self.timeout)
     return Counter(self.messages) == Counter(self.expected_msg)
 
-  def _get_subscription(self):
-    return pubsub.Client(project=self.project).subscription(self.sub_name)
-
-  def _wait_for_messages(self, subscription, expected_num, timeout):
+  def _wait_for_messages(self, expected_num, timeout):
     """Wait for messages from given subscription."""
-    logging.debug('Start pulling messages from %s', subscription.full_name)
     total_messages = []
+
+    sub_client = pubsub.SubscriberClient()
     start_time = time.time()
     while time.time() - start_time <= timeout:
-      pulled = subscription.pull(max_messages=MAX_MESSAGES_IN_ONE_PULL)
-      for ack_id, message in pulled:
-        subscription.acknowledge([ack_id])
+      response = sub_client.pull(self.sub_name,
+                                 max_messages=MAX_MESSAGES_IN_ONE_PULL,
+                                 return_immediately=True)
+      for rm in response.received_messages:
+        msg = PubsubMessage._from_message(rm.message)
         if not self.with_attributes:
-          total_messages.append(message.data)
+          total_messages.append(msg.data)
           continue
 
-        msg = PubsubMessage._from_message(message)
         if self.strip_attributes:
           for attr in self.strip_attributes:
             try:
@@ -117,12 +115,16 @@
                                       'expected attribute not found.')
         total_messages.append(msg)
 
+      ack_ids = [rm.ack_id for rm in response.received_messages]
+      if ack_ids:
+        sub_client.acknowledge(self.sub_name, ack_ids)
       if len(total_messages) >= expected_num:
-        return total_messages
+        break
       time.sleep(1)
 
-    logging.error('Timeout after %d sec. Received %d messages from %s.',
-                  timeout, len(total_messages), subscription.full_name)
+    if time.time() - start_time > timeout:
+      logging.error('Timeout after %d sec. Received %d messages from %s.',
+                    timeout, len(total_messages), self.sub_name)
     return total_messages
 
   def describe_to(self, description):
diff --git a/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher_test.py b/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher_test.py
index 0e59481..0f9351d 100644
--- a/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher_test.py
+++ b/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher_test.py
@@ -27,17 +27,19 @@
 
 from apache_beam.io.gcp.pubsub import PubsubMessage
 from apache_beam.io.gcp.tests.pubsub_matcher import PubSubMessageMatcher
+from apache_beam.testing.test_utils import PullResponseMessage
+from apache_beam.testing.test_utils import create_pull_response
 
 # Protect against environments where 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
 
 
 @unittest.skipIf(pubsub is None, 'PubSub dependencies are not installed.')
+@mock.patch('time.sleep', return_value=None)
+@mock.patch('google.cloud.pubsub.SubscriberClient')
 class PubSubMatcherTest(unittest.TestCase):
 
   def setUp(self):
@@ -48,90 +50,75 @@
         'mock_project', 'mock_sub_name', ['mock_expected_msg'],
         with_attributes=with_attributes, strip_attributes=strip_attributes)
 
-  @mock.patch('time.sleep', return_value=None)
-  @mock.patch('apache_beam.io.gcp.tests.pubsub_matcher.'
-              'PubSubMessageMatcher._get_subscription')
   def test_message_matcher_success(self, mock_get_sub, unsued_mock):
     self.init_matcher()
     self.pubsub_matcher.expected_msg = ['a', 'b']
     mock_sub = mock_get_sub.return_value
     mock_sub.pull.side_effect = [
-        [(1, pubsub.message.Message(b'a', 'unused_id'))],
-        [(2, pubsub.message.Message(b'b', 'unused_id'))],
+        create_pull_response([PullResponseMessage(b'a', {})]),
+        create_pull_response([PullResponseMessage(b'b', {})]),
     ]
     hc_assert_that(self.mock_presult, self.pubsub_matcher)
     self.assertEqual(mock_sub.pull.call_count, 2)
+    self.assertEqual(mock_sub.acknowledge.call_count, 2)
 
-  @mock.patch('time.sleep', return_value=None)
-  @mock.patch('apache_beam.io.gcp.tests.pubsub_matcher.'
-              'PubSubMessageMatcher._get_subscription')
   def test_message_matcher_attributes_success(self, mock_get_sub, unsued_mock):
     self.init_matcher(with_attributes=True)
     self.pubsub_matcher.expected_msg = [PubsubMessage('a', {'k': 'v'})]
     mock_sub = mock_get_sub.return_value
-    msg_a = pubsub.message.Message(b'a', 'unused_id')
-    msg_a.attributes['k'] = 'v'
-    mock_sub.pull.side_effect = [[(1, msg_a)]]
+    mock_sub.pull.side_effect = [
+        create_pull_response([PullResponseMessage(b'a', {'k': 'v'})])
+    ]
     hc_assert_that(self.mock_presult, self.pubsub_matcher)
     self.assertEqual(mock_sub.pull.call_count, 1)
+    self.assertEqual(mock_sub.acknowledge.call_count, 1)
 
-  @mock.patch('time.sleep', return_value=None)
-  @mock.patch('apache_beam.io.gcp.tests.pubsub_matcher.'
-              'PubSubMessageMatcher._get_subscription')
   def test_message_matcher_attributes_fail(self, mock_get_sub, unsued_mock):
     self.init_matcher(with_attributes=True)
     self.pubsub_matcher.expected_msg = [PubsubMessage('a', {})]
     mock_sub = mock_get_sub.return_value
-    msg_a = pubsub.message.Message(b'a', 'unused_id')
-    msg_a.attributes['k'] = 'v'  # Unexpected.
-    mock_sub.pull.side_effect = [[(1, msg_a)]]
+    # Unexpected attribute 'k'.
+    mock_sub.pull.side_effect = [
+        create_pull_response([PullResponseMessage(b'a', {'k': 'v'})])
+    ]
     with self.assertRaisesRegexp(AssertionError, r'Unexpected'):
       hc_assert_that(self.mock_presult, self.pubsub_matcher)
     self.assertEqual(mock_sub.pull.call_count, 1)
+    self.assertEqual(mock_sub.acknowledge.call_count, 1)
 
-  @mock.patch('time.sleep', return_value=None)
-  @mock.patch('apache_beam.io.gcp.tests.pubsub_matcher.'
-              'PubSubMessageMatcher._get_subscription')
   def test_message_matcher_strip_success(self, mock_get_sub, unsued_mock):
     self.init_matcher(with_attributes=True,
                       strip_attributes=['id', 'timestamp'])
     self.pubsub_matcher.expected_msg = [PubsubMessage('a', {'k': 'v'})]
     mock_sub = mock_get_sub.return_value
-    msg_a = pubsub.message.Message(b'a', 'unused_id')
-    msg_a.attributes['id'] = 'foo'
-    msg_a.attributes['timestamp'] = 'bar'
-    msg_a.attributes['k'] = 'v'
-    mock_sub.pull.side_effect = [[(1, msg_a)]]
+    mock_sub.pull.side_effect = [create_pull_response([
+        PullResponseMessage(b'a', {'id': 'foo', 'timestamp': 'bar', 'k': 'v'})
+    ])]
     hc_assert_that(self.mock_presult, self.pubsub_matcher)
     self.assertEqual(mock_sub.pull.call_count, 1)
+    self.assertEqual(mock_sub.acknowledge.call_count, 1)
 
-  @mock.patch('time.sleep', return_value=None)
-  @mock.patch('apache_beam.io.gcp.tests.pubsub_matcher.'
-              'PubSubMessageMatcher._get_subscription')
   def test_message_matcher_strip_fail(self, mock_get_sub, unsued_mock):
     self.init_matcher(with_attributes=True,
                       strip_attributes=['id', 'timestamp'])
     self.pubsub_matcher.expected_msg = [PubsubMessage('a', {'k': 'v'})]
     mock_sub = mock_get_sub.return_value
-    # msg_a is missing attribute 'timestamp'.
-    msg_a = pubsub.message.Message(b'a', 'unused_id')
-    msg_a.attributes['id'] = 'foo'
-    msg_a.attributes['k'] = 'v'
-    mock_sub.pull.side_effect = [[(1, msg_a)]]
+    # Message is missing attribute 'timestamp'.
+    mock_sub.pull.side_effect = [create_pull_response([
+        PullResponseMessage(b'a', {'id': 'foo', 'k': 'v'})
+    ])]
     with self.assertRaisesRegexp(AssertionError, r'Stripped attributes'):
       hc_assert_that(self.mock_presult, self.pubsub_matcher)
     self.assertEqual(mock_sub.pull.call_count, 1)
+    self.assertEqual(mock_sub.acknowledge.call_count, 1)
 
-  @mock.patch('time.sleep', return_value=None)
-  @mock.patch('apache_beam.io.gcp.tests.pubsub_matcher.'
-              'PubSubMessageMatcher._get_subscription')
   def test_message_matcher_mismatch(self, mock_get_sub, unused_mock):
     self.init_matcher()
     self.pubsub_matcher.expected_msg = ['a']
     mock_sub = mock_get_sub.return_value
-    mock_sub.pull.return_value = [
-        (1, pubsub.message.Message(b'c', 'unused_id')),
-        (1, pubsub.message.Message(b'd', 'unused_id')),
+    mock_sub.pull.side_effect = [
+        create_pull_response([PullResponseMessage(b'c', {}),
+                              PullResponseMessage(b'd', {})]),
     ]
     with self.assertRaises(AssertionError) as error:
       hc_assert_that(self.mock_presult, self.pubsub_matcher)
@@ -140,10 +127,9 @@
     self.assertTrue(
         '\nExpected: Expected 1 messages.\n     but: Got 2 messages.'
         in str(error.exception.args[0]))
+    self.assertEqual(mock_sub.pull.call_count, 1)
+    self.assertEqual(mock_sub.acknowledge.call_count, 1)
 
-  @mock.patch('time.sleep', return_value=None)
-  @mock.patch('apache_beam.io.gcp.tests.pubsub_matcher.'
-              'PubSubMessageMatcher._get_subscription')
   def test_message_matcher_timeout(self, mock_get_sub, unused_mock):
     self.init_matcher()
     mock_sub = mock_get_sub.return_value
@@ -152,6 +138,7 @@
     with self.assertRaisesRegexp(AssertionError, r'Expected 1.*\n.*Got 0'):
       hc_assert_that(self.mock_presult, self.pubsub_matcher)
     self.assertTrue(mock_sub.pull.called)
+    self.assertEqual(mock_sub.acknowledge.call_count, 0)
 
 
 if __name__ == '__main__':
diff --git a/sdks/python/apache_beam/runners/direct/direct_runner.py b/sdks/python/apache_beam/runners/direct/direct_runner.py
index 00e37f3..d410992 100644
--- a/sdks/python/apache_beam/runners/direct/direct_runner.py
+++ b/sdks/python/apache_beam/runners/direct/direct_runner.py
@@ -25,6 +25,7 @@
 
 import itertools
 import logging
+import time
 
 from google.protobuf import wrappers_pb2
 
@@ -264,11 +265,12 @@
 
 
 class _DirectWriteToPubSubFn(DoFn):
-  _topic = None
+  BUFFER_SIZE_ELEMENTS = 100
+  FLUSH_TIMEOUT_SECS = BUFFER_SIZE_ELEMENTS * 0.5
 
   def __init__(self, sink):
     self.project = sink.project
-    self.topic_name = sink.topic_name
+    self.short_topic_name = sink.topic_name
     self.id_label = sink.id_label
     self.timestamp_attribute = sink.timestamp_attribute
     self.with_attributes = sink.with_attributes
@@ -282,30 +284,33 @@
                                 'supported for PubSub writes')
 
   def start_bundle(self):
-    from google.cloud import pubsub
-
-    if self._topic is None:
-      self._topic = pubsub.Client(project=self.project).topic(
-          self.topic_name)
     self._buffer = []
 
   def process(self, elem):
     self._buffer.append(elem)
-    if len(self._buffer) >= 100:
+    if len(self._buffer) >= self.BUFFER_SIZE_ELEMENTS:
       self._flush()
 
   def finish_bundle(self):
     self._flush()
 
   def _flush(self):
-    if self._buffer:
-      with self._topic.batch() as batch:
-        for elem in self._buffer:
-          if self.with_attributes:
-            batch.publish(elem.data, **elem.attributes)
-          else:
-            batch.publish(elem)
-      self._buffer = []
+    from google.cloud import pubsub
+    pub_client = pubsub.PublisherClient()
+    topic = pub_client.topic_path(self.project, self.short_topic_name)
+
+    if self.with_attributes:
+      futures = [pub_client.publish(topic, elem.data, **elem.attributes)
+                 for elem in self._buffer]
+    else:
+      futures = [pub_client.publish(topic, elem)
+                 for elem in self._buffer]
+
+    timer_start = time.time()
+    for future in futures:
+      remaining = self.FLUSH_TIMEOUT_SECS - (time.time() - timer_start)
+      future.result(remaining)
+    self._buffer = []
 
 
 def _get_pubsub_transform_overrides(pipeline_options):
diff --git a/sdks/python/apache_beam/runners/direct/test_direct_runner.py b/sdks/python/apache_beam/runners/direct/test_direct_runner.py
index 8facca8..23dfeab 100644
--- a/sdks/python/apache_beam/runners/direct/test_direct_runner.py
+++ b/sdks/python/apache_beam/runners/direct/test_direct_runner.py
@@ -52,5 +52,6 @@
     finally:
       if not PipelineState.is_terminal(self.result.state):
         self.result.cancel()
+        self.result.wait_until_finish()
 
     return self.result
diff --git a/sdks/python/apache_beam/runners/direct/transform_evaluator.py b/sdks/python/apache_beam/runners/direct/transform_evaluator.py
index ef12e2c..fad0704 100644
--- a/sdks/python/apache_beam/runners/direct/transform_evaluator.py
+++ b/sdks/python/apache_beam/runners/direct/transform_evaluator.py
@@ -377,20 +377,44 @@
 
 
 class _PubSubSubscriptionWrapper(object):
-  """Wrapper for garbage-collecting temporary PubSub subscriptions."""
+  """Wrapper for managing temporary PubSub subscriptions."""
 
-  def __init__(self, subscription, should_cleanup):
-    self.subscription = subscription
-    self.should_cleanup = should_cleanup
+  def __init__(self, project, short_topic_name, short_sub_name):
+    """Initialize subscription wrapper.
+
+    If sub_name is None, will create a temporary subscription to topic_name.
+
+    Args:
+      project: GCP project name for topic and subscription. May be None.
+        Required if sub_name is None.
+      short_topic_name: Valid topic name without
+        'projects/{project}/topics/' prefix. May be None.
+        Required if sub_name is None.
+      short_sub_name: Valid subscription name without
+        'projects/{project}/subscriptions/' prefix. May be None.
+    """
+    from google.cloud import pubsub
+    self.sub_client = pubsub.SubscriberClient()
+
+    if short_sub_name is None:
+      self.sub_name = self.sub_client.subscription_path(
+          project, 'beam_%d_%x' % (int(time.time()), random.randrange(1 << 32)))
+      topic_name = self.sub_client.topic_path(project, short_topic_name)
+      self.sub_client.create_subscription(self.sub_name, topic_name)
+      self._should_cleanup = True
+    else:
+      self.sub_name = self.sub_client.subscription_path(project, short_sub_name)
+      self._should_cleanup = False
 
   def __del__(self):
-    if self.should_cleanup:
-      self.subscription.delete()
+    if self._should_cleanup:
+      self.sub_client.delete_subscription(self.sub_name)
 
 
 class _PubSubReadEvaluator(_TransformEvaluator):
   """TransformEvaluator for PubSub read."""
 
+  # A mapping of transform to _PubSubSubscriptionWrapper.
   _subscription_cache = {}
 
   def __init__(self, evaluation_context, applied_ptransform,
@@ -404,26 +428,16 @@
     if self.source.id_label:
       raise NotImplementedError(
           'DirectRunner: id_label is not supported for PubSub reads')
-    self._subscription = _PubSubReadEvaluator.get_subscription(
+    self._sub_name = _PubSubReadEvaluator.get_subscription(
         self._applied_ptransform, self.source.project, self.source.topic_name,
         self.source.subscription_name)
 
   @classmethod
-  def get_subscription(cls, transform, project, topic, subscription_name):
+  def get_subscription(cls, transform, project, topic, short_sub_name):
     if transform not in cls._subscription_cache:
-      from google.cloud import pubsub
-      should_create = not subscription_name
-      if should_create:
-        subscription_name = 'beam_%d_%x' % (
-            int(time.time()), random.randrange(1 << 32))
-      wrapper = _PubSubSubscriptionWrapper(
-          pubsub.Client(project=project).topic(topic).subscription(
-              subscription_name),
-          should_create)
-      if should_create:
-        wrapper.subscription.create()
+      wrapper = _PubSubSubscriptionWrapper(project, topic, short_sub_name)
       cls._subscription_cache[transform] = wrapper
-    return cls._subscription_cache[transform].subscription
+    return cls._subscription_cache[transform].sub_name
 
   def start_bundle(self):
     pass
@@ -438,28 +452,34 @@
     # evaluator fails with an exception before emitting a bundle. However,
     # the DirectRunner currently doesn't retry work items anyway, so the
     # pipeline would enter an inconsistent state on any error.
-    with pubsub.subscription.AutoAck(
-        self._subscription, return_immediately=True,
-        max_messages=10) as results:
-      def _get_element(message):
-        parsed_message = PubsubMessage._from_message(message)
-        if (timestamp_attribute and
-            timestamp_attribute in parsed_message.attributes):
-          rfc3339_or_milli = parsed_message.attributes[timestamp_attribute]
+    sub_client = pubsub.SubscriberClient()
+    response = sub_client.pull(self._sub_name, max_messages=10,
+                               return_immediately=True)
+
+    def _get_element(message):
+      parsed_message = PubsubMessage._from_message(message)
+      if (timestamp_attribute and
+          timestamp_attribute in parsed_message.attributes):
+        rfc3339_or_milli = parsed_message.attributes[timestamp_attribute]
+        try:
+          timestamp = Timestamp.from_rfc3339(rfc3339_or_milli)
+        except ValueError:
           try:
-            timestamp = Timestamp.from_rfc3339(rfc3339_or_milli)
-          except ValueError:
-            try:
-              timestamp = Timestamp(micros=int(rfc3339_or_milli) * 1000)
-            except ValueError as e:
-              raise ValueError('Bad timestamp value: %s' % e)
-        else:
-          timestamp = Timestamp.from_rfc3339(message.service_timestamp)
+            timestamp = Timestamp(micros=int(rfc3339_or_milli) * 1000)
+          except ValueError as e:
+            raise ValueError('Bad timestamp value: %s' % e)
+      else:
+        timestamp = Timestamp(message.publish_time.seconds,
+                              message.publish_time.nanos // 1000)
 
-        return timestamp, parsed_message
+      return timestamp, parsed_message
 
-      return [_get_element(message)
-              for unused_ack_id, message in iteritems(results)]
+    results = [_get_element(rm.message) for rm in response.received_messages]
+    ack_ids = [rm.ack_id for rm in response.received_messages]
+    if ack_ids:
+      sub_client.acknowledge(self._sub_name, ack_ids)
+
+    return results
 
   def finish_bundle(self):
     data = self._read_from_pubsub(self.source.timestamp_attribute)
diff --git a/sdks/python/apache_beam/testing/test_utils.py b/sdks/python/apache_beam/testing/test_utils.py
index 490d079..1f0e99e 100644
--- a/sdks/python/apache_beam/testing/test_utils.py
+++ b/sdks/python/apache_beam/testing/test_utils.py
@@ -24,11 +24,9 @@
 
 import hashlib
 import imp
-import logging
 import os
 import shutil
 import tempfile
-import time
 from builtins import object
 
 from mock import Mock
@@ -136,46 +134,61 @@
   FileSystems.delete(file_paths)
 
 
-def wait_for_subscriptions_created(subs, timeout=60):
-  """Wait for all PubSub subscriptions are created."""
-  return _wait_until_all_exist(subs, timeout)
-
-
-def wait_for_topics_created(topics, timeout=60):
-  """Wait for all PubSub topics are created."""
-  return _wait_until_all_exist(topics, timeout)
-
-
-def _wait_until_all_exist(components, timeout):
-  unchecked_components = set(components)
-  start_time = time.time()
-  while time.time() - start_time <= timeout:
-    unchecked_components = set(
-        [c for c in unchecked_components if not c.exists()])
-    if len(unchecked_components) == 0:
-      return True
-    time.sleep(2)
-
-  raise RuntimeError(
-      'Timeout after %d seconds. %d of %d topics/subscriptions not exist. '
-      'They are %s.' % (timeout, len(unchecked_components),
-                        len(components), list(unchecked_components)))
-
-
-def cleanup_subscriptions(subs):
+def cleanup_subscriptions(sub_client, subs):
   """Cleanup PubSub subscriptions if exist."""
-  _cleanup_pubsub(subs)
+  for sub in subs:
+    sub_client.delete_subscription(sub.name)
 
 
-def cleanup_topics(topics):
+def cleanup_topics(pub_client, topics):
   """Cleanup PubSub topics if exist."""
-  _cleanup_pubsub(topics)
+  for topic in topics:
+    pub_client.delete_topic(topic.name)
 
 
-def _cleanup_pubsub(components):
-  for c in components:
-    if c.exists():
-      c.delete()
-    else:
-      logging.debug('Cannot delete topic/subscription. %s does not exist.',
-                    c.full_name)
+class PullResponseMessage(object):
+  """Data representing a pull request response.
+
+  Utility class for ``create_pull_response``.
+  """
+  def __init__(self, data, attributes=None,
+               publish_time_secs=None, publish_time_nanos=None, ack_id=None):
+    self.data = data
+    self.attributes = attributes
+    self.publish_time_secs = publish_time_secs
+    self.publish_time_nanos = publish_time_nanos
+    self.ack_id = ack_id
+
+
+def create_pull_response(responses):
+  """Create an instance of ``google.cloud.pubsub.types.ReceivedMessage``.
+
+  Used to simulate the response from pubsub.SubscriberClient().pull().
+
+  Args:
+    responses: list of ``PullResponseMessage``
+
+  Returns:
+    An instance of ``google.cloud.pubsub.types.PullResponse`` populated with
+    responses.
+  """
+  from google.cloud import pubsub
+
+  res = pubsub.types.PullResponse()
+  for response in responses:
+    received_message = res.received_messages.add()
+
+    message = received_message.message
+    message.data = response.data
+    if response.attributes is not None:
+      for k, v in response.attributes.items():
+        message.attributes[k] = v
+    if response.publish_time_secs is not None:
+      message.publish_time.seconds = response.publish_time_secs
+    if response.publish_time_nanos is not None:
+      message.publish_time.nanos = response.publish_time_nanos
+
+    if response.ack_id is not None:
+      received_message.ack_id = response.ack_id
+
+  return res
diff --git a/sdks/python/apache_beam/testing/test_utils_test.py b/sdks/python/apache_beam/testing/test_utils_test.py
index bef4078..2b16c30 100644
--- a/sdks/python/apache_beam/testing/test_utils_test.py
+++ b/sdks/python/apache_beam/testing/test_utils_test.py
@@ -82,56 +82,19 @@
         self.assertEqual(f.readline(), b'line2\n')
         self.assertEqual(f.readline(), b'line3\n')
 
-  @mock.patch('time.sleep', return_value=None)
-  def test_wait_for_subscriptions_created_fails(self, patched_time_sleep):
-    sub1 = mock.MagicMock()
-    sub1.exists.return_value = True
-    sub2 = mock.MagicMock()
-    sub2.exists.return_value = False
-    with self.assertRaises(RuntimeError) as error:
-      utils.wait_for_subscriptions_created([sub1, sub2], timeout=0.1)
-    self.assertTrue(sub1.exists.called)
-    self.assertTrue(sub2.exists.called)
-    self.assertTrue(error.exception.args[0].startswith('Timeout after'))
-
-  @mock.patch('time.sleep', return_value=None)
-  def test_wait_for_topics_created_fails(self, patched_time_sleep):
-    topic1 = mock.MagicMock()
-    topic1.exists.return_value = True
-    topic2 = mock.MagicMock()
-    topic2.exists.return_value = False
-    with self.assertRaises(RuntimeError) as error:
-      utils.wait_for_subscriptions_created([topic1, topic2], timeout=0.1)
-    self.assertTrue(topic1.exists.called)
-    self.assertTrue(topic2.exists.called)
-    self.assertTrue(error.exception.args[0].startswith('Timeout after'))
-
-  @mock.patch('time.sleep', return_value=None)
-  def test_wait_for_subscriptions_created_succeeds(self, patched_time_sleep):
-    sub1 = mock.MagicMock()
-    sub1.exists.return_value = True
-    self.assertTrue(
-        utils.wait_for_subscriptions_created([sub1], timeout=0.1))
-
-  @mock.patch('time.sleep', return_value=None)
-  def test_wait_for_topics_created_succeeds(self, patched_time_sleep):
-    topic1 = mock.MagicMock()
-    topic1.exists.return_value = True
-    self.assertTrue(
-        utils.wait_for_subscriptions_created([topic1], timeout=0.1))
-    self.assertTrue(topic1.exists.called)
-
   def test_cleanup_subscriptions(self):
-    mock_sub = mock.MagicMock()
-    mock_sub.exist.return_value = True
-    utils.cleanup_subscriptions([mock_sub])
-    self.assertTrue(mock_sub.delete.called)
+    sub_client = mock.Mock()
+    sub = mock.Mock()
+    sub.name = 'test_sub'
+    utils.cleanup_subscriptions(sub_client, [sub])
+    sub_client.delete_subscription.assert_called_with(sub.name)
 
   def test_cleanup_topics(self):
-    mock_topics = mock.MagicMock()
-    mock_topics.exist.return_value = True
-    utils.cleanup_subscriptions([mock_topics])
-    self.assertTrue(mock_topics.delete.called)
+    pub_client = mock.Mock()
+    topic = mock.Mock()
+    topic.name = 'test_topic'
+    utils.cleanup_topics(pub_client, [topic])
+    pub_client.delete_topic.assert_called_with(topic.name)
 
 
 if __name__ == '__main__':
diff --git a/sdks/python/container/base_image_requirements.txt b/sdks/python/container/base_image_requirements.txt
index a32bc44..8e352a6 100644
--- a/sdks/python/container/base_image_requirements.txt
+++ b/sdks/python/container/base_image_requirements.txt
@@ -47,10 +47,9 @@
 # GCP extra features
 google-apitools==0.5.20
 googledatastore==7.0.1
-google-cloud-pubsub==0.26.0
+google-cloud-pubsub==0.35.4
 google-cloud-bigquery==0.25.0
 proto-google-cloud-datastore-v1==0.90.4
-proto-google-cloud-pubsub-v1==0.15.4
 
 # Optional packages
 cython==0.28.1
diff --git a/sdks/python/setup.py b/sdks/python/setup.py
index 9ee4bdf..ce94945 100644
--- a/sdks/python/setup.py
+++ b/sdks/python/setup.py
@@ -141,8 +141,7 @@
     'google-apitools>=0.5.18,<=0.5.20',
     'proto-google-cloud-datastore-v1>=0.90.0,<=0.90.4',
     'googledatastore==7.0.1',
-    'google-cloud-pubsub==0.26.0',
-    'proto-google-cloud-pubsub-v1==0.15.4',
+    'google-cloud-pubsub==0.35.4',
     # GCP packages required by tests
     'google-cloud-bigquery==0.25.0',
 ]