How to mock asyncio coroutines?

Question

The following code fails with TypeError: 'Mock' object is not iterable in ImBeingTested.i_call_other_coroutines because I've replaced ImGoingToBeMocked by a Mock object.

How can I mock coroutines?

class ImGoingToBeMocked:
    @asyncio.coroutine
    def yeah_im_not_going_to_run(self):
        yield from asyncio.sleep(1)
        return "sup"

class ImBeingTested:
    def __init__(self, hidude):
        self.hidude = hidude

    @asyncio.coroutine
    def i_call_other_coroutines(self):
        return (yield from self.hidude.yeah_im_not_going_to_run())

class TestImBeingTested(unittest.TestCase):

    def test_i_call_other_coroutines(self):
        mocked = Mock(ImGoingToBeMocked)
        ibt = ImBeingTested(mocked)

        ret = asyncio.get_event_loop().run_until_complete(ibt.i_call_other_coroutines())

Solution 1

Since mock library doesn't support coroutines I create mocked coroutines manually and assign those to mock object. A bit more verbose but it works.

Your example may look like this:

import asyncio
import unittest
from unittest.mock import Mock


class ImGoingToBeMocked:
    @asyncio.coroutine
    def yeah_im_not_going_to_run(self):
        yield from asyncio.sleep(1)
        return "sup"


class ImBeingTested:
    def __init__(self, hidude):
        self.hidude = hidude

    @asyncio.coroutine
    def i_call_other_coroutines(self):
        return (yield from self.hidude.yeah_im_not_going_to_run())


class TestImBeingTested(unittest.TestCase):

    def test_i_call_other_coroutines(self):
        mocked = Mock(ImGoingToBeMocked)
        ibt = ImBeingTested(mocked)

        @asyncio.coroutine
        def mock_coro():
            return "sup"
        mocked.yeah_im_not_going_to_run = mock_coro

        ret = asyncio.get_event_loop().run_until_complete(
            ibt.i_call_other_coroutines())
        self.assertEqual("sup", ret)


if __name__ == '__main__':
    unittest.main()

Solution 2

I am writting a wrapper to unittest which aims at cutting the boilerplate when writting tests for asyncio.

The code lives here: https://github.com/Martiusweb/asynctest

You can mock a coroutine with asynctest.CoroutineMock:

>>> mock = CoroutineMock(return_value='a result')
>>> asyncio.iscoroutinefunction(mock)
True
>>> asyncio.iscoroutine(mock())
True
>>> asyncio.run_until_complete(mock())
'a result'

It also works with the side_effect attribute, and an asynctest.Mock with a spec can return CoroutineMock:

>>> asyncio.iscoroutinefunction(Foo().coroutine)
True
>>> asyncio.iscoroutinefunction(Foo().function)
False
>>> asynctest.Mock(spec=Foo()).coroutine
<class 'asynctest.mock.CoroutineMock'>
>>> asynctest.Mock(spec=Foo()).function
<class 'asynctest.mock.Mock'>

All the features of unittest.Mock are expected to work correctly (patch(), etc).

Solution 3

Springing off of Andrew Svetlov's answer, I just wanted to share this helper function:

def get_mock_coro(return_value):
    @asyncio.coroutine
    def mock_coro(*args, **kwargs):
        return return_value

    return Mock(wraps=mock_coro)

This lets you use the standard assert_called_with, call_count and other methods and attributes a regular unittest.Mock gives you.

You can use this with code in the question like:

class ImGoingToBeMocked:
    @asyncio.coroutine
    def yeah_im_not_going_to_run(self):
        yield from asyncio.sleep(1)
        return "sup"

class ImBeingTested:
    def __init__(self, hidude):
        self.hidude = hidude

    @asyncio.coroutine
    def i_call_other_coroutines(self):
        return (yield from self.hidude.yeah_im_not_going_to_run())

class TestImBeingTested(unittest.TestCase):

    def test_i_call_other_coroutines(self):
        mocked = Mock(ImGoingToBeMocked)
        mocked.yeah_im_not_going_to_run = get_mock_coro()
        ibt = ImBeingTested(mocked)

        ret = asyncio.get_event_loop().run_until_complete(ibt.i_call_other_coroutines())
        self.assertEqual(mocked.yeah_im_not_going_to_run.call_count, 1)

Solution 4

You can create asynchronous mocks yourself:

import asyncio
from unittest.mock import Mock


class AsyncMock(Mock):

    def __call__(self, *args, **kwargs):
        sup = super(AsyncMock, self)
        async def coro():
            return sup.__call__(*args, **kwargs)
        return coro()

    def __await__(self):
        return self().__await__()

Solution 5

A slightly simplified example for python 3.6+ adapted from a few of the answers here:

import unittest

class MyUnittest()

  # your standard unittest function
  def test_myunittest(self):

    # define a local mock async function that does what you want, such as throw an exception. The signature should match the function you're mocking.
    async def mock_myasync_function():
      raise Exception('I am testing an exception within a coroutine here, do what you want')

    # patch the original function `myasync_function` with the one you just defined above, note the usage of `wrap`, which hasn't been used in other answers.
    with unittest.mock.patch('mymodule.MyClass.myasync_function', wraps=mock_myasync_function) as mock:
      with self.assertRaises(Exception):
        # call some complicated code that ultimately schedules your asyncio corotine mymodule.MyClass.myasync_function
        do_something_to_call_myasync_function()

View more solutions

Comments

[Dustin Wyatt]

Well this is obvious and now I feel dumb for asking the question! Thanks!

[Dustin Wyatt]

I expanded this with with a helper and stuff in this answer: stackoverflow.com/a/29905620/23972

[AlexandreH]

+1 for the wraps keyword, which did make me understand its purpose a bit more. Quick follow up question; how would you go about returning more than one value from a mocked coroutine? i.e. read() is a coroutine and you want to first return some data b'data', and then return an EOF-like condition (e.g. no data, b'' or None). AFAICS you can't use .return_value or .side_effect on the mocked coroutine (gives a bad yield error).

[Natim]

I tried that, but I got: ValueError: a coroutine was expected, got <_GatheringFuture pending> while using it in asyncio.gather(AsyncMock())

[Wowbagger and his liquid lunch]

AsyncMock acts like a coroutine function, but gather expects a coroutine. You probably need an extra set of parentheses, like asyncio.gather(AsyncMock()()).

[Natim]

What about from asynctest import CoroutineMock

[Natim]

With an extra set of parenthesis I get TypeError: 'generator' object is not callable

[Natim]

I put a reproducible code here: github.com/Martiusweb/asynctest/issues/114

[Natim]

I found out, that, for some reasons, it is crashing with asyncio.run while working with loop.run_until_complete — bugs.python.org/issue36222

[Natim]

All this is because gather returns a future not a coroutine.

[calvin]

This answer worked best for my use case, but doesn't this only work with python3.8+? AFAICT this only works because of the addition of AsyncMock in 3.8. (Note that the docs for patch state that "the target is replaced with an AsyncMock if the patched object is an async function".)