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

user_825920 0
Viewed 856

(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
pytest python-unittest python-asyncio python
2 Answers

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.
user_9298610

Need to change

a.f = MagicMock(return_value=future)

to

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