在单元测试中，如何在单元测试中传递模拟对象以便它可以在'with'上下文中使用？答案

【问题标题】：In unittest, how to pass mock object in unit tests so it can be used in 'with' context?在单元测试中，如何在单元测试中传递模拟对象以便它可以在'with'上下文中使用？
【发布时间】：2021-06-30 22:47:30
【问题描述】：

我正在尝试为以下功能编写单元测试：

def my_func():
    _session, _engine = get_session_and_engine()
    with _session.begin():
        # Some functionality
         pass

    return result

我的单元测试文件：

@mock.patch('core.my_func_file.get_session_and_engine')
def test_my_func(mock_get_func):
    mock_session = mock.Mock()
    mock_session.return_value = mock_session
    mock_session.__enter__ = mock.Mock(return_value=(mock_session, None))
    mock_session.__exit__ = mock.Mock(return_value=None)

    mock_get_func.return_value = (mock_session, mock.Mock())
    res = my_func()
    assert res is not None

当我在 pytest 中运行此测试时，我收到以下错误。

   with _session.begin():
       AttributeError: __enter__
         core.my_func_file.py:5024: AttributeError

这里的问题是我们正在传递一个模拟对象，该对象将在“with”上下文中使用。当时，我们得到了上述错误。

如何模拟此类对象，以便在with 上下文中使用？

版本信息：
蟒蛇 3.7
pytest 6.2.3

【问题讨论】：

docs.python.org/3/library/…

标签： python unit-testing pytest pytest-mock

【解决方案1】：

扩展@jonrsharpe 的评论

由于这里的with 上下文使用_session.begin()，我们需要传递该函数的返回值。所以添加后 mock_session.begin.return_value = mock_session ，测试运行良好。

    mock_session = mock.Mock()
    mock_session.__enter__ = mock.Mock(return_value=mock_session)
    mock_session.__exit__ = mock.Mock(return_value=None)
    mock_session.begin.return_value = mock_session

参考：https://docs.python.org/3/library/unittest.mock.html#mocking-magic-methods

【讨论】：