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
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
'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
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
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
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.
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
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 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
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
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
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.
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!