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

Question

我正在尝试为具有上下文管理器和许多调用的类函数的方法编写单元测试。我很难理解如何正确模拟该函数，以便我可以测试返回值。我要模拟的类是 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

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()

Answer 2

这是相同的测试，使用 pytest 和 mocker fixture:

def test(mocker):
    mock_db = mocker.MagicMock(name='db')
    mocker.patch('db_class.db', new=mock_db)
    expected_result = [5, 10]
    mock_db.a.return_value.__enter__.return_value.b.fetch.return_value = expected_result

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

您可能会发现我编写测试的方式比测试本身更有趣 - 我创建了一个 python library 来帮助我了解语法。

以下是我系统地解决您的问题的方法：

我们从你想要的测试和我的助手库开始：

import db_class

from mock_autogen.pytest_mocker import PytestMocker

def test(mocker):
    # this would output the mocks we need
    print(PytestMocker(db_class).mock_modules().prepare_asserts_calls().generate())

    # your original test, without the mocks
    expected_result = [5,10]
    foo = db_class.Foo()
    result = foo.method()
    assert result == expected_result

现在测试显然失败了（AttributeError: module 'db' has no attribute 'a'），但打印输出很有用：

# mocked modules
mock_db = mocker.MagicMock(name='db')
mocker.patch('db_class.db', new=mock_db)
# calls to generate_asserts, put this after the 'act'
import mock_autogen
print(mock_autogen.generator.generate_asserts(mock_db, name='mock_db'))

现在，我将模拟放在调用Foo() 之前和generate_asserts 之后，就在你的断言之前，就像这样（不需要之前的打印，所以我删除了它）：

def test(mocker):
    # mocked modules
    mock_db = mocker.MagicMock(name='db')
    mocker.patch('db_class.db', new=mock_db)

    # your original test, without the mocks
    expected_result = [5,10]
    foo = db_class.Foo()
    result = foo.method()

    # calls to generate_asserts, put this after the 'act'
    import mock_autogen
    print(mock_autogen.generator.generate_asserts(mock_db, name='mock_db'))

    assert result == expected_result

现在断言失败 (AssertionError: assert <MagicMock name='db.a().__enter__().b.fetch()' id='139996983259768'> == [5, 10])，但我们再次获得了一些有价值的输入：

mock_db.a.return_value.__enter__.assert_called_once_with()
mock_db.a.return_value.__enter__.return_value.b.fetch.assert_called_once_with()
mock_db.a.return_value.__exit__.assert_called_once_with(None, None, None)

请注意第二行，这几乎就是您需要模拟的内容。稍作改动，它看起来就像mock_db.a.return_value.__enter__.return_value.b.fetch.return_value = expected_result，这样，我们就可以得到最终版本的测试：

def test(mocker):
    mock_db = mocker.MagicMock(name='db')
    mocker.patch('db_class.db', new=mock_db)
    expected_result = [5, 10]
    mock_db.a.return_value.__enter__.return_value.b.fetch.return_value = expected_result

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

您可以添加其他自动生成的断言，或者如果您觉得有用，可以更改它们以包含其他断言。