如何用Mock单元测试创建返回PySpark DataFrame的代码？测试失败排查

1. createDataFrame调用参数断言不匹配 #

原业务代码的get_data方法里，调用spark.createDataFrame(data)时传入的是单个字典{"Name": self.name, "Count": count}，但你测试里写的断言是assert_called_once_with([expected_data])——把字典包在列表里了，两者参数完全不一致，这会直接导致断言失败。

2. 原业务代码存在冗余（顺带说明） #

你原代码里ctrl_data = spark.createDataFrame(data)这行生成的DataFrame根本没被使用，方法直接返回了原始的data字典，所以result == expected_data这个断言是能通过的，但Spark调用相关的断言会因为参数错误失败。

3. Mock写法可优化（可选） #

你直接给SparkSession.builder.getOrCreate赋值Mock的方式虽然能运行，但用unittest.mock.patch来模拟更稳妥，不会污染全局的SparkSession状态，比如用上下文管理器包裹模拟逻辑。

修正后的测试代码示例 #

from unittest.mock import MagicMock, patch
from get_data import Extract

def test_get_data():
    name = "John Doe"
    extract = Extract(NAME=name)
    count = 10
    expected_data = {"Name": name, "Count": count}

    # 用patch模拟SparkSession的获取逻辑
    with patch('get_data.SparkSession.builder.getOrCreate') as mock_get_session:
        spark_mock = MagicMock()
        mock_get_session.return_value = spark_mock

        result = extract.get_data(count)

        # 验证返回值符合预期
        assert result == expected_data
        # 验证SparkSession获取操作被调用一次
        mock_get_session.assert_called_once()
        # 验证createDataFrame的调用参数和原代码一致（单个字典）
        spark_mock.createDataFrame.assert_called_once_with(expected_data)

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