Python单元测试中如何Mock函数内的实例方法调用？

我明白你遇到的问题了——你现在是把mock的对象搞错了！你mock的是AccountDBConnection这个类，但get_query_results是调用在这个类创建的实例上的，直接断言类的方法调用，自然会提示从未被调用。

问题拆解 #

当你用@patch('package.module_that_defines_data_function.AccountDBConnection')装饰测试方法时，mock_AccountDBConnection代表的是被mock后的类。而你的业务代码里，with AccountDBConnection(account_id) as adb做了两件关键的事：

1. 调用AccountDBConnection(account_id)创建一个实例（对应mock类的return_value属性）
2. 调用这个实例的__enter__方法，把返回值赋值给adb（对应实例的__enter__.return_value属性）

你要断言的是adb.get_query_results，也就是这个实例的方法，而不是类本身的方法。

解决方案 #

调整你的测试代码，先获取到with语句中adb对应的mock实例，再对它的方法做断言：

from unittest.mock import patch, ANY
import unittest
from package.module_that_defines_data_function import data_function

class TestDataFunction(unittest.TestCase):
    @patch('package.module_that_defines_data_function.AccountDBConnection')
    def test_data_function(self, mock_AccountDBConnection):
        # 获取with语句里adb对应的mock实例
        mock_adb_instance = mock_AccountDBConnection.return_value.__enter__.return_value
        
        # 执行要测试的函数
        data_function()
        
        # 现在断言实例的get_query_results方法调用
        mock_adb_instance.get_query_results.assert_called_once_with(ANY)
        
        # 其他你需要的断言...

补充说明 #

如果你的AccountDBConnection上下文管理器的__enter__方法返回的不是自身（这种情况比较少见），你可以手动设置__enter__的返回值，比如：

from unittest.mock import MagicMock

mock_adb = MagicMock()
mock_AccountDBConnection.return_value.__enter__.return_value = mock_adb

不过默认的MagicMock已经会自动处理这种链式调用，所以大多数情况下直接用上面的标准写法就足够了。

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