如何为使用spark.sql的函数编写单元测试？解决Jenkins报错问题

问题场景 #

在Databricks环境配置Unity Catalog后，部分函数通过spark.sql查询数据表获取数据。编写单元测试时，在Databricks本地可正常执行，但部署到Jenkins服务器运行时，抛出**"spark未定义"**错误。现有测试代码通过setUp创建SparkSession并赋值给全局spark，原函数直接使用全局spark执行SQL查询。

根本原因 #

原函数硬编码依赖全局spark变量：Databricks环境默认初始化全局SparkSession，因此本地测试正常；但Jenkins测试环境中，全局变量的作用域或初始化时机问题，导致函数执行时无法找到spark实例。同时，直接依赖全局变量的代码耦合性高，不利于单元测试的隔离。

解决方案 #

1. 重构原函数，解耦Spark依赖 #

修改final_flag_generate函数，将SparkSession作为参数传入，避免依赖全局变量：

from pyspark.sql.functions import col

def final_flag_generate(spark, joined_sql, catalog_name, val_dt, iteration):
    mapping_id_list = [row.MAPPING_ID for row in joined_sql.collect()]
    id_tuple = tuple(mapping_id_list)
    print(id_tuple)
    
    final_calculation_df = spark.sql(f"""select * from {catalog_name}.db1.tbl1 
                                         where mapping_id in {id_tuple} 
                                           and val_dt='{val_dt}' 
                                           and iteration={iteration} """)

    count_ones = final_calculation_df.filter(col("balance_flag") == 1).count()
    print(count_ones)
    total_count = final_calculation_df.count()
    print(total_count)

    final_result_flag = 1 if count_ones == total_count else 0
    return final_result_flag, final_calculation_df

2. 修正单元测试，传递SparkSession并正确Mock #

调整测试用例，将setUp中创建的SparkSession传入函数，并确保mock_sql正确绑定到该SparkSession的sql方法：

import unittest
from pyspark.sql import SparkSession
from unittest.mock import patch

class TestMyFunction(unittest.TestCase):
    def setUp(self):
        self.spark = SparkSession.builder.master("local").appName("test").getOrCreate()
        # 初始化测试用的DataFrame
        self.final_flag_generate_df = self.spark.createDataFrame(
            [(1, 1), (2, 1)], ["MAPPING_ID", "balance_flag"]
        )
        # Mock spark.sql方法
        self.mock_sql = patch.object(self.spark, "sql").start()

    def tearDown(self):
        patch.stopall()
        self.spark.stop()

    def test_final_flag_generate(self):
        # 设置mock返回值
        self.mock_sql.return_value = self.final_flag_generate_df
        catalog_name = "test_catalog"
        val_dt = "2022-01-01"
        iteration = 1
        
        # 传入spark实例到函数
        result_flag, result_df = final_flag_generate(
            self.spark, self.final_flag_generate_df, catalog_name, val_dt, iteration
        )
        
        expected_flag = 1
        expected_columns = self.final_flag_generate_df.columns
        
        self.assertEqual(expected_flag, result_flag)
        self.assertListEqual(expected_columns, result_df.columns)
        # 验证sql方法是否被正确调用（可选）
        self.mock_sql.assert_called_once()

关键优化点 #

• 解耦全局依赖：函数不再依赖全局spark，改为显式传入SparkSession，提升代码可测试性和复用性。
• 隔离Mock：直接Mock测试用例中创建的SparkSession的sql方法，避免全局变量的作用域问题。
• 清理资源：在tearDown中停止SparkSession，避免测试资源泄漏。

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