Mocking async call in python 3.5

Question

How do I mock async call from one native coroutine to other one using unittest.mock.patch?

I currently have quite an awkward solution:

class CoroutineMock(MagicMock):
    def __await__(self, *args, **kwargs):
        future = Future()
        future.set_result(self)
        result = yield from future
        return result

Then

class TestCoroutines(TestCase):
    @patch('some.path', new_callable=CoroutineMock)
    def test(self, mock):
        some_action()
        mock.assert_called_with(1,2,3)

This works but looks ugly. Is there more pythonic way to do this?

Answer 1

Based on @scolvin answer I created this (imo) cleaner way:

def async_return(result):
    f = asyncio.Future()
    f.set_result(result)
    return f

That's it, just use it around whatever return you want to be async, as in

mock = MagicMock(return_value=async_return("Example return"))
await mock()

Answer 2

Everyone's missing what's probably the simplest and clearest solution:

@patch('some.path')
def test(self, mock):
    f = asyncio.Future()
    f.set_result('whatever result you want')
    mock.return_value = f
    mock.assert_called_with(1, 2, 3)

remember a coroutine can be thought of as just a function which is guaranteed to return a future which can, in turn be awaited.

Answer 3

The solution was actually quite simple: I just needed to convert __call__ method of mock into coroutine:

class AsyncMock(MagicMock):
    async def __call__(self, *args, **kwargs):
        return super(AsyncMock, self).__call__(*args, **kwargs)

This works perfectly, when mock is called, code receives native coroutine

Answer 4

One more variant of "simplest" solution to mock a async object, which is just a one liner.

In source:

class Yo:
    async foo(self):
        await self.bar()
    async bar(self):
        # Some code

In test:

from asyncio import coroutine

yo = Yo()
# Here bounded method bar is mocked and will return a customised result.
yo.bar = Mock(side_effect=coroutine(lambda:'the awaitable should return this'))
event_loop.run_until_complete(yo.foo())

Answer 5

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

Answer 6

Another way of mocking coroutine is to make coroutine, that returns mock. This way you can mock coroutines that will be passed into asyncio.wait or asyncio.wait_for.

This makes more universal coroutines though makes setup of tests more cumbersome:

def make_coroutine(mock)
    async def coroutine(*args, **kwargs):
        return mock(*args, **kwargs)
    return coroutine


class Test(TestCase):
    def setUp(self):
        self.coroutine_mock = Mock()
        self.patcher = patch('some.coroutine',
                             new=make_coroutine(self.coroutine_mock))
        self.patcher.start()

    def tearDown(self):
        self.patcher.stop()