Update Python unittest with asyncio tests for aiohttp and more

Add asynchronous IO to an older library

Recently I was tasked with adding full asyncio compatibility to an established library that was using requests. I chose to add an aiohttp option in addition to requests and unify the interfaces completely: the only difference for the new async interface would be adding a keyword to the class constructor and then using the await keyword for all IO. Fantastic! Or so I thought.

The unit testing problem

Like many libraries, this one has extensive unit tests in place using Python unittest. For any tests that did not directly use IO I was able to add a novel decorator discussed at the bottom of this post. The trouble came where the library was using mock with the @mock.patch("requests.post") decorator to stub requests. Did I need to mock aiohttp in the same way? The documentation discourages it. After experimenting with some proposed patterns and dependencies, I have settled on what I present below as being minimally invasive and avoiding too many new imports.

An aiohttp server

async_server.py

from aiohttp import web
from tests.helpers import SimpleMock, SimpleMockRequest


routes = web.RouteTableDef()

mock_resp = SimpleMock()
mock_req = SimpleMockRequest()


@routes.route("*", "/{_:.*}")
async def handler(request):
    mock_req.update(request)
    return web.Response(**mock_resp)


async def setup_web_server(app, host='localhost', port=8109):
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, host, port)
    await site.start()

app = web.Application()
app.add_routes(routes)

• Our first import is the only new import required to update our test cases to handle aiohttp: a simple server aiohttp.web
• Import SimpleMock and SimpleMockRequest, discussed below
• routes = web.RouteTableDef() gives us a route object
• mock_resp and mock_req will be used as pointers to the current instances of a mocked response or request
• The handler() function will match all methods and paths and:
  • update the mocked request with the current request mock_req.update(request)
  • return the mocked response return web.Response(**mock_resp)
• setup_web_server() will add the server to the event loop
• The simple web app is instantiated as app = web.Application() and registers the route to the handler app.add_routes(routes)

The mock response, "mock" request, async decorator

helpers.py

The mocked response was initially a plain dict, mock_resp = {} but I wanted to make the interface a bit more usable. SimpleMock is just a case insensitive dict that maps keys to attributes so that you can match the interface of mocked request.post more closely. For example, your unit test using requests might look like:

resp = mock.Mock()
resp.content = MOCK_RESPONSE

where the new interface will be:

mock_resp.text = MOCK_RESPONSE

class SimpleMock(dict):
    """Case insensitive dict to mock HTTP response."""
    def __init__(self, *args, **kwargs):
        super(SimpleMock, self).__init__(*args, **kwargs)
        for k in list(self.keys()):
            v = super(SimpleMock, self).pop(k)
            self.__setitem__(k, v)

    def __setitem__(self, key, value):
        super(SimpleMock, self).__setitem__(str(key).lower(), value)

    def __getitem__(self, key):
        if key.lower() not in self:
            return None
        return super(SimpleMock, self).__getitem__(key.lower())

    def __setattr__(self, key, value):
        self.__setitem__(key, value)

    def __getattr__(self, key):
        return self.__getitem__(key)

SimpleMockRequest on the other hand will probably need to be tinkered with for individual needs. I only needed method, host, url, headers, and body but added a few more entries in attributes.

class SimpleMockRequest(SimpleMock):
    """Case insensitive dict interface for an aiohttp Request object."""
    def update(self, request):
        self.clear()
        attributes = [
            "method",
            "host",
            "path",
            "path_qs",
            "query",
            "body"
        ]
        self.headers = SimpleMock(request.headers)
        self.query = SimpleMock(request.query)
        self.url = str(request.url)  # match requests interface
        self.url_object = request.url
        for attr in attributes:
            try:
                self[attr] = getattr(request, attr)
            except AttributeError:
                self[attr] = None

As you've gathered by now, we aren't mocking the request; we are sending a real request to a real server. The server will use the update() method on the pointer mock_req so that we have access to the real request that the server received shortly after it was sent. The server then responds with the truly mocked mock_resp.

• self.clear() is significant here - there is one mock_req for the whole test suite so we need to clear the old one!

We will also need this async decorator to run our unittest methods that are declared with async def:

from functools import wraps

