如何通过Pytest Fixture让Spark在单元测试中全局可用？

问题描述 #

我使用Databricks资产包，通过pytest为Notebook中的函数编写单元测试。我通过pytest.fixture创建了Spark会话，但遇到了作用域问题，代码如下：

@pytest.fixture(scope="session")
def spark_session(request):
    spark = get_spark()
    yield spark

当前Spark会话仅在单元测试函数内部可用，当测试用例调用Notebook中的函数时，会报错NameError: name 'spark' is not defined。请问如何在单元测试文件中定义Spark，使其对所有测试调用的函数都可用？

可行解决方案 #

方法1：将Spark会话作为参数传入Notebook函数 #

修改Notebook中的函数，显式接收spark参数，避免依赖全局变量：

# Notebook中的函数示例
def process_data(spark, input_path):
    df = spark.read.csv(input_path)
    # 业务处理逻辑
    return df

测试用例中直接传入fixture生成的spark会话：

def test_process_data(spark_session):
    result_df = process_data(spark_session, "test_input.csv")
    # 断言验证逻辑

方法2：将Spark会话注入Notebook模块的全局命名空间 #

利用pytest的autouse fixture，在测试启动时把spark会话注入到Notebook对应的模块中，无需修改原有函数代码：

import sys
from my_notebook_module import *  # 导入你的Notebook函数所在模块

@pytest.fixture(scope="session", autouse=True)
def setup_global_spark():
    spark = get_spark()
    # 把spark注入到Notebook模块的全局变量中
    sys.modules['my_notebook_module'].spark = spark
    yield spark
    # 测试结束后清理资源
    spark.stop()

这样Notebook中的函数就能直接使用全局的spark变量。

方法3：使用Databricks Connect自动注册全局Spark会话 #

如果测试环境配置了Databricks Connect，它会自动创建一个指向远端Databricks集群的Spark会话，并将其注册为全局spark变量。只需确保测试环境正确配置DATABRICKS_HOST、DATABRICKS_TOKEN等环境变量，Notebook中的函数即可直接使用全局spark。

内容的提问来源于stack exchange，提问作者Kylo