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
andmock_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)
- update the mocked request with the current request
-
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 handlerapp.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!
Top comments (1)
No need to add so much boilerplate, just use pytest and pytest-asyncio