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 one of a two-part supplement and addendum to my recent talk on “Strategies for Testing Async Code” at PyCon US 2019. I had the idea to propose this talk last November when I realized that my team and I have run into a number of challenges while writing tests for our asyncio codebases. I believed it was likely that there were others who would benefit from the solutions we found to these issues, so I quickly wrote up an outline of the challenges and solutions found, then hit submit on my session proposal. A few days later, I received an email from the PyCon program committee informing me of a typo in the proposed title. I fixed it, resubmitted, and then waited. To my surprise, I received an email in mid-February telling me that my talk was accepted!
Let’s start with a brief introduction to asynchronous programming and asyncio — the Python standard library API for writing asynchronous, concurrent application. Asyncio became part of the standard library in 3.4 in 2014, but the current API dates from 3.5. There are two primary concepts in asyncio.
Coroutines perform asynchronous work. They allow cooperative concurrency by ensuring that they perform I/O, and other non-CPU-intensive operations in a non-blocking manner. Specifically, they are functions that return objects that represent a computation and/or an I/O operation that will eventually complete.
The async keyword is used to define a coroutine function, and the await keyword is used to wait for the eventual results of the object returned by such a function. Here is an example coroutine:
In the above example, my_coroutine defines with the async keyword on line 1 returns the awaited results of another_coroutine on line 2.
Event Loops schedule asynchronous work. They are schedulers that run coroutines that are not waiting for I/O or another coroutine to finish their computation, as well as perform the I/O operations for the coroutines. The top-level asyncio library provides a few methods for managing event loops. The two we’ll see in this post include:
- asyncio.get_event_loop() returns the current event loop. If there is no current event loop set in the current OS thread, a new event loop will be created and returned.
- asyncio.new_event_loop() creates and returns a new event loop.
Example Asyncio Use Case: Herding Cats
In order to present the concepts of asyncio without relying on much understanding of a complex use case, I used the metaphor of herding cats, which implies many disjointed things happening at the same time, which may eventually complete. I designed a simple cat object:
Attempts to move() (line 14) a Cat will only succeed if the cat is PLEASED, and even then, it will take some time for the cat to move. When you pet() a Cat you make the cat PLEASED. Because of the asyncio.sleep() on line 16, this is defined as a coroutine.
With this object in mind, I created a simple coroutine to use in test cases:
This coroutine is given a cat and a direction to move it. It calls pet() on line 2 to make the cat amenable and then awaits the result of the move() coroutine on line 3.
Challenge: How to Test a Coroutine
You cannot just call the coroutine as a function and expect to test the results, because the results of calling it are an object of eventual completion. You cannot just await the results, because await can only be called from inside coroutine. What you have to do is ensure that you do the test from a coroutine running inside an event loop. Here is an example of how to schedule and run a coroutine inside of an event loop and test the response, using the standard library unittest framework:
On line 15 we get an event loop, which we use on line 16 to run our coroutine under test in. We are able to use the result in our test assertion on line 18.
I also provide an example of using a pytest plugin, pytest-asyncio, which allows tests to be defined as a coroutine and implicitly run in an event loop, as well as an example of using a new method asyncio.run() available as of 3.7. All three of these solutions allow you to unit test a single coroutine with relative ease.
Challenge: How to Mock a Coroutine
Mocking and patching are commonly used in python unit testing. However, you cannot mock a coroutine with the stdlib unittest.mock framework. I presented what is essentially a proxy mock which subclasses unittest.mock.MagicMock:
The one method present, __call__ on line 2, defines the behavior when the object is used with () — or called as a function. In this case, it is defined in a coroutine context via asyncdef, so it controls how the object behaves when called as a coroutine. In effect, it returns a coroutine object that returns the __call__ method of the superclass, which allows you to treat it as a MagicMock for assertion purposes. This is how to use it in a test case:
Challenge: How to Mock an Async Context Manager
Another useful feature of asyncio is asynchronous context managers. These allow context managers to suspend execution (via await) in their setup (enter) and teardown (exit) phases. This is done by defining __aenter__ and __aexit__ methods. Unittest.mock.MagicMockaccepts the standard python magic methods by default, but not the newer magic coroutines required for use in async with. Similar to our AsyncMock above, we can implement a mock context manager with defining some magic coroutines:
By defining the __aenter__ and __aexit__ magic coroutines, and simply passing them through to the standard context manager magic methods, we can use AsyncContextManager in any async with statement.
To learn more, view the code or the video, and stay tuned for more advanced use cases.
