saurabh18cs
Honored Contributor III

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()