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

Python - create mock test for class method that has context manager

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

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

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

这是相同的测试,使用 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

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