cloudfit-public-docs

Unit Testing Python Asyncio Code

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)

Async Test Cases

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 no asyncSetUp and asyncTearDown: instead setUp and tearDown are coroutines and you can’t supply both a synchronous and asynchronous version of each.

Basic Mocking

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 with MagicMock, 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)

Debugging your async code

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)

Controlling time and ordering

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

Summary

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.