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

Question

I am trying to write unit test for the following function:

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

    return result

My unit test file:

@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

When I run this test in pytest, I am getting the following error.

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

The issue here is we are passing a mock object which is to be used in 'with' context. At that time, we are getting the above mentioned error.

How to mock such objects so that it can be used in with context?

Version information:
python 3.7
pytest 6.2.3

Answer 1

Extending the comment of @jonrsharpe

Since with context here uses _session.begin() , we need to pass the return value for that function. So after adding mock_session.begin.return_value = mock_session , the test runs fine.

    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

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