Unit Testing AsyncIO Code

Posted by
on under

I'm currently in the process of adding asyncio support to my Socket.IO server. Being experienced in the eventlet and gevent way of doing async, this has been a very interesting project, and a great learning experience. At some point I reached a head scratching moment, when I tried to write some unit tests to exercise the new code I was writing, but found that the Python unittest and mock libraries do not offer any facilities specifically tailored to testing asyncio.

One of the aspects I'm most proud of regarding my Socket.IO server is how complete the unit test suite is, in spite of being a highly networked project that runs under multiple asynchronous and networking frameworks. Given the high complexity of this project, I considered it a requirement to properly test all this new asyncio code, so I spent some time thinking about ways to implement asyncio testing. In this article I want to share the solutions I came up with, which helped me reach 100% coverage of my asyncio code.

Testing Synchronous Code

Let me first show you what is my approach to unit testing in this project using a synchronous example. Let's assume we have to unit test the receive() function shown below, stored in a receive.py module:

def receive(packet_type, packet_data):
    """Receive packet from the client."""
    if packet_type == 'PING':
        send_to_client("PONG", packet_data)
    elif packet_type == 'MESSAGE':
        response = trigger_event('message', packet_data)
        send_to_client('MESSAGE', response)
    else:
        raise ValueError('Invalid packet type')

def send_to_client(packet_type, packet_data):
    """Implementation of this function not shown."""
    pass

def trigger_event(event_name, event_data):
    """Implementation of this function not shown."""
    pass

Spend a couple minutes reviewing the receive() function. One of the problems for testing this function is that the send_to_client() function writes some data to a connected party over the network, but during unit tests there is no other party, so calling this function will cause an error. Likewise, the trigger_event() function sends a notification that may be picked up by other subsystems, and that is also bound to fail because any notifications emitted by the test are not real. To make matters worse, the return value from trigger_event() is then sent back to the client.

Luckily, we have one easy case to test, so let's begin with that. If the function receives a packet type that is not 'PING' or 'MESSAGE', you can see that a ValueError exception is raised. Testing for this condition is fairly easy:

import unittest
from receive import receive

class TestReceive(unittest.TestCase):
    def test_invalid_packet(self):
        self.assertRaises(ValueError, receive, 'FOO', 'data')

The other two cases require a bit more care, because we do not want any actions triggered by the test to result in real network access, since that is bound to fail in the test environment. What we have to do is mock the functions that cannot properly run in a test environment.

Below you can see a unit test for the 'PING' case:

import unittest
from unittest import mock
from receive import receive

class TestReceive(unittest.TestCase):
    # ...
    @mock.patch('receive.send_to_client')
    def test_ping(self, send_to_client):
        receive('PING', 'data')
        send_to_client.assert_called_once_with('PONG', 'data')

The idea of mocking is to "hijack" a function that we do not want to be called, and replace it with an alternative version that just happily returns, without doing anything. In the above examples, I used the mock.patch decorator to mock send_to_client(). The patch decorator is convenient because the real function will automatically be restored when the decorated test function ends.

An argument that I frequently hear against mocking is that it makes tests "test less", since functions that are supposed to be called when the application runs for real are suppressed. I disagree with that view. The purpose of unit tests is to do focused tests, so in my opinion, if we want to ensure that the receive() function works correctly, we do not need functions send_to_client() and trigger_event() to run, and can be mocked. That doesn't mean that these two functions should not be tested, though. You should have more unit tests that focus on them. There is certainly a place for tests that run the application for real, but those are integration tests, and they are not the subject of this article.

What happens to a function when it is mocked? The mock library replaces it with a MagicMock object, which is a super flexible object that can be used as sort of a wildcard, in the sense that you can call it as a function, with any kind of arguments, and in all cases it will just return without causing any errors. Below you can see a MagicMock object in action:

>>> from unittest import mock
>>> f = mock.MagicMock()
>>> f()
<MagicMock name='mock()' id='4373365928'>
>>> f('hello', 'world')
<MagicMock name='mock()' id='4373365928'>
>>> f.some_method('foo')
<MagicMock name='mock.some_method()' id='4373489200'>

As you can see above, you can call the mock as a function, with random arguments or no arguments at all, and the function happily returns. You can also reference any attributes or methods on the mock object, and they will be created on the fly, as more mock objects.

The most interesting aspect of mock objects, is that they keep track of how they were called, and you can check the call history and ensure it is what is expected. Consider the following example:

