Subclassing MagicMock
will propagate your custom class for all the mocks generated from your coroutine mock. For instance, AsyncMock().__str__
will also become an AsyncMock
which is probably not what you're looking for.
Instead, you might want to define a factory that creates a Mock
(or a MagicMock
) with custom arguments, for instance side_effect=coroutine(coro)
. Also, it might be a good idea to separate the coroutine function from the coroutine (as explained in the documentation).
Here is what I came up with:
from asyncio import coroutine
def CoroMock():
coro = Mock(name="CoroutineResult")
corofunc = Mock(name="CoroutineFunction", side_effect=coroutine(coro))
corofunc.coro = coro
return corofunc
An explanation of the different objects:
corofunc
: the coroutine function mock
corofunc.side_effect()
: the coroutine, generated for each call
corofunc.coro
: the mock used by the coroutine to get the result
corofunc.coro.return_value
: the value returned by the coroutine
corofunc.coro.side_effect
: might be used to raise an exception
Example:
async def coro(a, b):
return await sleep(1, result=a+b)
def some_action(a, b):
return get_event_loop().run_until_complete(coro(a, b))
@patch('__main__.coro', new_callable=CoroMock)
def test(corofunc):
a, b, c = 1, 2, 3
corofunc.coro.return_value = c
result = some_action(a, b)
corofunc.assert_called_with(a, b)
assert result == c