In python unit, how to pass mock object in unit tests which can be used in ‘with’ context

  pytest, pytest-mock, python, unit-testing

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
         return result

My unit test file

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__ 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 informations
python 3.7
pytest 6.2.3

Source: Python Questions