>>> f('hello', 'world')
<MagicMock name='mock()' id='4373526568'>
>>> f.assert_called_once_with('hello', 'world')
>>> f.assert_called_once_with('bye', 'world')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File ".../unittest/mock.py", line 825, in assert_called_once_with
    return self.assert_called_with(*args, **kwargs)
  File ".../unittest/mock.py", line 814, in assert_called_with
    raise AssertionError(_error_message()) from cause
AssertionError: Expected call: mock('bye', 'world')
Actual call: mock('hello', 'world')

By invoking the assert_called_once_with() method on the mock, we can verify that the function was called, and with exactly what arguments. Note how when I use the correct arguments, the call returns silently, but when I use wrong arguments I get an assertion. Hopefully now you can understand the second unit test I presented above.

We finally have the most complex case, which is the one where a 'MESSAGE' package is received. In this case, an event is triggered, and the event recipient (of which we know nothing) will send a response, which is then forwarded back to the client. Below you can see how to fully unit test this case:

class TestReceive(unittest.TestCase):
    # ...
    @mock.patch('receive.trigger_event', return_value='my response')
    @mock.patch('receive.send_to_client')
    def test_message(self, send_to_client, trigger_event):
        receive('MESSAGE', 'data')
        trigger_event.assert_called_once_with('message', 'data')
        send_to_client.assert_called_once_with('MESSAGE', 'my response')

For this code path I had to mock both the trigger_event() and send_to_client() functions, so there are two decorators. To simulate an event recipient returning a response, I used the return_value argument of the mock.patch() decorator, which sets what value the mock returns when it is invoked (you saw in the console session above that by default, the mock returns itself when it is called). Using return_value, I can feed a value of my choice into the receive() function as a result of calling trigger_event(), and then verify that this special value was used correctly when invoking send_to_client() in the next line. Pretty cool, right?

Going Green Is Easy

One of the reasons I like greenlet based frameworks such as Eventlet and Gevent so much, is that the asynchronous operations are mostly hidden under the covers. You can write your code in the way you are used to when writing synchronous code, and as long as you make sure that all the I/O functions that you call are compatible with the framework, you will be okay.

For that reason, testing code that is written for greenlet frameworks is really not very different from testing regular code. The receive() function would not need to be changed at all to work under Eventlet or Gevent, and as such, the unit tests presented in the previous section are appropriate as they are.

AsyncIO Testing

And now we get to the main topic of this article, which is to repeat the same exercise we've done above for the receive() function, but do it on a fully asynchronous version of it.

Let's assume the worst possible case, in which receive(), send_to_client() and trigger_event() are all asynchronous functions. Below you can see the asyncio version of our functions, in a module we are going to call async_receive.py:

async def receive(packet_type, packet_data):
    """Receive packet from the client."""
    if packet_type == 'PING':
        await send_to_client("PONG", packet_data)
    elif packet_type == 'MESSAGE':
        response = await trigger_event('message', packet_data)
        await send_to_client('MESSAGE', response)
    else:
        raise ValueError('Invalid packet type')

async def send_to_client(packet_type, packet_data):
    """Implementation of this function not shown."""
    pass

async def trigger_event(event_name, event_data):
    """Implementation of this function not shown."""
    pass

The logic in this code is actually identical to that of the synchronous versions, but I sprinkled the async and await keywords to make everything work asynchronously.

The first problem is how to run the receive() function inside a unit test, because when you invoke an asynchronous function, the function does not run, and instead a coroutine object that represents the asynchronous task in a suspended state is returned:

>>> from async_receive import receive
>>> receive('FOO', 'data')
<coroutine object receive at 0x10e5b32b0>

This may seem weird if you are not familiar with asyncio. Coroutines are scheduled in and out of the CPU by the event loop, so the only way to get the actual function to run is to instantiate a loop and run it on it:

>>> import asyncio
>>> asyncio.get_event_loop().run_until_complete(receive('FOO', 'data'))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File ".../asyncio/base_events.py", line 466, in run_until_complete
    return future.result()
  File ".../async_receive.py", line 19, in receive
    raise ValueError('Invalid packet type')
ValueError: Invalid packet type

Here I'm getting a default event loop with asyncio.get_event_loop(), and then using its run_until_complete() method to run the receive() coroutine and wait for it to end. You can see that when I passed an invalid packet type the ValueError exception was raised, just like in the synchronous case.

To be able to run asynchronous functions in a unit test, we have to run them just like the above example. To make this a bit simpler, I created a helper function that runs any coroutine in an event loop and passes its return value back to the caller:

import asyncio

def _run(coro):
    return asyncio.get_event_loop().run_until_complete(coro)

