fix: Replace asserts with None checks for graceful shutdown (#1244)

Author: mukund-ananthu

google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py

@@ -1104,8 +1104,11 @@ def _on_response(self, response: gapic_types.StreamingPullResponse) -> None:
         )
 
         with self._pause_resume_lock:
-            assert self._scheduler is not None
-            assert self._leaser is not None
+            if self._scheduler is None or self._leaser is None:
+                _LOGGER.debug(
+                    f"self._scheduler={self._scheduler} or self._leaser={self._leaser} is None. Stopping further processing."
+                )
+                return
 
             for received_message in received_messages:
                 if (

tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py

@@ -1375,6 +1375,62 @@ def test_close_blocking_scheduler_shutdown():
     scheduler.shutdown.assert_called_once_with(await_msg_callbacks=True)
 
 
+def test__on_response_none_scheduler():
+    manager, _, _, _, _, _ = make_running_manager()
+
+    manager._callback = mock.sentinel.callback
+    manager._scheduler = None
+    # Set up the messages.
+    response = gapic_types.StreamingPullResponse(
+        received_messages=[
+            gapic_types.ReceivedMessage(
+                ack_id="ack1",
+                message=gapic_types.PubsubMessage(data=b"foo", message_id="1"),
+            ),
+            gapic_types.ReceivedMessage(
+                ack_id="ack2",
+                message=gapic_types.PubsubMessage(data=b"bar", message_id="2"),
+                delivery_attempt=6,
+            ),
+        ]
+    )
+
+    manager._maybe_release_messages = mock.Mock()
+
+    # adjust message bookkeeping in leaser
+    fake_leaser_add(leaser, init_msg_count=0, assumed_msg_size=42)
+    manager._on_response(response)
+
+    manager._maybe_release_messages.assert_not_called
+
+
+def test__on_response_none_leaser():
+    manager, _, _, _, _, _ = make_running_manager()
+
+    manager._callback = mock.sentinel.callback
+    manager._leaser = None
+    # Set up the messages.
+    response = gapic_types.StreamingPullResponse(
+        received_messages=[
+            gapic_types.ReceivedMessage(
+                ack_id="ack1",
+                message=gapic_types.PubsubMessage(data=b"foo", message_id="1"),
+            ),
+            gapic_types.ReceivedMessage(
+                ack_id="ack2",
+                message=gapic_types.PubsubMessage(data=b"bar", message_id="2"),
+                delivery_attempt=6,
+            ),
+        ]
+    )
+
+    manager._maybe_release_messages = mock.Mock()
+
+    manager._on_response(response)
+
+    manager._maybe_release_messages.assert_not_called
+
+
 def test_close_nonblocking_scheduler_shutdown():
     manager, _, _, _, _, _ = make_running_manager(await_callbacks_on_shutdown=False)
     scheduler = manager._scheduler