Options
- Mark as New
- Bookmark
- Subscribe
- Mute
- Subscribe to RSS Feed
- Permalink
- Report Inappropriate Content
09-19-2025 02:51 AM
Hi Utu,
try doing something like this , wrap import within a fixture itself.
import os
import pytest
from pyspark.sql import SparkSession
_local_test = True
@pytest.fixture(scope='session')
def spark():
if 'DATABRICKS_RUNTIME_VERSION' in os.environ:
# Running on Databricks
print("Running on Databricks with SparkSession")
spark_session = SparkSession.builder.getOrCreate()
yield spark_session
spark_session.stop()
elif _local_test:
spark_session = SparkSession.builder.master("local[1]").getOrCreate()
yield spark_session
spark_session.stop()
else:
# Running outside Databricks, use Databricks Connect
print("Running outside Databricks with Databricks Connect")
from databricks.connect import DatabricksSession
from databricks.sdk.core import Config
spark_session = DatabricksSession.builder.remote(
cluster_id="0805-092156",
token=""
).getOrCreate()
yield spark_session
spark_session.stop()