How to mock an async instance method of a patched class?

Question

(The following code can be run in Jupyter.) I have a class B, which uses class A, needs to be tested.

class A:
    async def f(self):
        pass

class B:
    async def f(self):
        a = A()
        x = await a.f()  # need to be patched/mocked

And I have the following test code. It seems it mocked the class method of A instead of the instance method.

from asyncio import Future
from unittest.mock import MagicMock, Mock, patch

async def test():
    sut = B()
    with patch('__main__.A') as a:  # it's __main__ in Jupyter
        future = Future()
        future.set_result('result')
        a.f = MagicMock(return_value=future)
        await sut.f()

await test()

However, the code got the error of:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
C:\Users\X~1\AppData\Local\Temp\1/ipykernel_36576/3227724090.py in <module>
     20         await sut.f()
     21 
---> 22 await test()

C:\Users\X~1\AppData\Local\Temp\1/ipykernel_36576/3227724090.py in test()
     18         future.set_result('result')
     19         a.f = MagicMock(return_value=future)
---> 20         await sut.f()
     21 
     22 await test()

C:\Users\X~1\AppData\Local\Temp\1/ipykernel_36576/3227724090.py in f(self)
      6     async def f(self):
      7         a = A()
----> 8         x = await a.f()  # need to be patched/mocked
      9 
     10 from asyncio import Future

TypeError: object MagicMock can't be used in 'await' expression

Flinty

Answer 1

In Python 3.8+, patching an async method gives you an AsyncMock, so providing a result is a little more straightforward.

In the docs of the patch method itself:

If new is omitted, then the target is replaced with an AsyncMock if the patched object is an async function or a MagicMock otherwise.

AsyncMock lets you supply a return value in a much more straightforward manner:

import asyncio
from unittest.mock import patch


class A:
    async def f(self):
        return "foo"


class B:
    async def f(self):
        return await A().f()


async def main():
    print(await B().f())

    with patch("__main__.A.f", return_value="bar") as p:
        print(await B().f())


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        sys.exit(1)

....prints:

$ python example.py
foo
bar

The side_effect kwarg covers most kinds of values you'd be looking to return (e.g. if you need your mock function await something).

• if side_effect is a function, the async function will return the result of that function,
• if side_effect is an exception, the async function will raise the exception,
• if side_effect is an iterable, the async function will return the next value of the iterable, however, if the sequence of result is exhausted, StopAsyncIteration is raised immediately,
• if side_effect is not defined, the async function will return the value defined by return_value, hence, by default, the async function returns a new AsyncMock object.

Babineaux

Nice solution. However, it need to import the module A (if A is in another module), which means it may run some code in module A while testing B? – Flinty

If you have code other than function and class definitions in a module, then yes, it would execute when imported by your test module. You can also use patch directly, which would not require a direct reference to the class object. I will update the example. – Babineaux

Answer 2

Need to change

a.f = MagicMock(return_value=future)

to

a().f = MagicMock(return_value=future)

Flinty

Answer 3

I was facing similar problem and I had to return value for the patched object function call. For example I had to mock OAuthClient has a async function

I have used new_callable=AsyncMock to modify the default behaviour of @patch decorator https://docs.python.org/3/library/unittest.mock.html#unittest.mock.patch:~:text=These%20will%20be%20passed%20to%20AsyncMock%20if%20the%20patched%20object%20is%20asynchronous%2C%20to%20MagicMock%20otherwise%20or%20to%20new_callable%20if%20specified.

class AuthService:
    def __init__(self) -> None:
        self.auth_client = AuthClient()
    
    async def get_token(payload):
        return await self.auth_client.get_token(payload)


class AuthServiceTest:
    @pytest.mark.usefixtures('auth_token_response')
    @patch("app.client.AuthClient.__init__", return_value=None, new_callable=AsyncMock)
    async def test_get_token(auth_client_async_mock, auth_token_response):
        auth_client_async_mock.get_token.return_value = auth_token_response
        service = AuthService(auth_client = auth_client_async_mock)
        response = service.get_token(
            payload={"username": "alex", "password": "strong_password", "grant_type": "password"}
        )
        # assert response

Rebirth