Update Python unittest with asyncio tests for aiohttp and more
Tags: python, testing, asyncio, aiohttp
Published: 2020-12-08 02:56:21

# 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.](https://docs.aiohttp.org/en/stable/testing.html#faking-request-object)  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
```python
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:
```python
resp = mock.Mock()
resp.content = MOCK_RESPONSE
```
where the new interface will be:
```python
mock_resp.text = MOCK_RESPONSE
```
```python
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`.
```python
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`:
```python
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
```python
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
```python
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:

```python
@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:

```python
@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.
```python
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!