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

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

单位_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()

---

下面是相同的测试，使用pytest和：

您可能会发现，我编写测试的方式比测试本身更有趣——我创建了一个用于帮助我理解语法的工具

以下是我如何系统地处理您的问题：

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

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'没有属性'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'))

现在，我将mock放在调用

Foo（）

之前，并将

generate\u断言放在您的断言之前，就像这样（不需要上一次打印，所以我删除了它）：

现在assert失败了（
AssertionError:assert=[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\u db.a.return\u value.\uu输入\uuuuuu.return\u value.b.fetch.return\u value=expected\u 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

您可以添加其他自动生成的断言，或者如果您觉得有用，可以修改它们以包含其他断言。
您到底想模拟什么？请检查您的代码好吗？Foo不是db的一部分，而是根据您编写的db_类。您似乎还试图模拟整个db模块，但在with语句中，您需要的是模拟db.a（）方法。最后，您误解了如何在模拟中指定方法。这里enter和exit是属性，但应该是方法。对于b.@Cilyan规范中的fetch也是一样-感谢您的响应-我更新了代码以正确导入Foo方法。您的响应很有意义，我猜首先我应该添加另一个补丁@patch（'db_class.db.a'）以进行模拟。我仍然不清楚如何实现fetch、enter和exit作为方法和属性。@vks我试图模拟db_类foo.method。我想检查它返回的结果是否是预期结果。@Cilyan好的，我想我已经得到了解决方案，谢谢你的见解！很好的上下文模拟方法。
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)

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