feat(pubsub): add stop method

pubsub/google/cloud/pubsub_v1/futures.py

@@ -74,9 +74,7 @@ def running(self):
             bool: ``True`` if this method has not yet completed, or
                 ``False`` if it has completed.
         """
-        if self.done():
-            return False
-        return True
+        return not self.done()
 
     def done(self):
         """Return True the future is done, False otherwise.

pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py

@@ -75,6 +75,9 @@ def __init__(self, client, topic, settings, autocommit=True):
         # These members are all communicated between threads; ensure that
         # any writes to them use the "state lock" to remain atomic.
         self._futures = []
+        # _futures list should remain unchanged after batch
+        # status changed from ACCEPTING_MESSAGES to any other
+        # in order to avoid race conditions
         self._messages = []
         self._size = 0
         self._status = base.BatchStatus.ACCEPTING_MESSAGES
@@ -137,6 +140,17 @@ def status(self):
         """
         return self._status
 
+    def wait(self):
+        """If commit is in progress, waits until all of the futures resolved.
+
+        .. note::
+
+            This method blocks until all futures of this batch resolved.
+        """
+        if self._status != base.BatchStatus.ACCEPTING_MESSAGES:
+            for future in self._futures:
+                future.result()
+
     def commit(self):
         """Actually publish all of the messages on the active batch.
 

pubsub/google/cloud/pubsub_v1/publisher/client.py

@@ -274,3 +274,22 @@ def publish(self, topic, data, **attrs):
                 batch = self._batch(topic, create=True)
 
         return future
+
+    def stop(self):
+        """Immediately publish all outstanding batches.
+
+        This asynchronously pushes all outstanding messages
+        and waits until all futures resolved. Method should be
+        invoked prior to deleting this Client object in order
+        to ensure that no pending messages are lost.
+
+        .. note::
+
+            This method blocks until all futures of all
+            batches resolved.
+        """
+        for topic in self._batches:
+            self._batches[topic].commit()
+
+        for topic in self._batches:
+            self._batches[topic].wait()

pubsub/tests/unit/pubsub_v1/publisher/batch/test_thread.py

@@ -252,6 +252,31 @@ def test_block__commmit_api_error():
         assert future.exception() == error
 
 
+def test_wait():
+    batch = create_batch()
+    with mock.patch(
+        "google.cloud.pubsub_v1.publisher.futures.Future.result", return_value=True
+    ) as result_mock:
+        futures = (batch.publish({"data": b"msg"}),)
+
+        batch._status = BatchStatus.IN_PROGRESS
+        futures[0]._completed.set()
+        batch.wait()
+
+        result_mock.assert_called()
+
+
+def test_wait_not_started():
+    batch = create_batch()
+    with mock.patch(
+        "google.cloud.pubsub_v1.publisher.futures.Future.result", return_value=True
+    ) as result_mock:
+        batch.publish({"data": b"msg"})
+        batch.wait()
+
+        result_mock.assert_not_called()
+
+
 def test_block__commmit_retry_error():
     batch = create_batch()
     futures = (

pubsub/tests/unit/pubsub_v1/publisher/test_publisher_client.py

@@ -201,6 +201,35 @@ def test_publish_attrs_type_error():
         client.publish(topic, b"foo", answer=42)
 
 
+def test_stop():
+    creds = mock.Mock(spec=credentials.Credentials)
+    client = publisher.Client(credentials=creds)
+
+    batch = client._batch("topic1", autocommit=False)
+    batch2 = client._batch("topic2", autocommit=False)
+
+    pubsub_msg = types.PubsubMessage(data=b"msg")
+
+    cp = mock.patch.object(batch, "commit")
+    wp = mock.patch.object(batch, "wait")
+    cp2 = mock.patch.object(batch2, "commit")
+    wp2 = mock.patch.object(batch2, "wait")
+
+    with cp as c_mock, cp2 as c_mock2, wp as w_mock, wp2 as w_mock2:
+        batch.publish(pubsub_msg)
+        batch2.publish(pubsub_msg)
+
+        client.stop()
+
+        # check if commit() called
+        c_mock.assert_called()
+        c_mock2.assert_called()
+
+        # check if wait() called
+        w_mock.assert_called()
+        w_mock2.assert_called()
+
+
 def test_gapic_instance_method():
     creds = mock.Mock(spec=credentials.Credentials)
     client = publisher.Client(credentials=creds)