unit testing with asyncio

Jan 26, 2017   (11 minute read)   #python  #testing  #asyncio 

The Problem

# One method on this class is an asyncio coroutine.
class PersonClass(object):

    def __init__(self, name=None, cursor=None):
        """Constructor for PersonClass."""
        self.name = name
        self.cursor = cursor
        self.insert_sql = 'INSERT INTO person VALUES ?;'

    async def create(self):  # ....HOW TO TEST THIS METHOD??
        """Persist the person to the database."""
        return await self.cursor.execute(self.insert_sql, (self.name,))

Testing and even stubbing asyncio code can be confusing. And frustrating. I spent an embarrassing four hours yesterday trying to unit test my async/await code. After much googling (and screaming) I discovered two methods that are both neat and clean.

My Environment

Inspiration for writing this article comes from my project veggiecron-server, which is written with async/await syntax from Python 3.5 and 3.6. For testing I am using pytest as the test runner and MagicMock for stubbing. I also highly recommend the pytest plugin pytest-asyncio which reduces the event loop boilerplate code needed for each asynchronous test.

With unittest.TestClass

import asyncio

from unittest import TestCase
from unittest.mock import MagicMock

from src.person import PersonClass

class TestPersonClass(TestCase):

    def test_send_query_on_create(self):
        """Should send insert query to database on create()"""

        event_loop = asyncio.new_event_loop()
        asyncio.set_event_loop(event_loop)

        async def run_test():

            # Stub the database cursor
            database_cursor = MagicMock()

            # Stub the execute function with mocked results from the database
            execute_stub = MagicMock(return_value='future result!')
            # Wrap the stub in a coroutine (so it can be awaited)
            execute_coro = asyncio.coroutine(execute_stub)
            database_cursor.execute = execute_coro

            # Instantiate new person obj
            person = PersonClass(name='Johnny', cursor=database_cursor)

            # Call person.create() to trigger the database call
            person_create_response = await person.create()

            # Assert the response from person.create() is our predefined future
            assert person_create_response == 'future result!'
            # Assert the database cursor was called once
            execute_stub.assert_called_once_with(person.insert_sql, ('Johnny',))

        # Run the async test
        coro = asyncio.coroutine(run_test)
        event_loop.run_until_complete(coro())
        event_loop.close()

This is one of the simplest ways to test asynchronous code without external dependencies. What is going on here:

  • A new event loop should be created for every test, otherwise it could become a point of mutated state across multiple tests and cause strange issues.
  • await can only be called from a function marked with async, so a coroutine must be created inside the test function to await the function call person.create().
  • Again, I cannot await run_test() in the test function so I transform it into a coroutine object that I can send directly to the event loop.

With pytest.mark.asyncio

import asyncio
import pytest

from unittest.mock import MagicMock

from my_code import PersonClass

class TestPersonClass(object):

    @pytest.mark.asyncio
    async def test_send_query_on_create(self):
        """Should send insert query to database on create()"""

        # Stub the database cursor
        database_cursor = MagicMock()

        # Stub the execute function with mocked results from the database
        execute_stub = MagicMock(return_value='future result!')
        # Wrap the stub in a coroutine (so it can be awaited)
        execute_coro = asyncio.coroutine(execute_stub)
        database_cursor.execute = execute_coro

        # Instantiate new person obj
        person = PersonClass(name='Johnny', cursor=database_cursor)

        # Call person.create() to trigger the database call
        person_create_response = await person.create()

        # Assert the response from person.create() is our predefined future
        assert person_create_response == 'future result!'
        # Assert the database cursor was called once
        execute_stub.assert_called_once_with(person.insert_sql, ('Johnny',))

This code is much cleaner, and resembles any other unit test in the project. No setting up event loops or managing coroutines. Just one decorator and we are gold.

Important! Danger! EXTERMINATE!!

Please take note in the above code that TestPersonClass is not a child class of unittest.TestCase. If it was, the test would still succeed – but the success would be a false positive because code after the await expression would not run.

Why is this happening? The answer is complex enough that it deserves a separate post, but the tl;dr version is that on line 93 of pytest-asyncio’s source the author is expecting the event loop to be passed into the test from a pytest fixture, while unittest.TestCase methods cannot directly receive fixture function arguments. Whew.

External References

Articles

Docs