Python Mock在内部方法调用场景下无法断言调用的问题

我在测试存在内部方法调用的类时，无法完成assert_called()断言操作。示例代码如下：

from unittest.mock import Mock

class Foo:
    def print1(self) -> None:
        print("print1 called")  # 能看到这个打印信息
        self.print2()
        
    def print2(self) -> None:
        print("print2 called")  # 能看到这个打印信息

foo = Foo()
foo_mock = Mock(wraps=foo)
foo_mock.print1()
foo_mock.print1.assert_called()  # 断言通过
foo_mock.print2.assert_called()  # 断言失败

运行后，foo_mock.print1.assert_called()可正常通过，但foo_mock.print2.assert_called()无法通过。原因在于print1方法中的self是原始Foo类型实例而非Mock实例，导致Mock无法记录print2的调用。

这是预期行为吗？ #

是的，这是预期行为。当使用Mock(wraps=foo)创建包装Mock时，调用foo_mock.print1()会将请求转发给原始foo实例的print1方法，而原始方法内部的self绑定的是原始Foo对象，不是Mock对象。因此内部调用self.print2()实际是调用原始实例的方法，不会被Mock捕获记录。

解决方案 #

方法一：手动绑定方法到Mock实例 #

你自己发现的修复方式有效，通过types.MethodType将原始类的方法绑定到Mock实例上，这样方法内部的self就会指向Mock对象，调用行为就能被正确记录：

import types
from unittest.mock import Mock

class Foo:
    def print1(self) -> None:
        print("print1 called", type(self))
        self.print2()
        
    def print2(self) -> None:
        print("print2 called")

foo = Foo()
foo_mock = Mock(wraps=foo)

setattr(foo_mock, "print1", types.MethodType(Mock(wraps=Foo.print1), foo_mock))  # 修复
setattr(foo_mock, "print2", types.MethodType(Mock(wraps=Foo.print2), foo_mock))  # 修复

foo_mock.print1()
foo_mock.print1.assert_called()
foo_mock.print2.assert_called()

方法二：使用patch.object替换实例方法 #

更简便的方式是用unittest.mock.patch.object直接替换实例的方法为Mock，自动处理self的绑定问题：

from unittest.mock import patch

class Foo:
    def print1(self) -> None:
        print("print1 called")
        self.print2()
        
    def print2(self) -> None:
        print("print2 called")

foo = Foo()
with patch.object(foo, 'print1', wraps=foo.print1) as mock_print1, \
     patch.object(foo, 'print2', wraps=foo.print2) as mock_print2:
    foo.print1()
    mock_print1.assert_called()
    mock_print2.assert_called()

方法三：批量绑定类方法到Mock实例 #

如果需要为类的所有方法创建Mock并绑定到Mock实例，可以遍历类的方法批量处理：

import types
from unittest.mock import Mock

class Foo:
    def print1(self) -> None:
        print("print1 called")
        self.print2()
        
    def print2(self) -> None:
        print("print2 called")

foo = Foo()
foo_mock = Mock(spec=Foo)

# 遍历类的所有可调用方法，绑定到Mock实例
for attr_name, attr_value in Foo.__dict__.items():
    if callable(attr_value):
        wrapped_mock = Mock(wraps=attr_value)
        setattr(foo_mock, attr_name, types.MethodType(wrapped_mock, foo_mock))

foo_mock.print1()
foo_mock.print1.assert_called()
foo_mock.print2.assert_called()

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