.

Author: sentrivana

sentry_sdk/integrations/spark/spark_driver.py

@@ -296,16 +296,14 @@ def onStageCompleted(self, stageCompleted):  # noqa: N802,N803
 
 def _get_attempt_id(stage_info):
     # type: (Any) -> Optional[int]
-    from py4j.protocol import Py4JJavaError
-
     try:
         return stage_info.attemptId()
-    except Py4JJavaError:
+    except Exception:
         pass
 
     try:
         return stage_info.attemptNumber()
-    except Py4JJavaError:
+    except Exception:
         pass
 
     return None

tests/integrations/spark/test_spark.py

@@ -14,6 +14,7 @@
 
 from py4j.protocol import Py4JJavaError
 
+
 ################
 # DRIVER TESTS #
 ################
@@ -166,6 +167,65 @@ def stageInfo(self):  # noqa: N802
         assert mock_hub.kwargs["data"]["name"] == "run-job"
 
 
+def test_sentry_listener_on_stage_submitted_no_attempt_id(sentry_listener):
+    listener = sentry_listener
+    with patch.object(listener, "_add_breadcrumb") as mock_add_breadcrumb:
+
+        class StageInfo:
+            def stageId(self):  # noqa: N802
+                return "sample-stage-id-submit"
+
+            def name(self):
+                return "run-job"
+
+            def attemptNumber(self):  # noqa: N802
+                return 14
+
+        class MockStageSubmitted:
+            def stageInfo(self):  # noqa: N802
+                stageinf = StageInfo()
+                return stageinf
+
+        mock_stage_submitted = MockStageSubmitted()
+        listener.onStageSubmitted(mock_stage_submitted)
+
+        mock_add_breadcrumb.assert_called_once()
+        mock_hub = mock_add_breadcrumb.call_args
+
+        assert mock_hub.kwargs["level"] == "info"
+        assert "sample-stage-id-submit" in mock_hub.kwargs["message"]
+        assert mock_hub.kwargs["data"]["attemptId"] == 14
+        assert mock_hub.kwargs["data"]["name"] == "run-job"
+
+
+def test_sentry_listener_on_stage_submitted_no_attempt_id_or_number(sentry_listener):
+    listener = sentry_listener
+    with patch.object(listener, "_add_breadcrumb") as mock_add_breadcrumb:
+
+        class StageInfo:
+            def stageId(self):  # noqa: N802
+                return "sample-stage-id-submit"
+
+            def name(self):
+                return "run-job"
+
+        class MockStageSubmitted:
+            def stageInfo(self):  # noqa: N802
+                stageinf = StageInfo()
+                return stageinf
+
+        mock_stage_submitted = MockStageSubmitted()
+        listener.onStageSubmitted(mock_stage_submitted)
+
+        mock_add_breadcrumb.assert_called_once()
+        mock_hub = mock_add_breadcrumb.call_args
+
+        assert mock_hub.kwargs["level"] == "info"
+        assert "sample-stage-id-submit" in mock_hub.kwargs["message"]
+        assert "attemptId" not in mock_hub.kwargs["data"]
+        assert mock_hub.kwargs["data"]["name"] == "run-job"
+
+
 @pytest.fixture
 def get_mock_stage_completed():
     def _inner(failure_reason):
@@ -282,71 +342,3 @@ def mock_main():
         "partitionId": "2",
         "taskAttemptId": "3",
     }
-
-
-def test_sentry_listener_on_stage_submitted_no_attempt_id(sentry_listener):
-    listener = sentry_listener
-    with patch.object(listener, "_add_breadcrumb") as mock_add_breadcrumb:
-
-        class StageInfo:
-            def stageId(self):  # noqa: N802
-                return "sample-stage-id-submit"
-
-            def name(self):
-                return "run-job"
-
-            def attemptId(self):  # noqa: N802
-                raise Py4JJavaError
-
-            def attemptNumber(self):  # noqa: N802
-                return 14
-
-        class MockStageSubmitted:
-            def stageInfo(self):  # noqa: N802
-                stageinf = StageInfo()
-                return stageinf
-
-        mock_stage_submitted = MockStageSubmitted()
-        listener.onStageSubmitted(mock_stage_submitted)
-
-        mock_add_breadcrumb.assert_called_once()
-        mock_hub = mock_add_breadcrumb.call_args
-
-        assert mock_hub.kwargs["level"] == "info"
-        assert "sample-stage-id-submit" in mock_hub.kwargs["message"]
-        assert mock_hub.kwargs["data"]["attemptId"] == 14
-        assert mock_hub.kwargs["data"]["name"] == "run-job"
-
-
-def test_sentry_listener_on_stage_submitted_no_attempt_id_or_number(sentry_listener):
-    listener = sentry_listener
-    with patch.object(listener, "_add_breadcrumb") as mock_add_breadcrumb:
-
-        class StageInfo:
-            def stageId(self):  # noqa: N802
-                return "sample-stage-id-submit"
-
-            def name(self):
-                return "run-job"
-
-            def attemptId(self):  # noqa: N802
-                raise Py4JJavaError
-
-            def attemptNumber(self):  # noqa: N802
-                raise Py4JJavaError
-
-        class MockStageSubmitted:
-            def stageInfo(self):  # noqa: N802
-                stageinf = StageInfo()
-                return stageinf
-
-        mock_stage_submitted = MockStageSubmitted()
-        listener.onStageSubmitted(mock_stage_submitted)
-
-        mock_add_breadcrumb.assert_called_once()
-        mock_hub = mock_add_breadcrumb.call_args
-
-        assert mock_hub.kwargs["level"] == "info"
-        assert "sample-stage-id-submit" in mock_hub.kwargs["message"]
-        assert "attemptId" not in mock_hub.kwargs["data"]
-        assert mock_hub.kwargs["data"]["name"] == "run-job"