Add: tests cases for before ad after SparkContext creation

Author: seyoon-lim

Diff for: tests/integrations/spark/test_spark.py

@@ -1,6 +1,7 @@
 import pytest
 import sys
 from unittest.mock import patch
+
 from sentry_sdk.integrations.spark.spark_driver import (
     _set_app_properties,
     _start_sentry_listener,
@@ -18,8 +19,15 @@
 ################
 
 
-def test_set_app_properties():
-    spark_context = SparkContext(appName="Testing123")
+@pytest.fixture(scope="function")
+def create_spark_context():
+    yield lambda: SparkContext(appName="Testing123")
+    if SparkContext._active_spark_context:
+        SparkContext._active_spark_context.stop()
+
+
+def test_set_app_properties(create_spark_context):
+    spark_context = create_spark_context()
     _set_app_properties()
 
     assert spark_context.getLocalProperty("sentry_app_name") == "Testing123"
@@ -30,20 +38,48 @@ def test_set_app_properties():
     )
 
 
-def test_start_sentry_listener():
-    spark_context = SparkContext.getOrCreate()
-
+def test_start_sentry_listener(create_spark_context):
+    spark_context = create_spark_context()
     gateway = spark_context._gateway
     assert gateway._callback_server is None
 
-    _start_sentry_listener(spark_context)
+    _start_sentry_listener()
 
     assert gateway._callback_server is not None
 
 
-def test_initialize_spark_integration(sentry_init):
+def test_initialize_spark_integration(
+    sentry_init, create_spark_context, reset_integrations
+):
     sentry_init(integrations=[SparkIntegration()])
-    SparkContext.getOrCreate()
+    create_spark_context()
+
+
+@patch("sentry_sdk.integrations.spark.spark_driver._patch_spark_context_init")
+def test_initialize_spark_integration_before_spark_context_init(
+    mock_patch_spark_context_init,
+    sentry_init,
+    create_spark_context,
+):
+    sentry_init(integrations=[SparkIntegration()])
+    create_spark_context()
+
+    mock_patch_spark_context_init.assert_called_once()
+
+
+@patch("sentry_sdk.integrations.spark.spark_driver._activate_integration")
+@patch("sentry_sdk.integrations.spark.spark_driver._patch_spark_context_init")
+def test_initialize_spark_integration_after_spark_context_init(
+    mock_patch_spark_context_init,
+    mock_activate_integration,
+    create_spark_context,
+    sentry_init,
+):
+    create_spark_context()
+    sentry_init(integrations=[SparkIntegration()])
+
+    mock_activate_integration.assert_called_once()
+    mock_patch_spark_context_init.assert_not_called()
 
 
 @pytest.fixture