如何在测试函数中替换方法为自定义函数（不使用patch.object）

你遇到的问题很典型——MongoDB的游标对象（find()返回的）的count()是无参方法，用来返回匹配文档的总数；而Python列表的count()是用来统计某个元素出现的次数，参数必填。要解决这个问题，我们需要构建一个模拟游标对象，既保留类似列表的迭代能力，又能提供无参的count()方法返回预设数据的长度。

下面给你两种可行的实现方式：

方法一：使用unittest.mock.MagicMock快速构建 #

我们可以用MagicMock来模拟游标，给它添加自定义的count()方法，同时让它的迭代行为和预设的列表一致：

from unittest.mock import patch, MagicMock
import some_module

def test_somefunc():
    # 预设的模拟数据
    mock_items = [{"some_field": True}, {"some_field": True}]
    
    # 创建模拟游标对象
    mock_cursor = MagicMock()
    # 让游标可以迭代，返回预设的mock_items
    mock_cursor.__iter__.return_value = iter(mock_items)
    # 自定义count()方法，返回预设列表的长度
    mock_cursor.count.return_value = len(mock_items)
    # 支持batch_size链式调用，返回自身
    mock_cursor.batch_size.return_value = mock_cursor
    
    # patch目标集合的find方法，返回模拟游标
    with patch.object(some_module, 'mongo_collection') as mock_collection:
        mock_collection.find.return_value = mock_cursor
        
        # 调用要测试的函数
        some_module.somefunc()
        
        # 这里可以添加断言，比如验证count()被调用，或者counter的变化等
        mock_cursor.count.assert_called_once()

方法二：自定义模拟游标类（更直观可控） #

如果需要更清晰的结构，也可以自己定义一个类来模拟Mongo游标，完全按照需求实现方法：

from unittest.mock import patch
import some_module

class MockMongoCursor:
    def __init__(self, data):
        self.data = data
        # 支持迭代
        self._iterator = iter(data)
    
    def count(self):
        # 无参count()，返回数据长度
        return len(self.data)
    
    def batch_size(self, size):
        # 链式调用，返回自身（实际测试中可以忽略size参数）
        return self
    
    def __iter__(self):
        return self
    
    def __next__(self):
        return next(self._iterator)

def test_somefunc():
    mock_data = [{"some_field": True}, {"some_field": True}]
    
    with patch.object(some_module, 'mongo_collection') as mock_collection:
        # 让find返回我们自定义的模拟游标
        mock_collection.find.return_value = MockMongoCursor(mock_data)
        
        some_module.somefunc()
        
        # 验证逻辑，比如count是否被调用
        assert mock_collection.find.return_value.count() == len(mock_data)

关键说明： #

• 不管哪种方法，都要确保batch_size()方法返回模拟对象本身，因为原代码是链式调用find(...).batch_size(50)
• 模拟的游标需要支持迭代，这样原代码中如果有遍历items的逻辑也能正常运行
• count()方法要实现无参，返回预设数据的长度，和Mongo游标行为一致

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