如何在类方法中使用mock.patch结合wraps？

场景复现 #

普通函数模拟成功示例 #

from unittest import mock

def f(x, y):
    return x * y

print(f(2, 3))  # output: 6

original_f = f
def m(x, y):
    return original_f(x, y + 1)

with mock.patch('__main__.f', wraps=m) as mk:
    print(f(2, 3))  # output: 8

类方法模拟失败示例 #

from unittest import mock

class A:
    def __init__(self, x):
        self.x = x
    def f(self, y):
        return self.x * y

a = A(2)

print(a.f(3))  # output: 6

original_f = A.f
def m(self, y):
    return original_f(self, y + 1)

with mock.patch('__main__.A.f', wraps=m) as mk:
    print(a.f(3))  # Exception

报错信息

Traceback (most recent call last):
  File "a.py", line 18, in <module>
    print(a.f(3))
          ^^^^^^
  File "/usr/lib/python3.11/unittest/mock.py", line 1118, in __call__
    return self._mock_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.11/unittest/mock.py", line 1122, in _mock_call
    return self._execute_mock_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.11/unittest/mock.py", line 1192, in _execute_mock_call
    return self._mock_wraps(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: m() missing 1 required positional argument: 'y'

现象显示：mock错误地将3作为self参数传入m，导致y参数缺失，引发异常。

补充复杂场景 #

编辑补充：
以下示例用于说明两个限制条件：

• 无法修改a.f(3)这类调用代码；
• 测试函数run_test无法访问test_target内部创建的实例a、b。

from unittest import mock

class A:
    # 此类代码不可修改
    def __init__(self, x):
        self.x = x
    def f(self, y):
        return self.x * y

def test_target():
    # 此函数代码不可修改
    a = A(2)
    b = A(3)
    return a.f(3), b.f(3)

def run_test():
    print(test_target())  # output: (6, 9)

    original_f = A.f
    def m(self, y):
        return original_f(self, y + 1)

    with mock.patch('__main__.A.f', wraps=m) as mk:
        print(test_target())  # Exception

if __name__ == '__main__':
    run_test()

原因分析 #

当用mock.patch替换类方法时，wraps参数没有正确处理Python的方法绑定机制：

• 正常调用实例方法a.f(3)时，Python会自动把实例a作为第一个参数self传递给方法；
• 但使用wraps=m时，mock将m视为普通函数，直接把调用时的参数原封不动传递，导致3被错误当作self，y参数缺失。

本质是wraps没有把m绑定为类方法，破坏了Python的方法调用逻辑。

修复方案 #

方案1：改用side_effect替代wraps（推荐） #

side_effect会正确接收方法调用的所有参数（包括自动传递的self），并转发给自定义函数执行，完美避开方法绑定问题：

针对基础失败示例的修复：

from unittest import mock

class A:
    def __init__(self, x):
        self.x = x
    def f(self, y):
        return self.x * y

a = A(2)

print(a.f(3))  # output: 6

original_f = A.f
def m(self, y):
    return original_f(self, y + 1)

with mock.patch('__main__.A.f', side_effect=m) as mk:
    print(a.f(3))  # output: 8

针对复杂补充场景的修复：

def run_test():
    print(test_target())  # output: (6, 9)

    original_f = A.f
    def m(self, y):
        return original_f(self, y + 1)

    with mock.patch('__main__.A.f', side_effect=m) as mk:
        print(test_target())  # output: (8, 12)

方案2：将m绑定为类方法 #

通过types.MethodType把自定义函数m绑定为类的方法，让mock能正确处理参数传递：

from unittest import mock
import types

class A:
    def __init__(self, x):
        self.x = x
    def f(self, y):
        return self.x * y

a = A(2)

print(a.f(3))  # output: 6

original_f = A.f
def m(self, y):
    return original_f(self, y + 1)

with mock.patch('__main__.A.f', wraps=types.MethodType(m, A)) as mk:
    print(a.f(3))  # output: 8

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