This post provides some approaches and examples for unit testing Python asyncio-based code, building on the brief introduction in Python Asyncio Part 4 - Library Support. It describes how to write asynchronous unit tests, provides patterns for mocking various async constructs and offers some tips and tricks for testing and debugging in general.
Some familiarity with unit testing in Python is assumed, and it may be helpful to read this post in conjunction with the documentation for the unittest module. This post will also deal exclusively with the built-in unittest
module (as opposed to tools like pytest)
The unittest
module in the standard library contains an IsolatedAsyncioTestCase
, similar to the TestCase
module for synchronous code. It allows you to write coroutines for each test, which in turn can call other coroutines, and more generally use the await
keyword. The test case class will handle finding your tests, setting up an event loop and calling them, and reporting whether they succeed or fail. It is also possible to mix synchronous and asynchronous tests in the same test class.
This is a simple example of an async test:
import asyncio
import unittest
async def my_func():
await asyncio.sleep(0.1)
return True
class TestStuff(unittest.IsolatedAsyncioTestCase):
async def test_my_func(self):
r = await my_func()
self.assertTrue(r)
When run with something like python3 -m unittest -v test-example.py
, the following output is produced:
❯ python3 -m unittest -v test-example.py
test_my_func (test-example.TestStuff) ... ok
----------------------------------------------------------------------
Ran 1 test in 0.104s
OK
The async test case supports the setUp
and tearDown
method from TestCase
, along with asyncSetUp
and asyncTearDown
, which accept coroutines and are called after setUp
and before tearDown
, respectively, although in general it probably only makes sense to use one of asyncSetUp
or setUp
, and one of asyncTearDown
or tearDown
. In addition IsolatedAsyncioTestCase
supports setUpClass
and tearDownClass
.
NOTE:
IsolatedAsyncioTestCase
was added in Python 3.8. Prior to this, the asynctest module can be used and works in a similar way, although note that there is noasyncSetUp
andasyncTearDown
: insteadsetUp
andtearDown
are coroutines and you can’t supply both a synchronous and asynchronous version of each.
The unittest.mock
module contains an async version of MagicMock
, called AsyncMock
, which creates a mock that behaves like an async function. It implements similar side_effect
and return_value
behaviour as its synchronous counterpart, so awaiting a mock will return the return_value
or result of the side_effect
(see https://docs.python.org/3/library/unittest.mock.html#unittest.mock.AsyncMock for the list of those behaviours).
This simple example demonstrates an AsyncMock
object in action. Notice that the type returned from the call to my_mock()
is a coroutine, which must be awaited to produce a result, as in standard async code.
import unittest
from unittest.mock import AsyncMock
class TestMockingDemo(unittest.IsolatedAsyncioTestCase):
async def test_mocking_demo(self):
my_mock = AsyncMock()
my_mock.return_value = 3
r = my_mock()
print(type(r))
awaited_result = await r
print(awaited_result)
self.assertEqual(3, await my_mock())
❯ python3 -m unittest test-example.py
<class 'coroutine'>
3
.
----------------------------------------------------------------------
Ran 1 test in 0.006s
OK
When mocking an object (e.g. a class) that contains async methods, if the original object is provided as a spec
for the mock, then the different function types will be detected and correctly set as either a MagicMock
or an AsyncMock
, which should mean Python can do most of the mocking work for you in practice.
IMPORTANT!: The instances of the class
AsyncMock
are callable mock objects that will behave like coroutine functions, not coroutine objects. This works well for the most common sorts of async code you will likely write most of the time. If you need to mock more complex patterns involving passing around coroutine objects or futures you may have to define your own mocks starting withMagicMock
, but this is not particularly difficult to do.
More complex async constructs such as generators and context managers can also be mocked with AsyncMock, as in the examples below. AsyncMock also has some async-specific assertions, see https://docs.python.org/3/library/unittest.mock.html#unittest.mock.AsyncMock for the full list and further documentation.
import unittest
from unittest.mock import AsyncMock, Mock
class TestMockingDemo(unittest.IsolatedAsyncioTestCase):
async def test_mock_generator(self):
expected_values = ["foo", "bar", "baz"]
my_mock_generator = AsyncMock()
my_mock_generator.__aiter__.return_value = expected_values
actual_values = []
async for value in my_mock_generator:
actual_values.append(value)
self.assertListEqual(expected_values, actual_values)
async def test_mock_context_manager(self):
mock_cm = AsyncMock()
# Note that by default an AsyncMock returns more AsyncMocks - we have to replace it with a Mock if we want a
# synchronous function
mock_cm.get_state = Mock(return_value="Not entered")
# Get a context object as a result of entering the context manager. Alternatively, __aenter__ could return
# mock_cm, to emulate the behaviour of returning self upon entering the context manager
mock_ctx = AsyncMock()
mock_cm.__aenter__.return_value = mock_ctx
mock_ctx.get_state = Mock(return_value="Entered")
print(mock_cm.get_state())
self.assertEqual("Not entered", mock_cm.get_state())
async with mock_cm as entered_ctx:
print(entered_ctx.get_state())
self.assertEqual("Entered", entered_ctx.get_state())
async def test_mock_has_awaits(self):
my_mock = AsyncMock()
my_mock.assert_not_awaited()
await my_mock(27)
my_mock.assert_awaited_once_with(27)
Sometimes it’s also useful to temporarily patch a piece of code with a fairly complex mock, in a way that gets re-used
across multiple tests. One way to achieve this is to write a context manager that constructs the mock, patches the class
under test and yields the mock object. A simple example is shown below: note that the patch only applies inside the
context manager, the use of mock.create_autospec
to construct a mock object with sync and async methods that match
the original class, and that mock_class
remains in scope outside the context manager (because control blocks don’t
create a new scope in Python). Of course, this is also possible in purely synchronous code.
from contextlib import contextmanager
from typing import Generator
from unittest import mock
import unittest
class ClassUnderTest:
async def func(self):
return "original"
@contextmanager
def complex_mock(option_1, option_2=None) -> Generator[mock.MagicMock, None, None]:
# Note the mock instance has to be created before patching for autospec to work
mock_instance = mock.create_autospec(ClassUnderTest)
mock_instance.func.return_value = option_1
with mock.patch("test_example.ClassUnderTest") as class_mock:
# When constructing as `ClassUnderTest()`, return the mock instance
class_mock.return_value = mock_instance
yield mock_instance
class TestContextManagerMock(unittest.IsolatedAsyncioTestCase):
async def test_complex_mock(self):
with complex_mock("mocked") as mock_class:
class_under_test = ClassUnderTest()
return_val = await class_under_test.func()
self.assertEqual("mocked", return_val)
mock_class.func.assert_awaited_once()
class_under_test_unmocked = ClassUnderTest()
unmocked_return_val = await class_under_test_unmocked.func()
self.assertEqual("original", unmocked_return_val)
Just like in synchronous Python code, an interactive debugger can be started at any time by calling the breakpoint()
built in, which will drop into the pdb debugger with the ability to inspect variables, set breakpoints and step through code (there’s a useful guide to the commands at https://nblock.org/2011/11/15/pdb-cheatsheet/). This also works well with async code, however more care is needed, because the execution flow will move around when a Task yields and another one takes over.
Consider this simple example. Note that here the debugger is called using a different form to allow certain modules to be skipped, which will become important later. For comparison, calling breakpoint()
effectively calls import pdb; pdb.set_trace()
.
import asyncio
async def my_sleeping_coro():
for i in range(0, 4):
import pdb; pdb.Pdb(skip=["asyncio.*"]).set_trace()
print(f"Wake 1. Iteration {i}")
await asyncio.sleep(5)
async def my_other_routine():
for i in range(0, 10):
print(f"Wake 2. Iteration {i}")
await asyncio.sleep(2)
async def run_coroutines():
await asyncio.gather(my_sleeping_coro(), my_other_routine())
asyncio.run(run_coroutines())
This is what happens when the above code is run and stepped through. Notice how my_other_routine
runs several iterations without stopping, because the next
command is moving to the next statement in my_sleeping_coro
.
❯ python3 test-example.py
> /cloudfit-public-docs/asyncio/test-example.py(5)my_sleeping_coro()
-> for i in range(0, 4):
(Pdb) next
> /cloudfit-public-docs/asyncio/test-example.py(6)my_sleeping_coro()
-> print(f"Wake 1. Iteration {i}")
(Pdb) next
Wake 1. Iteration 0
> /cloudfit-public-docs/asyncio/test-example.py(7)my_sleeping_coro()
-> await asyncio.sleep(5)
(Pdb) next
Wake 2. Iteration 0
Wake 2. Iteration 1
Wake 2. Iteration 2
> /cloudfit-public-docs/asyncio/test-example.py(5)my_sleeping_coro()
-> for i in range(0, 4):
(Pdb) next
> /cloudfit-public-docs/asyncio/test-example.py(6)my_sleeping_coro()
-> print(f"Wake 1. Iteration {i}")
(Pdb) next
Wake 1. Iteration 1
> /cloudfit-public-docs/asyncio/test-example.py(7)my_sleeping_coro()
-> await asyncio.sleep(5)
(Pdb)
However, if instead we single-step through the code, after skipping all the asyncio
library code the step
moves into the other coroutine, because it is now executing while my_sleeping_coro
sleeps. However my_other_routine
only runs once before another iteration of my_sleeping_coro
happens, because pausing in the debugger caused the 5-second timeout to expire and adjusted the order of execution.
❯ python3 test-example.py
> /cloudfit-public-docs/asyncio/test-example.py(5)my_sleeping_coro()
-> for i in range(0, 4):
(Pdb) step
> /cloudfit-public-docs/asyncio/test-example.py(6)my_sleeping_coro()
-> print(f"Wake 1. Iteration {i}")
(Pdb) step
Wake 1. Iteration 0
> /cloudfit-public-docs/asyncio/test-example.py(7)my_sleeping_coro()
-> await asyncio.sleep(5)
(Pdb) step
--Return--
> /cloudfit-public-docs/asyncio/test-example.py(7)my_sleeping_coro()-><Future pending>
-> await asyncio.sleep(5)
(Pdb) step
--Call--
> /cloudfit-public-docs/asyncio/test-example.py(9)my_other_routine()
-> async def my_other_routine():
(Pdb) step
> /cloudfit-public-docs/asyncio/test-example.py(10)my_other_routine()
-> for i in range(0, 10):
(Pdb) step
> /cloudfit-public-docs/asyncio/test-example.py(11)my_other_routine()
-> print(f"Wake 2. Iteration {i}")
(Pdb) step
Wake 2. Iteration 0
> /cloudfit-public-docs/asyncio/test-example.py(12)my_other_routine()
-> await asyncio.sleep(2)
(Pdb) step
--Return--
> /cloudfit-public-docs/asyncio/test-example.py(12)my_other_routine()-><Future pending>
-> await asyncio.sleep(2)
(Pdb) step
--Call--
> /cloudfit-public-docs/asyncio/test-example.py(7)my_sleeping_coro()-><Future finished result=None>
-> await asyncio.sleep(5)
(Pdb) step
> /cloudfit-public-docs/asyncio/test-example.py(5)my_sleeping_coro()-><Future finished result=None>
-> for i in range(0, 4):
(Pdb) step
> /cloudfit-public-docs/asyncio/test-example.py(6)my_sleeping_coro()-><Future finished result=None>
-> print(f"Wake 1. Iteration {i}")
(Pdb) step
Wake 1. Iteration 1
> /cloudfit-public-docs/asyncio/test-example.py(7)my_sleeping_coro()-><Future finished result=None>
-> await asyncio.sleep(5)
(Pdb)
Python also provides a “debug mode” for asyncio, which will report of coroutines that were not awaited, and calls which took too long (and probably should have used run_in_executor
). It can be enabled in various ways as documented in https://docs.python.org/3/library/asyncio-dev.html (and is enabled by default when running unit tests), and produces output like the below, caused by debugger making that coroutine take a long time.
❯ PYTHONDEVMODE=1 python3 test-example.py
> /cloudfit-public-docs/asyncio/test-example.py(5)my_sleeping_coro()
-> for i in range(0, 4):
(Pdb) n
> /cloudfit-public-docs/asyncio/test-example.py(6)my_sleeping_coro()
-> print(f"Wake 1. Iteration {i}")
(Pdb) n
Wake 1. Iteration 0
> /cloudfit-public-docs/asyncio/test-example.py(7)my_sleeping_coro()
-> await asyncio.sleep(5)
(Pdb) n
Executing <Task pending name='Task-2' coro=<my_sleeping_coro() running at /cloudfit-public-docs/asyncio/test-example.py:7> wait_for=<Future pending cb=[Task.task_wakeup()] created at /home/samn/.pyenv/versions/3.10.0/lib/python3.10/asyncio/base_events.py:424> cb=[gather.<locals>._done_callback() at /home/samn/.pyenv/versions/3.10.0/lib/python3.10/asyncio/tasks.py:718] created at /home/samn/.pyenv/versions/3.10.0/lib/python3.10/asyncio/tasks.py:638> took 4.414 seconds
Wake 2. Iteration 0
Wake 2. Iteration 1
Wake 2. Iteration 2
> /cloudfit-public-docs/asyncio/test-example.py(5)my_sleeping_coro()
-> for i in range(0, 4):
(Pdb)
The previous section noted that sometimes the effect of debugging will change the order in which code runs, and this is also true of writing unit tests: the act of testing may change the order and timing of various tasks. However this is also the case if the timing of IO-bound calls changes (e.g. due to network congestion), or the underlying OS is busy and your process is paused. In any case, code which relies asyncio.sleep()
to ensure tasks run at specific times or in a specific order is likely to be very brittle and prone to race conditions.
For the case where tasks need to run in a specific order, asyncio provides synchronisation primitives, similar to those used when writing multi-threaded code. See https://docs.python.org/3/library/asyncio-sync.html#asyncio.Event for documentation and examples.
For controlling time, a better option is to write your code to monitor a system clock using the time
module, and perform actions once a specific time has passed, for example by waking up periodically (with asyncio.sleep(0.1)
for example) and checking the current time. This has the advantage that calls to the time
module can be patched in tests, to directly control how much time your code thinks has passed. This patching approach is equally valid in synchronous and asynchronous, as in the example, which uses synchronous code for simplicity.
import time
import unittest
from unittest import mock
def my_timed_function():
start_time = time.monotonic()
# Do some work
end_time = time.monotonic()
elapsed_time = end_time - start_time
return elapsed_time
class TestTimeDemo(unittest.IsolatedAsyncioTestCase):
def test_patch_time(self):
actual_start_time = time.monotonic()
with mock.patch('time.monotonic') as mock_monotonic:
mock_monotonic.side_effect = [actual_start_time, actual_start_time + 3600]
elapsed_time = my_timed_function()
self.assertEqual(3600, elapsed_time)
actual_elapsed_time = time.monotonic() - actual_start_time
print(f"Function elapsed time: {elapsed_time}. Actual elapsed time {actual_elapsed_time}")
Running these tests produces:
❯ python3 -m unittest test-example.py
Function elapsed time: 3600.0. Actual elapsed time 0.0002748540009633871
.
----------------------------------------------------------------------
Ran 1 test in 0.003s
OK
In newer versions of Python the unittest
module provides fairly comprehensive support for writing tests of asyncio-based code, and should make writing those tests very similar to writing them for synchronous code. This post has also provided some templates and examples for how to mock out more complex asyncio constructs.