fix(pubsub): handle None in on response callback (#9982)

plamut · pradn · commit d372f66168ff · 2019-12-16T19:26:18.000-03:00
If the underlying RPC is shut down while pulling the messages with a
streming pull, the StreamingPullManager's _on_response() method is
invoked with None (as opposed to a StreamingPullResponse instance).

This commit handles this case and prevents an error in a background
thread on streaming pull manager shutdown.
diff --git a/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py b/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py
@@ -542,6 +542,13 @@ def _on_response(self, response):
         After the messages have all had their ack deadline updated, execute
         the callback for each message using the executor.
         """
+        if response is None:
+            _LOGGER.debug(
+                "Response callback invoked with None, likely due to a "
+                "transport shutdown."
+            )
+            return
+
         _LOGGER.debug(
             "Processing %s received message(s), currenty on hold %s (bytes %s).",
             len(response.received_messages),
diff --git a/pubsub/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py b/pubsub/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py
@@ -721,6 +721,21 @@ def test__on_response_with_leaser_overload():
             assert msg.message_id in ("2", "3")
 
 
+def test__on_response_none_data(caplog):
+    caplog.set_level(logging.DEBUG)
+
+    manager, _, dispatcher, leaser, _, scheduler = make_running_manager()
+    manager._callback = mock.sentinel.callback
+
+    # adjust message bookkeeping in leaser
+    fake_leaser_add(leaser, init_msg_count=0, assumed_msg_size=10)
+
+    manager._on_response(response=None)
+
+    scheduler.schedule.assert_not_called()
+    assert "callback invoked with None" in caplog.text
+
+
 def test_retryable_stream_errors():
     # Make sure the config matches our hard-coded tuple of exceptions.
     interfaces = subscriber_client_config.config["interfaces"]
