Python Mock进阶：如何让返回的Mock实例存储生成它的函数参数？

问题描述 #

我们需要验证bar.func是否以foo.a(1, 2)和foo.b(3, 4)的返回值作为参数被调用，但默认的MagicMock无法直接关联这两个调用的上下文：

• bar.method_calls仅记录传入的Mock实例，无法获取该实例对应的foo.a/foo.b调用参数
• foo.method_calls的记录与bar.func的调用解耦，无法区分哪些foo.a调用是实际传入bar.func的

默认情况下，foo.a返回的Mock实例不存储创建时的参数，且foo.a本身是复用的Mock对象，多次调用会累积记录，导致无法精准关联到bar.func的调用参数。

自定义Traceable Mock实现 #

通过自定义MagicMock子类，让每个由foo方法调用生成的Mock实例独立记录自身被创建时的调用参数和上下文，解决复用和关联问题：

from unittest.mock import MagicMock, call

class TraceableChildMock(MagicMock):
    def __init__(self, call_args, call_kwargs, method_name=None, parent_mock=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 存储当前实例对应的原调用参数
        self.call_args = call_args
        self.call_kwargs = call_kwargs
        # 存储所属的父Mock和方法名
        self.method_name = method_name
        self.parent_mock = parent_mock

class TraceableMock(MagicMock):
    def __getattr__(self, name):
        # 重写属性获取逻辑，避免复用同一个方法Mock
        # 每次获取foo.a时，返回一个临时的调用包装器
        def method_wrapper(*args, **kwargs):
            # 生成带调用信息的子Mock
            return TraceableChildMock(
                call_args=args,
                call_kwargs=kwargs,
                method_name=name,
                parent_mock=self,
                name=f"{self._mock_name}.{name}()"
            )
        # 用MagicMock包装这个方法，保持Mock的其他特性
        method_mock = MagicMock(side_effect=method_wrapper)
        setattr(self, name, method_mock)
        return method_mock

使用示例与验证 #

业务代码无需修改，直接使用自定义的TraceableMock替代普通MagicMock：

# 初始化自定义Mock
foo = TraceableMock()
bar = MagicMock()

# 业务代码逻辑（完全无感知Mock实现）
a = foo.a(1, 2)
a2 = foo.b(5, 6)  # 无关的额外调用
b = foo.b(3, 4)
bar.func(a, b)

# 验证bar.func的调用参数是否符合预期
# 验证第一个参数是foo.a(1,2)的返回值
first_arg = bar.func.call_args.args[0]
assert first_arg.method_name == 'a'
assert first_arg.call_args == (1, 2)
assert first_arg.parent_mock is foo

# 验证第二个参数是foo.b(3,4)的返回值
second_arg = bar.func.call_args.args[1]
assert second_arg.method_name == 'b'
assert second_arg.call_args == (3, 4)
assert second_arg.parent_mock is foo

方案优势 #

• 业务代码无侵入：原有的foo.a(1,2)、bar.func(a,b)逻辑完全不需要修改
• 精准关联上下文：每个传入bar.func的Mock实例都独立记录自己的创建参数，不受同方法其他调用的干扰
• 保留Mock特性：自定义Mock继承自MagicMock，依然支持普通Mock的所有功能（如断言调用次数、设置返回值等）

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