Unit testing Python Coroutines

Photo by Content Pixie on Unsplash

In my last post we created python coroutines and executed them using asyncio. In this post we will look at how do we unit test such a code.

In my post Learn Python Coroutines we created generator and divisibility test functions. We worked with various examples and performed data exchange between the two. In a nutshell we have developed the following coroutines :

As discussed previously coroutines are executed in asynchronous manner. They can yield CPU for other executions. These behaviours makes them difficult to test.

Unit Testing

There are couple of ways we can try unit testing our code. I will stick to pytest for unit testing our code.

First Attempt

I am going to write unit test for countdown_async function. The function is responsible for generating a number and invoking the test function with the generated number. As stated earlier the challenge is to make the coroutine execute in a synchronous manner. The asyncio.run method comes to our rescue. It executes a coroutine in a blocking manner. Now, my test case looks like the following :

In the above code I have passed a mock method, where I collect all data. Since the test coroutine executes an await on the passed method, so I need to add asyncio.sleep to return back an awaitable

The above code collects all generated numbers. I am asserting just the size . More assertions can be added based on data.

Second Attempt

I am going to change the current testcase as a coroutine. The coroutine test will invoke the countdown_async coroutine and await on the result. I will need to added pytest-asyncio in the requirement.txt to support this.

I have made a couple of changes in the above testcase :

• Define the testcase as a coroutine with async keyword in the method definition
• Execute await on the countdown_async instead of asyncio.run
• decorate the testcase with @pytest.mark.asyncio

The rest of the testcase and asserts remain the same. The above changes make the testcase as an awaitable.

Mocking

Now, I will try to add a test case for divisibleBy_async coroutine. If we look into the implementation, it performs a divisibility check and then prints a message to the console. The method does not return any thing. So the first question to answer is “how do I validate the behaviour ?”

I can do assertions if I can get hold of the print function. But then it is not easy to mock the print method. Alternatively , I added an output method which prints to the console. I can mock this method. Now, my test case looks like the following :

In my test case I have done the following :

• mocked the output method and captured the invoked message. Since I am not doing an await on the output, so unlike above I do not need asyncio.sleep
• await on the method under test
• assert the data, captured in the mocked method

I will try to test the async_tasks coroutine. In this method I am executing tasks using asyncio.task api. Thus the mock method must be an awaitable. As a result my test case looks like the following :

I have used a nonlocal variable to count the invocations and validated it post the test execution.

Working with Queues

Let’s now build testcase for queues. I can create and use queue directly in my testcase. The testcase for countdown_async_with_queue is fairly simple.

I am asserting queue size for verification.

The above test case was quite simple. I will not build the test case for divisibleBy_with_queue. In the test case I need to create a queue with some data. The coroutine has the following interesting challenge :

The coroutines is a consumer which works in a while loop. It keeps executing and waits for the next data in the queue. If we invoke this directly then the execution may get stuck. Alternatively I can execute the coroutine as a task, which will be executed by the event-loop.

The complete test case looks like the following.

The above test case shows the following things :

• I have added a fixtureto setup the queue. The fixture is also marked as async. I have to do this as I can only add data in queue from a coroutine.
• asyncio.create_task is used to schedule the task. I need to capture the task reference. Post test the reference can be used to cancel the task
• I have to join on the queue

In summary testing python coroutines is complex. It will make you think about execution and data assertion. pytest-asyncio is an essential unit testing framework which offers good support for test coroutines.