Python unittest中如何为PySpark链式调用打补丁返回指定值？

你模拟不成功的核心是完全搞错了collect()方法的返回值类型，把调用链路给中断了。
原函数的完整执行逻辑是顺着一层一层的方法调用走的：

1. spark.sql('query') 调用完，拿到一个DataFrame对象
2. 对这个DataFrame调用.select('value')，又拿到一个新的DataFrame对象
3. 对这个新的DataFrame调用.collect()，返回的是pyspark.sql.Row对象组成的列表，不是DataFrame本身
4. 对这个列表用[0]取首元素，拿到对应的Row对象
5. 对Row对象用['value']按键取值，最终得到赋值给sql_output的值

你之前两次尝试的错配：

• 把collect()的返回值设为None：后面紧接着要对这个返回值执行[0]索引，但None不支持下标取值，直接报类型错误
• 把collect()的返回值设为自己构造的测试DataFrame：真实PySpark逻辑里collect()从来不会返回DataFrame，你传进去的DataFrame本身也不支持用整数下标取某一行，自然走不通后面的取值流程

顺着调用链路逐层匹配返回值类型就行，完全不需要启动真实的Spark环境：

轻量模拟版本 #

Row对象本身只支持['key']式的按键取值，普通Python字典就能完美模拟这个行为，写法最省事：

from unittest.mock import patch

@patch("my_function.spark.sql")
def test_my_function(self, mock_spark_sql):
    # 逐层匹配链式调用的返回值
    # 对应spark.sql()返回的DataFrame
    mock_df_after_sql = mock_spark_sql.return_value
    # 对应上述DataFrame调用.select()后返回的DataFrame
    mock_df_after_select = mock_df_after_sql.select.return_value
    # 对应上述DataFrame调用.collect()后返回的行列表
    # 用字典模拟Row对象，给value键赋值None
    mock_df_after_select.collect.return_value = [{"value": None}]

跑通这个测试的时候，原函数里的sql_output就会被正确赋值为None，不会触发任何索引报错。如果需要模拟其他类型的返回值，只需要修改列表里字典的value就行，比如要模拟返回值是"test_str"，就把collect的返回值改成[{"value": "test_str"}]。

贴合真实运行环境的模拟版本 #

如果想要和线上运行时的返回类型完全一致，不需要启动Spark集群，直接导入原生Row类实例化就可以：

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

@patch("my_function.spark.sql")
def test_my_function(self, mock_spark_sql):
    mock_df_after_sql = mock_spark_sql.return_value
    mock_df_after_select = mock_df_after_sql.select.return_value
    # 用原生Row对象构造列表，和真实collect()的返回值完全一致
    mock_df_after_select.collect.return_value = [Row(value=None)]

• 做多级链式调用的Mock时，一定要顺着实际代码的调用顺序，逐次给每一层方法的return_value匹配对应的真实类型，不能跨层赋值
• 别给collect()返回DataFrame对象，这完全和PySpark的API设计相悖，就算你造了完整的测试DataFrame，也接不住后续的下标索引流程

其实不用显式拆分每一层的中间Mock对象，MagicMock本身支持自动生成链式调用的子Mock，只要在collect对应的返回值位置设置成支持下标的列表就行，本质逻辑和上面的写法没有区别。

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