Search code examples
pythonmockingcontextmanager

Writing a contextmanager that patches an object

In my Python 3 test code, I have a lot of these statements:

from unittest.mock import patch

user = User(...)
with patch.object(MyAuthenticationClass, 'authenticate', return_value=(user, 'token'):
    # do something

Now I want to write this as:

with request_user(user):
    # do something

How would I write a method request_user as context manager such that it patches the authentication in this way, and removes the patch after the with block?

Solution

  • You can write a simple wrapper like this:

    def request_user(user):
    return patch.object(MyAuthenticationClass, 'authenticate', return_value=(user, 'token')

    And use it:

    with request_user(user):
    # ...