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

Question

我正在尝试为以下函数编写单元测试：

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 上下文中使用？

版本信息：
python3.7
测试 6.2.3

Answer 1

扩展@jonrsharpe 的评论

由于with上下文在这里使用_session.begin()，我们需要为该函数传递return值。所以添加后 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

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

In unittest, how to pass mock object in unit tests so it can be used in 'with' context?

python

unit-testing

pytest

pytest-mock