fix: shutdown error on streaming pull callback error (googleapis#40)

plamut · web-flow · commit 552539e7beb3 · 2020-02-28T11:21:57.000-05:00
* fix: shutdown error on streaming pull callback error

* fix incorrect comment
diff --git a/google/cloud/pubsub_v1/subscriber/futures.py b/google/cloud/pubsub_v1/subscriber/futures.py
@@ -33,6 +33,11 @@ def __init__(self, manager):
         self._cancelled = False
 
     def _on_close_callback(self, manager, result):
+        if self.done():
+            # The future has already been resolved in a different thread,
+            # nothing to do on the streaming pull manager shutdown.
+            return
+
         if result is None:
             self.set_result(True)
         else:
diff --git a/tests/unit/pubsub_v1/subscriber/test_futures_subscriber.py b/tests/unit/pubsub_v1/subscriber/test_futures_subscriber.py
@@ -57,6 +57,18 @@ def test__on_close_callback_failure(self):
 
         assert not future.running()
 
+    def test__on_close_callback_future_already_done(self):
+        future = self.make_future()
+
+        future.set_result("foo")
+        assert future.done()
+
+        # invoking on close callback should not result in an error
+        future._on_close_callback(mock.sentinel.manager, "bar")
+
+        result = future.result()
+        assert result == "foo"  # on close callback was a no-op
+
     def test_cancel(self):
         future = self.make_future()
 
