Search code examples
pythonpytestpython-asynciopython-unittest

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


(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

Solution

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