You might have already heard about Python's asyncio
module, it allows you to easily run concurrent
code using Python.
In the last few months I've worked in some codebases which take advantage of the benefits of
asyncio
, mainly through their use of aiohttp.
When using asyncio
you'll use the async/await
keywords to both define and call
asynchronous functions.
This also means changes in the way you test your code because, unlike ordinary functions,
asynchronous functions always return a coroutine object, which needs to be awaited, using the
await
keyword in order to actually schedule it, run it and get the actual return value.
As such, let's take a quick look into how we can easily test asynchronous functions by leveraging Futures in Python 3.7 and AsyncMock in Python 3.8
In this example, we're going to be mocking a simple function which adds two integers, its parameters,
while resorting to asyncio.sleep
to simulate IO heavy tasks, for example, HTTP requests or a
database calls.
import asyncio
async def sum(x, y):
await asyncio.sleep(1)
return x + y
Asynchronous functions in Python return what's known as a Future
object, which contains the
result of calling the asynchronous function.
As such, the "secret" to mocking these functions is to make the patched function return a
Future
object with the result we're expecting, as one can see in the example below.
import pytest
import asyncio
@pytest.fixture()
def mock_sum(mocker):
future = asyncio.Future()
future.set_result(4)
mocker.patch('app.sum', return_value=future)
As you can see in the example above, we're creating a pytest fixture, namely mock_sum
that
patches the function we created at the beginning of the post and specifies that the function call
will return a Future
object, with a result of 4
.
In your own tests you will, of course, need to change the call to set_result
to return whatever
value you're expecting, maybe a HTTP response or some database query result.
With this done we can now create a simple test case that tests the sum
function:
import pytest import asyncio
@pytest.mark.asyncio
async def test_sum(mock_sum):
result = await sum(1, 2)
# I know 1+2 is equal to 3 but one man can only dream!
assert result == 4
There's also a few different things happening here when compared to a regular test function:
@pytest.mark.asyncio
decorator - This tells pytest that this is an asynchronous test function,
otherwise pytest will skip it.async def test_sum(mock_sum)
- Defines the asynchronous test function while at the same time
calls the pytest fixture, mock_sum
, so that it successfully mocks the sum
function's result.result = await sum(1,2)
- Correctly calls the asynchronous function using the await
keyword.Although 1 + 2
is equal to 3
I'm purposefully asserting that this returns 4
so as to make sure
that the fixture is indeed called.
If you go ahead and run pytest
now with the code shown above you should see that indeed it
executes successfully, passing the tests.
However, imagine that you want to mock the sum
function multiple times while having a different
value provided to set_result
in the Future
object. It doesn't make sense to create multiple
fixtures since we'll be repeatedly patching the same function. In this case we'll return the
Future
object and call the set_result
function in the test function, thus,
our fixture we'll now look like:
import pytest
import asyncio
@pytest.fixture()
def mock_sum(mocker):
future = asyncio.Future()
mocker.patch('app.sum', return_value=future)
return future
Notice that we're not calling set_result
in the fixture this time around. With the updated
fixture we now need to update the test function to look like this:
import pytest
import app
@pytest.mark.asyncio
async def test_sum(mock_sum):
mock_sum.set_result(4)
result = await app.sum(1, 2)
# I know 1+2 is equal to 3 but one man can only dream!
assert result == 4
Finally, notice now how we're calling mock_sum.set_result(4)
. If we want the mock to return
different values we now just need to change the value provided to set_result
instead of having to
create multiple fixture for different tests!
The code above only works for versions of Python <3.8. In Python 3.8 we need to change the code
slightly because
AsyncMock
has
been introduced.
With that said, we can simply change the mocking function to return the AsyncMock
instance instead
of the Future
instance.
from unittest.mock import AsyncMock
@pytest.fixture()
def mock_sum(mocker):
async_mock = AsyncMock(return_value=4)
mocker.patch('app.sum', side_effect=async_mock)
As you can see in the code above, the main change is that the return value is now set as an
AsyncMock
instance instead of a Future
instance, and we can also now use the return_value
argument in the AsyncMock
instantiation instead of needing to call a function afterwards to set
its result.
With the code above our test function will look like the first showed in this blog post, where we
don't change the result of the mock. However, as we did in the end of the previous section, if we
need to mock the same function multiple times while having different results it's better if we
just return the AsyncMock
instance from the fixture and set the return_value
in the test
function. As such, our fixture would now look like this:
from unittest.mock import AsyncMock
@pytest.fixture()
def mock_sum(mocker):
async_mock = AsyncMock()
mocker.patch('app.sum', side_effect=async_mock)
return async_mock
And with this fixture we could simply update our test function to the following:
@pytest.mark.asyncio
async def test_sum(mock_sum):
mock_sum.return_value = 4
result = await app.sum(1, 2)
assert result == 4
Notice that the only change compared to the previous section is that we now set the return_value
attribute of the mock instead of calling the set_result
function seeing as we're now working with
AsyncMock
instead of Future
. Aside that, the test function looks exactly the same.
In conclusion mocking asynchronous functions in Python is actually easier than I expected at first, mostly because I didn't really understood how asyncio worked. After some reading and experimentation it turns out it's quick and easy to do, and it allows you to run concurrent tests, which should speed up your test suite!
If you're reading this for a quick solution and don't really known what's going on I'd advise reading up on asyncio.
I hope this blogpost has helped you! 👋