PyTest：如何正确模拟导入的ContextManager类及指定方法返回值？

问题背景 #

业务代码依赖DBInterface上下文管理器类，需要在单元测试中完全Mock该类以避免操作真实数据库，同时设置get方法的返回值来验证业务逻辑：

业务代码示例：

from path.lib import DBInterface

class MyClass:

    def __init__(self):
        self.something = "something"

    def _my_method(self, some_key, new_setup):
        with DBInterface(self.something) as ic:
            current_setup = ic.get(some_key)

        if current_setup != new_setup:
            with DBInterface(self.something) as ic:
                ic.set(new_setup)

    def public_method(self, some_key, new_setup):
        return self._my_method(some_key, new_setup)

当前测试代码尝试多种方式设置get返回值均无效：

fake_db.get.return_value = None
fake_db.__enter__().get.return_value = None
fake_db.__enter__.get = Mock(return_value=None)
mocker.patch.object(MyClass.DBInterface, "get").return_value = None

现有测试代码框架：

class TestMyClass:

    def test_extractor_register(self, mocker):
        fake_db = mocker.patch.object('my_path.my_lib.DBInterface')
        fake_db.get.return_value = None
        test_class = MyClass()

        result = test_class.public_method(Tconf.test_key, Tconf.test_key_setup)
        fake_db.assert_has_calls([call().__enter__().get(Tconf.test_key),
                                  call().__enter__().set(Tconf.test_key, Tconf.test_key_setup)])
        assert result.result_status.status_code == Tconf.status_ok.status_code
        assert result.result_data == MyMethodResult.new_reg

解决方法 #

问题核心是DBInterface作为上下文管理器，with语句的执行流程是：DBInterface(...)实例化（调用类的__call__）→ 实例调用__enter__返回上下文对象ic。需要针对这个流程正确Mock：

方法1：分步Mock上下文对象 #

def test_extractor_register(self, mocker):
    # Mock DBInterface类
    fake_db_cls = mocker.patch('my_path.my_lib.DBInterface')
    
    # 创建Mock的上下文实例ic
    fake_ic = mocker.Mock()
    # 设置get方法的返回值
    fake_ic.get.return_value = None
    
    # 让DBInterface的实例返回fake_ic作为上下文对象
    fake_db_instance = mocker.Mock()
    fake_db_instance.__enter__.return_value = fake_ic
    fake_db_cls.return_value = fake_db_instance

    test_class = MyClass()
    result = test_class.public_method(Tconf.test_key, Tconf.test_key_setup)

    # 验证调用逻辑
    fake_db_cls.assert_called_with("something")
    fake_ic.get.assert_called_with(Tconf.test_key)
    fake_ic.set.assert_called_with(Tconf.test_key_setup)

    # 断言业务结果
    assert result.result_status.status_code == Tconf.status_ok.status_code
    assert result.result_data == MyMethodResult.new_reg

方法2：链式调用简化Mock #

def test_extractor_register(self, mocker):
    fake_db_cls = mocker.patch('my_path.my_lib.DBInterface')
    # 直接通过链式调用设置返回值：类实例化 → __enter__返回对象 → get方法返回值
    fake_db_cls.return_value.__enter__.return_value.get.return_value = None

    test_class = MyClass()
    result = test_class.public_method(Tconf.test_key, Tconf.test_key_setup)

    # 验证完整调用链
    fake_db_cls.assert_has_calls([
        call("something"),
        call().__enter__().get(Tconf.test_key),
        call("something"),
        call().__enter__().set(Tconf.test_key_setup)
    ])

    # 断言业务结果
    assert result.result_status.status_code == Tconf.status_ok.status_code
    assert result.result_data == MyMethodResult.new_reg

关键要点 #

• with DBInterface(...) as ic的本质是：DBInterface()生成实例 → 实例调用__enter__返回ic对象
• 之前的错误是直接给Mock的类设置get返回值，而非给__enter__返回的上下文对象设置
• 两种方法都需要确保__enter__返回的对象拥有正确的get方法返回值

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