Python - 为具有上下文管理器的类方法创建模拟测试

Question

我正在尝试为具有上下文管理器和许多调用的类函数的方法编写单元测试。我很难理解如何正确模拟函数以便测试返回值。我试图模拟的类是 db.正如您在下面看到的，我正在使用一个补丁，但我无法弄清楚如何让它返回正确的方法调用。我得到的是通用模拟函数，而不是我期望的返回值。

数据库类.py

import db

class Foo():
    def __init__(self):
        pass
    def method(self):
        with db.a() as a:
            b = a.b
            return b.fetch()

unit_db.py

 from mock import Mock, patch, MagicMock
 from db_class import Foo

 @patch('db_class.db')
 def test(db_mock):
     expected_result = [5,10]
     db_mock.return_value = Mock(__enter__ = db_mock,
                                 __exit___ = Mock(),
                                 b = Mock(fetch=expected_result))

     foo = Foo()
     result = foo.method()
     assert result == expected_result

Answer 1

感谢评论者，我找到了适合我的解决方案。诀窍是修补正确的类，在这种情况下，我想修补 db_class.db.a 而不是 db_class.db。之后，重要的是要确保 fetch() 调用是一种方法(我认为我是正确的)。对我来说，这个问题的棘手部分是修补正确的东西以及处理需要一些额外修补的上下文管理器。

@patch('db_class.db.a')
def test(db_a):
    expected_result = [5,10]
    b_fetch = MagicMock()
    b_fetch.fetch.return_value = expected_result 
    db_a.return_value = Mock(b = b_fetch,
                         __enter__= db_a,
                         __exit__ =Mock())
    foo = Foo()
    result = foo.method()
    assert result == expected_result

if __name__ == "__main__":
    test()