getsentry · szokeasaurusrex · Aug 27, 2024 · Aug 7, 2024 · Aug 9, 2024 · Aug 9, 2024
@@ -0,0 +1,73 @@
+import sys
+
+import sentry_sdk
+from sentry_sdk.utils import (
+    ensure_integration_enabled,
+    capture_internal_exceptions,
+    event_from_exception,
+)
+from sentry_sdk.integrations import Integration
+from sentry_sdk._types import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import NoReturn, Union
+
+
+class SysExitIntegration(Integration):
+    """Captures sys.exit calls and sends them as events to Sentry.
+
+    By default, SystemExit exceptions are not captured by the SDK. Enabling this integration will capture SystemExit
+    exceptions generated by sys.exit calls and send them to Sentry.
+
+    This integration, in its default configuration, only captures the sys.exit call if the exit code is a non-zero and
+    non-None value (unsuccessful exits). Pass `capture_successful_exits=True` to capture successful exits as well.
+    Note that the integration does not capture SystemExit exceptions raised outside a call to sys.exit.
+    """
+
+    identifier = "sys_exit"
+
+    def __init__(self, *, capture_successful_exits=False):
+        # type: (bool) -> None
+        self._capture_successful_exits = capture_successful_exits
+
+    @staticmethod
+    def setup_once():
+        # type: () -> None
+        SysExitIntegration._patch_sys_exit()
+
+    @staticmethod
+    def _patch_sys_exit():
+        # type: () -> None
+        old_exit = sys.exit  # type: Callable[[Union[str, int, None]], NoReturn]
+
+        @ensure_integration_enabled(SysExitIntegration, old_exit)
+        def sentry_patched_exit(__status=0):
+            # type: (Union[str, int, None]) -> NoReturn
+            # @ensure_integration_enabled ensures that this is non-None
+            integration = sentry_sdk.get_client().get_integration(
+                SysExitIntegration
+            )  # type: SysExitIntegration
+
+            try:
+                old_exit(__status)
+            except SystemExit as e:
+                with capture_internal_exceptions():
+                    if integration._capture_successful_exits or __status not in (
+                        0,
+                        None,
+                    ):
+                        _capture_exception(e)
+                raise e
+
+        sys.exit = sentry_patched_exit  # type: ignore
+
+
+def _capture_exception(exc):
+    # type: (SystemExit) -> None
+    event, hint = event_from_exception(
+        exc,
+        client_options=sentry_sdk.get_client().options,
+        mechanism={"type": SysExitIntegration.identifier, "handled": False},
+    )
+    sentry_sdk.capture_event(event, hint=hint)
@@ -0,0 +1,71 @@
+import sys
+
+import pytest
+
+from sentry_sdk.integrations.sys_exit import SysExitIntegration
+
+
+@pytest.mark.parametrize(
+    ("integration_params", "exit_status", "should_capture"),
+    (
+        ({}, 0, False),
+        ({}, 1, True),
+        ({}, None, False),
+        ({}, "unsuccessful exit", True),
+        ({"capture_successful_exits": False}, 0, False),
+        ({"capture_successful_exits": False}, 1, True),
+        ({"capture_successful_exits": False}, None, False),
+        ({"capture_successful_exits": False}, "unsuccessful exit", True),
+        ({"capture_successful_exits": True}, 0, True),
+        ({"capture_successful_exits": True}, 1, True),
+        ({"capture_successful_exits": True}, None, True),
+        ({"capture_successful_exits": True}, "unsuccessful exit", True),
+    ),
+)
+def test_sys_exit(
+    sentry_init, capture_events, integration_params, exit_status, should_capture
+):
+    sentry_init(integrations=[SysExitIntegration(**integration_params)])
+
+    events = capture_events()
+
+    # Manually catch the sys.exit rather than using pytest.raises because IDE does not recognize that pytest.raises
+    # will catch SystemExit.
+    try:
+        sys.exit(exit_status)
+    except SystemExit:
+        ...
+    else:
+        pytest.fail("Patched sys.exit did not raise SystemExit")
+
+    if should_capture:
+        (event,) = events
+        (exception_value,) = event["exception"]["values"]
+
+        assert exception_value["type"] == "SystemExit"
+        assert exception_value["value"] == (
+            str(exit_status) if exit_status is not None else ""
+        )
+    else:
+        assert len(events) == 0
+
+
+def test_sys_exit_integration_not_auto_enabled(sentry_init, capture_events):
+    sentry_init()  # No SysExitIntegration
+
+    events = capture_events()
+
+    # Manually catch the sys.exit rather than using pytest.raises because IDE does not recognize that pytest.raises
+    # will catch SystemExit.
+    try:
+        sys.exit(1)
+    except SystemExit:
+        ...
+    else:
+        pytest.fail(
+            "sys.exit should not be patched, but it must have been because it did not raise SystemExit"
+        )
+
+    assert (
+        len(events) == 0
+    ), "No events should have been captured because sys.exit should not have been patched"