Python: Mocking a Context Manager

Python: Mocking a context manager

You are setting the wrong mock: mock_tmp is not the context manager, but instead returns a context manager. Replace your setup line with:

mock_tmp.return_value.__enter__.return_value.name = mytmpname

and your test will work.

Python patch context manager to return object

The following should work:

_mocked_ctx_manager.return_value.__enter__.return_value = client

Your _mocked_ctx_manager returns a context manager. Therefore you need to set the __enter__.return_value of _mocked_ctx_manager.return_value.

I found the following article to be helpful: Surrender Python Mocking! I Have You Now.

How do you use patch() as a context manager?

I have figured out the issue. When patch is given a class, it will return a class, not an object of that class. So mock in my example is not a DataBaseMock object, but a reference to the class. This is why class level variables are visible, but not object fields. In order get my desired functionality, I did this:

class TestFunctions():
    def test_function(self):
        with patch("path.to.database.call") as mock:
            mock.return_value = DataBaseMock()
            result = function_i_am_testing()
            assert mock.return_value.x == result    

Now, mock is a MagicMock object, whose return value is the object I need.

Mock requests.Session.get in a nested context manager

You got it almost right, the correct way to mock is:

with patch("main.Session") as mocked_session:
    mocked_session.return_value.__enter__.return_value.get.return_value.__enter__.return_value = "MOCKED"

I will break this down for clarification.

Your session_mock refers to the Session class. To get the Session instance, you use return_value (as you do for every call).

The with statement internally calls __enter__() on that session to get the session instance, so you have to add __enter__ for the method and return_value for the call.

Now you call get() on that object, which by the same logic needs an added get.return_value.

And finally, this is used in another with statement to create response, which again needs an added __enter__.return_value.

The result gets you the response object for your mocked Session, and you can set another value.

If breaking that down in code, it would look something like:

with patch("main.Session") as mocked_session:
    mocked_session_instance = mocked_session.return_value
    mocked_session_object = mocked_session_instance.__enter__.return_value
    mocked_get = mocked_session_object.get.return_value
    mocked_response = mocked_get.__enter__
    mocked_response.return_value = "MOCKED"

Python - create mock test for class method that has context manager

Thanks to the commenters I have found a solution that works for me. The trick was to patch the correct class, in this case I wanted to patch db_class.db.a instead of db_class.db. After that, it is important to make sure that the fetch() call is a method (I think I'm getting that correct). The tricky part about this problem for me was patching the correct thing as well as dealing with the context manager which requires a bit of extra tinkering.

@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()

---

Related Topics