Using the helper function, we can write our first asynchronous test, for the invalid packet type case:

import unittest
from unittest import mock
from async_receive import receive

class TestReceive(unittest.TestCase):
    def test_invalid_packet(self):
        self.assertRaises(ValueError, _run, receive('FOO', 'data'))

Note how the _run() helper function receives as a single argument a coroutine object, which is returned by receive('FOO', 'data').

The remaining two cases get much more complicated, unfortunately. For the 'PING' case, we need to mock the send_to_client() call, like we did in the synchronous case, but the problem is that this is an asynchronous function and as such it is awaited by the receive() function. Sadly, a mock object cannot be used with the await keyword, only awaitable things such as coroutines can.

So how can we mock this function if we can't use a mock object? It took me a while to figure out a solution to this problem. The send_to_client() function returns a coroutine when it is invoked, so our mocked function needs to behave in the same way. But we don't want the coroutine that is returned during unit tests to represent the real function, since we can't allow that function to run. What we want, is for the mocked coroutine to invoke a MagicMock object, so that we can then ensure that the function was called as expected.

Did you get a headache trying to follow me in this reasoning? Well, I got one too when I was trying to figure this thing out. What I ended up doing is concocting another helper function, which returns a mock coroutine:

def AsyncMock(*args, **kwargs):
    m = mock.MagicMock(*args, **kwargs)

    async def mock_coro(*args, **kwargs):
        return m(*args, **kwargs)

    mock_coro.mock = m
    return mock_coro

Let's look at this function one part at a time. In the middle of the function body, there is an inner asynchronous function, called mock_coro(), that accepts any arguments it's given. Function AsyncMock() returns this inner async function, you can see that in the last line. I said above that we needed a mock asynchronous function that behaves like a real one, and this mock_coro() function does, simply because it is a real asynchronous function.

When this coroutine gets to run, we want to invoke a MagicMock() object. The m object initialized in the first line of AsyncMock() creates this mock object. Since a mock constructor accepts a variety of arguments (such as the return_value I used for one of the sync tests), I pass any arguments given to AsyncMock() directly to the mock object so that it can be configured. The body of the mock_coro inner async function invokes this mock, with all the arguments that were passed by the caller.

The last thing we need, is a way for the test code to get at this m object. To make it accessible from the outside, I added a mock attribute to the mock_coro function. In case you find this strange, in Python, functions are objects, so you can add custom attributes to them.

To help you understand how this async mock function works, let's look at an example:

>>> import asyncio
>>> from test_async_receive import AsyncMock
>>> f = AsyncMock(return_value='hello!')
>>> f('foo', 'bar')
<coroutine object AsyncMock.<locals>.mock_coro at 0x10ef84ca8>
>>> asyncio.get_event_loop().run_until_complete(f('foo', 'bar'))
'hello!'
>>> f.mock.assert_called_once_with('foo', 'bar')
>>> f.mock.assert_called_once_with('foo')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File ".../unittest/mock.py", line 825, in assert_called_once_with
    return self.assert_called_with(*args, **kwargs)
  File ".../unittest/mock.py", line 814, in assert_called_with
    raise AssertionError(_error_message()) from cause
AssertionError: Expected call: mock('foo')
Actual call: mock('foo', 'bar')

As you can see in this example, I assigned one of these mock async functions to f. If I invoke f() as a function, I get back a coroutine, exactly what we need to mimic real async functions. When run a loop on f('foo', 'bar') I get back the return value that I configured on the mock. Finally, to verify how the function was called, I have to use f.mock instead of f, to access the MagicMock instance that was called in the body of the coroutine. It's pretty complicated for such a short piece of code, but it appears to work great!

Here is the 'PING' unit test for the asynchronous version of receive():

class TestReceive(unittest.TestCase):
    # ...
    @mock.patch('async_receive.send_to_client', new=AsyncMock())
    def test_ping(self):
        _run(receive('PING', 'data'))
        from async_receive import send_to_client
        send_to_client.mock.assert_called_once_with('PONG', 'data')

Note how in the mock.patch decorator I'm now using the new argument, which replaces the target function with an object of my choice (the mock coroutine, in this case), instead of assigning a MagicMock object like before. It is unfortunate that the patch decorator does not pass the mocked object as an argument to the test when you use the new argument, so I had to import it from its module to then be able to access the inner mock and check that the call was made.

At this point the worse is over. The 'MESSAGE' case is a bit more complicated than the previous one, but really not that much:

