如何通过Mock/Patch修改内部函数入参且保留原函数执行逻辑

问题背景 #

简化示例 #

有如下嵌套调用的代码：

# my_module.my_submodule
from some_library import some_module as x

def do_thing(a, b=None):
    return a + x.random_number(b)

# my_module.my_other_submodule
from my_module.my_submodule import do_thing

def do_more_complex_thing(a):
    # 此处不传入参数b
    return do_thing(a)

测试时需要确保x.random_number(b)接收特定参数，但不想模拟返回值，要保留原函数逻辑。

实际Spark场景 #

在深层嵌套代码中调用pyspark.sql.functions.rand(seed)，需要固定seed保证测试的确定性，且无法修改函数签名传递seed：

# my_module.spark
import pyspark.sql.functions as f

def do_spark_thing():
    ...
    a = 1  # 无法通过方法签名设置该值
    f.rand(a)

之前的错误尝试 #

尝试用自定义函数patch时出现递归错误，原因是自定义函数内部又调用了被patch后的函数：

# test.py
import pyspark.sql.functions as f
from my_module.spark import do_spark_thing
from unittest.mock import patch

def test_do_more_complex_thing():
    def _set_seed(*args, **kwargs):
        return f.rand(1)  # 这里调用的是已经被patch成_set_seed的函数，导致递归

    with patch('my_module.spark.f.rand', _set_seed):
        do_spark_thing()  # 触发无限递归

解决方案 #

核心思路是提前保存原函数的引用，在patch函数中调用原函数并修改入参，避免递归。

方法1：保存原函数后自定义patch逻辑 #

import pyspark.sql.functions as f
from my_module.spark import do_spark_thing
from unittest.mock import patch

def test_do_spark_thing():
    # 提前保存原rand函数的引用，不受后续patch影响
    original_rand = f.rand

    def _override_seed(*args, **kwargs):
        # 强制将seed替换为固定值（比如42），保留原函数逻辑
        return original_rand(42)

    # 注意patch的路径是被测试代码中使用rand的路径
    with patch('my_module.spark.f.rand', _override_seed):
        result = do_spark_thing()
        # 添加断言验证确定性结果，比如：
        # assert result == expected_deterministic_value

方法2：使用Mock的side_effect包装原函数 #

import pyspark.sql.functions as f
from my_module.spark import do_spark_thing
from unittest.mock import patch, Mock

def test_do_spark_thing():
    original_rand = f.rand
    # 用side_effect自定义入参，调用原函数
    mock_rand = Mock(side_effect=lambda *args, **kwargs: original_rand(42))

    with patch('my_module.spark.f.rand', mock_rand):
        do_spark_thing()
        # 可选：验证rand被调用的次数或参数
        # mock_rand.assert_called_once()

关键注意事项 #

• 必须在patch前保存原函数：如果在patch之后获取原函数，拿到的是被替换后的函数，依然会触发递归。
• patch路径要正确：必须patch被测试代码中实际导入并使用目标函数的路径（比如my_module.spark.f.rand），而不是原库的路径（pyspark.sql.functions.rand），这是mock的核心规则——patch对象被使用的位置。

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