def async_test(f):
    """
    Decorator to create asyncio context for asyncio methods or functions.
    """
    @wraps(f)
    def g(*args, **kwargs):
        args[0].loop.run_until_complete(f(*args, **kwargs))
    return g

args[0] will be "self" when a method is called.

Integrate with unittest

With the helpers in place, we can import them into a test file, test_api.py

test_api.py - imports

from tests.helpers import async_test
from tests.async_server import (
    mock_resp,
    mock_req,
    setup_web_server,
    app,
)

test_api.py - unit test setup and teardown

Get event loop, start async server, init anything else for tests

class TestSomeEndpointsAndParsers(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # you probably have some existing code above here
        cls.loop = asyncio.get_event_loop()
        cls.future = cls.loop.run_until_complete(
            setup_web_server(app, host=LOCALHOST, port=ASYNC_SERVER_PORT)
        )

# What you need may be different, just a quick example
# I init an "async version" of the library I'm testing
# This gives me clever trick I'll discuss at end

    def setUp(self):
        self.server = Device()

        async def run():
            self.session = aiohttp.ClientSession()
            self.async_server = Device(use_async=True, session=self.session)
            await self.async_server.async_init()
        self.loop.run_until_complete(run())

# You do need to clear the mock_resp after each test case

    def tearDown(self):
        mock_resp.clear()  

        async def run():
            await self.session.close()
        self.loop.run_until_complete(run())

Now we will use the pattern below to add async tests

In addition to using the async keyword, make sure to prefix or postfix async to the test name itself so you do not override the synchronous one!

Test that an endpoint is called and the mocked response is parsed correctly:

@async_test
async def test_callaction_param_async(self):
    """
    Call an action with parameters and get the results.
    """
    mock_resp.text = TEST_CALLACTION_PARAM
    response = await self.async_server.GetPortMapping(NewPortMappingIndex=0)
    self.assertEqual(response["NewInternalClient"], "10.0.0.1")
    self.assertEqual(response["NewExternalPort"], 51773)
    self.assertEqual(response["NewEnabled"], True)

Test that an endpoint is called with a well formed request:

@async_test
async def test_subscribe_async(self):
    """
    Should perform a well formed HTTP SUBSCRIBE request.
    """
    cb_url = "http://127.0.0.1:5005"
    try:
        await self.async_server.async_subscribe(cb_url, timeout=123)
    except UnexpectedResponse:
        pass
    self.assertEqual(mock_req.method, "SUBSCRIBE")
    self.assertEqual(mock_req.body, None)
    self.assertEqual(mock_req.headers["SCRINGLE"], "scrumple")
    self.assertEqual(mock_req.headers["CALLBACK"], "<%s>" % cb_url)
    self.assertEqual(mock_req.headers["HOST"], ASYNC_HOST)
    self.assertEqual(mock_req.headers["TIMEOUT"], "123")

Conclusion

I enjoyed using this pattern to add async tests to a large set of tests. Comment below with what you've found to be the best pattern for testing a code base that maintains support for synchronous as well as asynchronous IO.

Bonus

Many of the test cases I was working on didn't actually do IO. However, since an async instance of the main class used a different constructor and an async_init() that DID make IO, it seemed worthwhile to add tests for these test cases as well. What I came up with is a decorator that:

• runs the test case on the synchronous instance
• runs the test case on the asynchronous instance

This is not plug and play - it uses hard coded names: self.server and self.async_server are defined in the unittest setUp() method. The other downside is that if a test case fails you won't be able to tell whether it was the synchronous or asynchronous instance that had an error.

from functools import wraps

def add_async_test(f):
    """
    Test both the synchronous and async methods of the device (server).
    """
    @wraps(f)
    def g(*args, **kwargs):
        f(*args, **kwargs)  # run the original test
        async_args = [a for a in args]  # make mutable copy of args
        server = async_args[0].server  # save reference to self.server
        async_args[0].server = async_args[0].async_server  # set copy.server to async_server
        f(*async_args, **kwargs)  # run the test using the async instance
        async_args[0].server = server  # point self.server back to original
    return g

That's it! Decorate all the test cases that don't do IO with @add_async_test and save a bunch of code and time - if you're lucky!

Comments

[slidenerd]

No need to add so much boilerplate, just use pytest and pytest-asyncio