class TestReceive(unittest.TestCase):
    # ...
    @mock.patch('async_receive.send_to_client', new=AsyncMock())
    @mock.patch('async_receive.trigger_event', new=AsyncMock(return_value='my response'))
    def test_message(self):
        _run(receive('MESSAGE', 'data'))
        from async_receive import send_to_client, trigger_event
        trigger_event.mock.assert_called_once_with('message', 'data')
        send_to_client.mock.assert_called_once_with('MESSAGE', 'my response')

In this one I'm passing return_value to the trigger_event() mock, and that in turn will send the argument to the inner mock, but other than that it works as the previous test.

Conclusion

I hope this gives you some ideas on how to tackle unit testing and mocking of asyncio code for your own projects. I have put all the code featured in this article in a GitHub repository, in case you want to play with it: https://github.com/miguelgrinberg/asyncio-testing.

Do you have other tricks you use when you write asyncio unit tests? Let me know in the comments below!

Become a Patron!

Hello, and thank you for visiting my blog! If you enjoyed this article, please consider supporting my work on this blog on Patreon!

17 comments
  • #1 Timo said

    Great article! Did you already check out asynctest? https://github.com/Martiusweb/asynctest

  • #2 Miguel Grinberg said

    @Timo: I did, but it seemed extremely complex and overkill, in my opinion. I did not like that it reaches inside the unittest library and uses private stuff from it, as that could break as new releases of unittest come out. The CoroutineMock class they it offers could have been useful, but once again, very difficult to read implementation, doesn't give me confidence. My implementation, which is just a few lines of code, uses only public asyncio and unittest functions, so my impression is that it is more robust.

  • #3 StefaanG said

    Thanks for the great blog, Miguel. I'm using the mock library from time to time but it's always nice to see some advanced examples.

  • #4 olujedai said

    Hello Manuel Grinberg, great tutorial as always. Quick question though does this implementation make flask-socketio compatible with asyncio?

  • #5 Miguel Grinberg said

    @olujedai: It's Miguel :)

    Flask is incompatible with Asyncio, so no, you cannot use Flask-SocketIO with Asyncio. But Flask-SocketIO is a wrapper around python-socketio, which is a generic Socket.IO server that is not specific to any framework. The python-socketio package works for any WSGI framework, and also for aiohttp and sanic under Asyncio. You can find python-socketio at https://github.com/miguelgrinberg/python-socketio. There are several examples in that repository.

  • #6 Chris J. said

    Thanks for your post. To avoid having to set mock_coro.mock, which doesn't seem necessary, you could rename the AsyncMock function to "make_async_mock" and have it return the pair (mock_coro, m).

  • #7 Miguel Grinberg said

    @Chris: this is a matter of preference. To me it makes more sense to have the coroutine and the mock in a single object, but returning the two as separate object is an acceptable solution as well.

  • #8 Dan V said

    Wouldn't it possible to simply inject the functions/methods? Then you wouldn't have to monkey patch them.

  • #9 Miguel Grinberg said

    @Dan: what do you mean exactly by "injecting"? Patching functions with the mock library is a pretty standard way to unit test.

  • #10 Glin said

    Nice article, thanks!

    BTW, Maybe it would be nice to make your own decorator async_patch which would replace the target with AsyncMock and pass it as an argument to inner function, then the async test would look pretty much the same as your synchronous test.

  • #11 Miguel Grinberg said

    @Glin: yeah, that could totally work and would simplify the test. Great idea!

  • #12 me said

    Very useful if you know what you are doing! Thanks.

  • #13 Christoforus Surjoputro said

    Great solution Miguel. I've one question about this approach. How can I mock with statement with this? For example I want to mock:
    async with aiohttp.ClientSession(headers=headers) as session:
    async with session.post(HELIO_LOGIN_URL, json=data) as resp:
    response = await resp.json()

    How can I mock the response variable? Thank you.

  • #14 Miguel Grinberg said

    @Christoforus: You would need to create a fake context manager that returns mocks for the aenter and aexit methods.

  • #15 Jude Pereira said

    This is awesome! It saved me quite some time! Thank you!

  • #16 Adam P said

    AsyncMock is super helpful. Thanks for sharing it!

    Still a little annoying to have to use a hand-rolled wrapper instead of Mock directly, perhaps someday they'll include an AsyncMock class in the unittest.mock package.

  • #17 Dan Vasquez said

    I was able to use your mock with the context manager to avoid having import statement.

    with patch.object(MyClass, 'my_method',  AsyncMock()) as my_mock:
        ...
        self.asertTrue(my_mock.mock.called)
    

Leave a Comment