如何正确测试调用pandas.read_sql的PySpark自定义read_data函数

我们的测试目标是验证read_data函数的两个核心逻辑，不需要真实连接数据库、启动Spark集群，完全通过mock依赖实现单元测试，运行速度更快且不受外部依赖影响：

• 验证pandas.read_sql使用了预期的SQL语句、数据库连接参数执行查询
• 验证read_sql返回的pandas DataFrame被正确传入spark.createDataFrame，且函数返回值就是createDataFrame生成的Spark DataFrame

import pandas as pd
from unittest import mock
# 此处替换为你自己的函数导入路径
from your_module import read_data, conn

# 注意patch路径要匹配read_data函数所在模块引用pd.read_sql的路径，不可直接patch全局pandas.read_sql
@mock.patch("your_module.pd.read_sql")
@mock.patch("pyspark.sql.SparkSession", autospec=True)
def test_read_data(spark_session, pandas_read_sql):
    # 构造mock返回值：模拟pandas.read_sql返回的测试DataFrame
    mock_pd_df = pd.DataFrame({"id": [1,2,3], "value": ["test1", "test2", "test3"]})
    pandas_read_sql.return_value = mock_pd_df

    # 构造mock返回值：模拟spark.createDataFrame返回的Spark DataFrame
    mock_spark_df = mock.Mock()
    spark_session.createDataFrame.return_value = mock_spark_df

    # 执行测试函数
    result = read_data(spark_session)

    # 断言1：验证pandas.read_sql调用参数符合预期
    pandas_read_sql.assert_called_once_with("SELECT * FROM table", conn)
    # 如果无法导入全局conn，也可以只校验SQL参数：
    # assert pandas_read_sql.call_args.args[0] == "SELECT * FROM table"

    # 断言2：验证spark.createDataFrame被传入了read_sql返回的pandas DataFrame
    spark_session.createDataFrame.assert_called_once_with(mock_pd_df)

    # 断言3：验证函数返回值和预期一致
    assert result == mock_spark_df

补充说明：如果你的函数逻辑后续有修改，比如SQL改为动态拼接、转换逻辑新增字段处理，只需要对应调整mock入参和断言逻辑即可。

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