Skip to content

Commit

Permalink
Add: tests cases for before ad after SparkContext creation
Browse files Browse the repository at this point in the history
  • Loading branch information
seyoon-lim committed Aug 7, 2024
1 parent 74edf67 commit 4c673d6
Showing 1 changed file with 44 additions and 8 deletions.
52 changes: 44 additions & 8 deletions tests/integrations/spark/test_spark.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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"
Expand All @@ -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
Expand Down

0 comments on commit 4c673d6

Please sign in to comment.