Editor’s Note: This blog post was originally found on the Agari Email Security blog.
By Neil Chazin
Creating a future where all of our customers can trust their inbox can push Agari engineers to the limits of available technologies. In fact, handling the scaling requirements of Agari Advanced Threat Protection has led our Sensor team to test some of the most advanced features of the Python programming language. To maintain quality while using these features, our team created some of the first approaches to unit testing “asynchronous” Python programs. This work was recognized internationally at the 2019 PyCon convention and is now being shared here.
This is part two of an addendum to my recent talk on “Strategies for Testing Async Code” at PyCon US 2019. Part one of this series gave the background on why we needed this solution, as well as an introduction asyncio and basic testing of it. In this second part, we will discuss some more complex examples of testing.
While examples provided in the first part of this series work well for most basic unit testing of async code, there are some other circumstances that need advanced testing capabilities. These include:
- Testing coroutines that are run in event loops outside of the MainThread
- Testing synchronous functions that call into event loops
- Functional tests, which might including testing non-terminating coroutines
- Alternate event loops
I created a helper class that can assist in most of these circumstances, which I call LoopRunner:
The run() method of this Thread subclass, defined on line 10, essentially ensures that the loop passed in via __init__ is always running. In particular, it binds the loop to its own thread and then runs the loop “forever” via run_forever(). The primary external interface to this class is run_coroutine(), defined on line 18. The given coroutine is run in the LoopRunner’s loop in a threadsafe manner.
The primary value of this interface is that we have an event loop that is always running that we can run coroutines in, while not blocking the rest of our test. A secondary benefit is that it ensures the coroutines under test are not run in MainThread. This is especially important in cases where asyncio code is added to an existent code base. Running in the loop runner ensures that the test environment is close to the actual use environment.
A very simple example of using this helper class looks like this:
setUp() creates a LoopRunner, self.runner, with a new event loop on line 16. On line 17, the self.runner’s thread is started, which ensures that the loop is running. Inside our unit test on line 23 the run_coroutine() method is used to obtain the results of our test method, by running it inside the self.runner event loop. This usage is fairly similar to the tests we saw in the previous post, and using a LoopRunner is probably not necessary for such a simple case.
Here is a somewhat more complex example which requires running the coroutine in a separate thread. It is based loosely off of a pattern in one of our codebases. I call the new test class HerdRouter:
Setting up an object of this class requires passing in an event loop, as can be seen on line 6. In addition to the loop, the object also creates a lock on line 8, to be used in the thread the loop is running in, and the data set self.commands, on line 9. It is important to note that HerdRouter objects maintains its data set, self.commands, in a threadsafe manner. The external interface to objects of this class is the add_command() method on line 19. This thread-safe method ensures that our data is safely added to the set, by using asyncio.run_coroutine_threadsafe() to call the _add_command() coroutine, and ensuring it is run in the object’s event loop, self.loop. add_command(), defined on line 13 utilizes the lock, self._rlock to ensure thread safety. Note that in the general use case, there would likely be some more asynchronous work done inside this coroutine, likely before acquiring the lock.
This interface assumes that HerdRouters are used in a multi-threaded environment. In fact, asyncio.run_coroutine_threadsafe() must be called from a thread other than the one bound to its target event loop. Therefore to test this object we need to have an event loop in a thread isolated from where the test is being run. LoopRunner is the perfect tool for this task:
setUp() creates a new event loop on line 11, and then uses loop when creating a new HerdRouter, self.herd_router on line 12 and LoopRunner on 13. On line 14 the runner thread is started, which also stars the loop running. In our test, we can simply call self.router.add_command() on line 21 and check the results on line 22.
More Advanced Strategies
The example above is the final one given in the PyCon US presentation, but there are two bonus examples included in this post. First, there is an example use case that is more complex. The code for this example can be found here. Consider the following simple network server:
The RequestServer is defined by an event loop and a port. Once the loop is started, the server can be started by calling the start_server() coroutine on line 16. This uses the asyncio.start_server() method to create a coroutine that waits for connections to self.porton line 17. Awaiting the coroutine on line 21 sets up a task on self.loop, which waits for connections to self.port and then serves the requests with the handler, handle_request(). This coroutine, defined on line 9, reads in a request on line 10. It then writes an Ack response on line 12, before draining and closing the connection on lines 13 and 14.
Running a functional test against this network server requires starting it in a running event loop, which works like this using LoopRunner:
The setUp() method is similar to what we’ve seen before, where a new event loop is created and run in a LoopRunner, self.runner. In our test method, we create a network service on line 13 and start it in the event loop running in self.runner’s thread on line 14. Then we use Pythons socket API to create a connection and send a test request on lines 16 through 19. On line 21 we get a response, and then compare it to our expected response. This does a fully functional, end-to-end test of the RequestServer class!
The second example is regarding the final bullet point in the list of more advanced topics — alternate event loops. For instance to test in uvloop using LoopRunner:
On line 6 we use uvloop.new_event_loop() instead of the native asyncio method to create the loop passed the LoopRunner. With that small change, we are able to test our code in an alternate event loop implementation.
To learn more, view the code or watch the video. Interested in working with our Engineering team? Check out open positions at Agari!
