我正在尝试为具有上下文管理器和许多调用的类函数的方法编写单元测试.我很难理解如何正确模拟该函数,以便可以测试返回值.我要模拟的类是db.正如您在下面看到的那样,我正在使用补丁程序,但是我不知道如何获取它以返回正确的方法调用.我正在得到一个通用的模拟函数,而不是我期望的返回值.

db_class.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

解决方法:

感谢评论者,我找到了适合我的解决方案.技巧是修补正确的类,在这种情况下,我想修补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()