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

时间:2015-05-25 16:56:22

标签: python unit-testing mocking patch with-statement

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

2 个答案:

答案 0 :(得分:3)

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

答案 1 :(得分:1)

这是使用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

您可以添加其他自动生成的断言,或者在发现有用时对其进行更改以包括其他断言。