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())

Answer 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()

Answer 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).

Answer 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)

Answer 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__()

Answer 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()

Answer 6

您可以使用asynctest并导入CoroutineMock或使用asynctest.mock.patch

Answer 7

Dustin's answer is probably the right one in the vast majority of cases. I had a different issue where the coroutine needed to return more than one value, eg simulating a read() operation, as briefly described in my comment .

After some more testing, the code below worked for me, by defining an iterator outside the mocking function, effectively remembering the last value returned to send the next one:

def test_some_read_operation(self):
    #...
    data = iter([b'data', b''])
    @asyncio.coroutine
    def read(*args):
        return next(data)
    mocked.read = Mock(wraps=read)
    # Here, the business class would use its .read() method which
    # would first read 4 bytes of data, and then no data
    # on its second read.

So, expanding on Dustin's answer, it would look like:

def get_mock_coro(return_values):
    values = iter(return_values)
    @asyncio.coroutine
    def mock_coro(*args, **kwargs):
        return next(values)

    return Mock(wraps=mock_coro)

The two immediate downsides I can see in this approach are:

1. It doesn't allow for raising exceptions easily (eg first return some data, then raise an error on second read operation).
2. I haven't found a way to use the standard Mock .side_effect or .return_value attributes to make it more obvious and readable.

Answer 8

Well, there are a bunch of answers here already, but I'll contribute my expanded version of e-satis's answer . This class mocks an async function and tracks call count and call args, just like the Mock class does for sync functions.

Tested on Python 3.7.0.

class AsyncMock:
    ''' A mock that acts like an async def function. '''
    def __init__(self, return_value=None, return_values=None):
        if return_values is not None:
            self._return_value = return_values
            self._index = 0
        else:
            self._return_value = return_value
            self._index = None
        self._call_count = 0
        self._call_args = None
        self._call_kwargs = None

    @property
    def call_args(self):
        return self._call_args

    @property
    def call_kwargs(self):
        return self._call_kwargs

    @property
    def called(self):
        return self._call_count > 0

    @property
    def call_count(self):
        return self._call_count

    async def __call__(self, *args, **kwargs):
        self._call_args = args
        self._call_kwargs = kwargs
        self._call_count += 1
        if self._index is not None:
            return_index = self._index
            self._index += 1
            return self._return_value[return_index]
        else:
            return self._return_value

Example usage:

async def test_async_mock():
    foo = AsyncMock(return_values=(1,2,3))
    assert await foo() == 1
    assert await foo() == 2
    assert await foo() == 3

Answer 9

You can subclass Mock to act like a coroutine function:

class CoroMock(Mock):
    async def __call__(self, *args, **kwargs):
        return super(CoroMock, self).__call__(*args, **kwargs)

    def _get_child_mock(self, **kw):
        return Mock(**kw)

You can use CoroMock pretty much like a normal mock, with the caveat that calls will not be recorded until the coroutine is executed by an event loop.

If you have a mock object and you want to make a particular method a coroutine, you can use Mock.attach_mock like this:

mock.attach_mock(CoroMock(), 'method_name')

Answer 10

New in Python 3.8 AsyncMock

Setting the spec of a Mock, MagicMock, or AsyncMock to a class with asynchronous and synchronous functions will automatically detect the synchronous functions and set them as MagicMock (if the parent mock is AsyncMock or MagicMock) or Mock (if the parent mock is Mock). All asynchronous functions will be AsyncMock .

class ExampleClass:

    def sync_foo():
        pass

    async def async_foo():
        pass


a_mock = AsyncMock(ExampleClass)
a_mock.sync_foo
>>> <MagicMock name='mock.sync_foo' id='...'>
a_mock.async_foo
>>> <AsyncMock name='mock.async_foo' id='...'>

mock = Mock(ExampleClass)
mock.sync_foo
>>> <Mock name='mock.sync_foo' id='...'>
mock.async_foo
>>> <AsyncMock name='mock.async_foo' id='...'>

Special Attributes for AsyncMock:

• assert_awaited
• assert_awaited_once
• assert_awaited_with
• assert_awaited_once_with
• assert_any_await
• assert_has_awaits
• assert_not_awaited
• reset_mock
• await_count
• await_args
• await_args_list