fix: Catch and surface BaseException() (#1108)

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
googleapis · Mar 6, 2024 · 07e427f · 07e427f
1 parent 08c5149
commit 07e427f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py b/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py
@@ -112,7 +112,7 @@ def _wrap_as_exception(maybe_exception: Any) -> BaseException:
 
 def _wrap_callback_errors(
     callback: Callable[["google.cloud.pubsub_v1.subscriber.message.Message"], Any],
-    on_callback_error: Callable[[Exception], Any],
+    on_callback_error: Callable[[BaseException], Any],
     message: "google.cloud.pubsub_v1.subscriber.message.Message",
 ):
     """Wraps a user callback so that if an exception occurs the message is
@@ -124,7 +124,7 @@ def _wrap_callback_errors(
     """
     try:
         callback(message)
-    except Exception as exc:
+    except BaseException as exc:
         # Note: the likelihood of this failing is extremely low. This just adds
         # a message to a queue, so if this doesn't work the world is in an
         # unrecoverable state and this thread should just bail.

diff --git a/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py b/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py
@@ -77,9 +77,14 @@ def test__wrap_callback_errors_no_error():
     on_callback_error.assert_not_called()
 
 
-def test__wrap_callback_errors_error():
-    callback_error = ValueError("meep")
-
+@pytest.mark.parametrize(
+    "callback_error",
+    [
+        (ValueError("ValueError")),
+        (BaseException("BaseException")),
+    ],
+)
+def test__wrap_callback_errors_error(callback_error):
     msg = mock.create_autospec(message.Message, instance=True)
     callback = mock.Mock(side_effect=callback_error)
     on_callback_error = mock.Mock()