In my last article, I showed how pytest’s fixture system and plug-in infrastructure can help you with writing cleaner and better tests. Fixtures allow you to create a clean event loop instance for every test case. The plug-in system allows you to write test functions that are actually asyncio coroutines. While I was working on that articel, Tin Tvrtkovic created the plug-in pytest-asyncio.
In short, it lets you do this:
importasyncioimporttimeimportpytest@pytest.mark.asynciodeftest_coro(event_loop):before=time.monotonic()awaitasyncio.sleep(0.1,loop=event_loop)after=time.monotonic()assertafter-before>=0.1Instead of this:
importasyncioimporttimedeftest_coro():loop=asyncio.new_event_loop()try:asyncio.set_event_loop(loop)before=time.monotonic()loop.run_until_complete(asyncio.sleep(0.1,loop=loop))after=time.monotonic()assertafter-before>=0.1finally:loop.close()So using pytest-asyncio clearly improves your test (and there is even more, what this plug-in does!).
While I have been working on aiomas, some additional requirements came up that were not so easily covered. What aiomas basically does is adding three layers of abstraction around the asyncio transports:
- The channel layer lets you send JSON or MsgPack encoded messages in a request-reply manner. This layer uses a custom protocol that works with different kinds of transports: TCP sockets, Unix domain sockets and custom transport called local queue.
- The RPC layer creates a remote-procedure-call system on top of the channel layer.
- The agent layer (for multi-agent systems) hides even more of the networking-related stuff and lets you basically write classes that call methods of other classes over a network connection.
Here is a simple example of how the channel layer works:
importaiomasasyncdefhandle_client(channel):"""Handle a client connection."""req=awaitchannel.recv()print(req.content)awaitreq.reply('cya')awaitchannel.close()asyncdefclient():"""Client coroutine: Send a greeting to the server and wait for a reply."""channel=awaitaiomas.channel.open_connection(('localhost',5555))rep=awaitchannel.send('ohai')print(rep)awaitchannel.close()server=aiomas.run(aiomas.channel.start_server(('localhost',5555),handle_client))aiomas.run(client())server.close()aiomas.run(server.wait_closed())Requirements for our tests
So with this in mind, I had the following requirements for my tests:
I need a clean event loop instance for every test.
This can be solved with the event_loop fixture provided by pytest-asyncio.
Every test should be run with every transport available (TCP socket, Unix domain socket, …).
This could in theory be solved with the pytest.mark.parametrize() decorator (but not in my case as we will see later).
Every test needs a client coroutine. Ideally, this would be the test itself.
pytest-asyncio’spytest.mark.asyncio decorator solves this.
Every test needs a server with a custom callback for client connections. Servers must be cleanly shut down no matter what the outcome of the test is.
It would seem that a fixture would do the job, but every server needs a test specific callback for handling client connections. This makes it a lot harder.
I don’t want any “address already in use” errors if one test fails badly.
pytest-asyncio’sunused_tcp_port fixture comes to help.
I don’t want to use loop.run_until_complete() all the time.
Again, the pytest.mark.asyncio decorator solves this.
To wrap up what remains to be solved: Every test needs at least two fixtures (one for the event loop, one for the address type), but I want to combine them as a single fixture. Creating a fixture for setting up a server would also be nice, but how can we do this?
Our first approach
The first thing we can do is to wrap the loop and the address type in a fixture. We’ll call it ctx (short for test context). With fixture parameters, it is also easy to create one fixture instance for every address type:
importtempfileimportpyimportpytestclassContext:def__init__(self,loop,addr):self.loop=loopself.addr=addr@pytest.fixture(params=['tcp','unix'])defctx(request,event_loop,unused_tcp_port,short_tmpdir):"""Generate tests with TCP sockets and Unix domain sockets."""addr_type=request.paramifaddr_type=='tcp':addr=('127.0.0.1',unused_tcp_port)elifaddr_type=='unix':addr=short_tmpdir.join('sock').strpathelse:raiseRuntimeError('Unknown addr type: %s'%addr_type)ctx=Context(event_loop,addr)returnctx@pytest.yield_fixture()defshort_tmpdir():"""Generate a short temp. dir for Unix domain sockets. The paths provided by ptest's tmpdir fixture are too long on some platforms."""withtempfile.TemporaryDirectory()astdir:yieldpy.path.local(tdir)This lets us write our tests like this:
importaiomas@pytest.mark.asyncioasyncdeftest_channel(ctx):results=[]asyncdefhandle_client(channel):req=awaitchannel.recv()results.append(req.content)awaitreq.reply('cya')awaitchannel.close()server=awaitaiomas.channel.start_server(ctx.addr,handle_client)try:channel=awaitaiomas.channel.open_connection(ctx.addr)rep=awaitchannel.send('ohai')results.append(rep)awaitchannel.close()finally:server.close()awaitserver.wait_closed()assertresults==['ohai','cya']This works already very nicely and every test using the ctx fixture is run once for every address type.
However, two problems remain:
- Our ctx fixture always requires an unused TCP port and a temporary directory – although we only need one of both in each case.
- Setting up the server (and closing it) also involves some code which will be the same for every test and should thus be moved into a fixture. However, a server fixture won’t work directly, because every server needs a test specific callback as you can see in the line where we create the server (server = await ...). But without a server fixture, we can’t have a tear-down method for it …
Let’s see how we can tackle these issues.
Approach number two
The first problem can be solved by using the getfuncargvalue() method of the request object that our fixture receives. Using this method, we can manually call a fixture function:
@pytest.fixture(params=['tcp','unix'])defctx(request,event_loop):"""Generate tests with TCP sockets and Unix domain sockets."""addr_type=request.paramifaddr_type=='tcp':port=request.getfuncargvalue('unused_tcp_port')addr=('127.0.0.1',port)elifaddr_type=='unix':tmpdir=request.getfuncargvalue('short_tmpdir')addr=tmpdir.join('sock').strpathelse:raiseRuntimeError('Unknown addr type: %s'%addr_type)ctx=Context(event_loop,addr)returnctxTo help with issue number two, we can extend our Context class that is passed into every test. We add a method Context.start_server(client_handler) that we can call from within our tests. We also add a finalize/teardown part to our ctx fixture that will close the server once we are done. And while we are at it, we’ll also create some more shortcut functions:
importasyncioimporttempfileimportpyimportpytestclassContext:def__init__(self,loop,addr):self.loop=loopself.addr=addrself.server=Noneasyncdefconnect(self,**kwargs):"""Create and return a connection to "self.addr"."""return(awaitaiomas.channel.open_connection(self.addr,loop=self.loop,**kwargs))asyncdefstart_server(self,handle_client,**kwargs):"""Start a server with the callback *handle_client* listening on"self.addr"."""self.server=awaitaiomas.channel.start_server(self.addr,handle_client,loop=self.loop,**kwargs)asyncdefstart_server_and_connect(self,handle_client,server_kwargs=None,client_kwargs=None):"""Shortcut for:: await ctx.start_server(...) channel = await ctx.connect()""""ifserver_kwargsisNone:server_kwargs={}ifclient_kwargsisNone:client_kwargs={}awaitself.start_server(handle_client,**server_kwargs)return(awaitself.connect(**client_kwargs))asyncdefclose_server(self):"""Close the server."""ifself.serverisnotNone:server,self.server=self.server,Noneserver.close()awaitserver.wait_closed()@pytest.yield_fixture(params=['tcp','unix'])defctx(request,event_loop):"""Generate tests with TCP sockets and Unix domain sockets."""addr_type=request.paramifaddr_type=='tcp':port=request.getfuncargvalue('unused_tcp_port')addr=('127.0.0.1',port)elifaddr_type=='unix':tmpdir=request.getfuncargvalue('short_tmpdir')addr=tmpdir.join('sock').strpathelse:raiseRuntimeError('Unknown addr type: %s'%addr_type)ctx=Context(event_loop,addr)yieldctx# Shutdown the server and wait for all pending tasks to finish:aiomas.run(ctx.close_server())aiomas.run(asyncio.gather(*asyncio.Task.all_tasks(event_loop),return_exceptions=True))With this extra functionality, our test case becomes a lot shorter, easier to read, and more reliable:
importaiomas@pytest.mark.asyncioasyncdeftest_channel(ctx):results=[]asyncdefhandle_client(channel):req=awaitchannel.recv()results.append(req.content)awaitreq.reply('cya')awaitchannel.close()channel=awaitctx.start_server_and_connect(handle_client)rep=awaitchannel.send('ohai')results.append(rep)awaitchannel.close()assertresults==['ohai','cya']The ctx fixture (and the associated Context class) is indeed not the shortest fixture I ever wrote, but it helped me to remove approx. 200 lines of boilerplate code from my tests (apart from making them more readable and maintainable).