为异步Python代码编写单元测试

最近使用 FastAPI 框架开发了一个 WEB 服务。

为了充分利用 FastAPI 作为一个 ASGI 框架的原生异步支持特性，很多业务代码也改成了异步函数，并且使用了异步的 HTTP 库httpx和 MongoDB 的异步 Python drivermotor。

由此带来的一个问题就是异步 Python 代码的单元测试的编写问题。

测试异步函数

编写测试代码

Python 的异步函数返回的是一个协程对象（coroutine），需要在前面加await才能获取异步函数的返回值，而只有在异步函数中才能使用await语句，这也意味着一般异步函数的测试代码本身也需要是一个异步函数。

async def add(a:int, b:int):
    return a + b

async def testAdd():
    ret = await add(1, 2)
    assert ret == 3

运行测试代码

与 Javascript 不同，Python 的异步代码需要显示地运行在事件循环中。

Python3.7 以上的版本中可以直接调用asyncio.run函。

如果使用的是更早的 Python 版本，就需要指定一个事件循环对象来运行异步代码。

import asyncio

# Python3.7+
asyncio.run(testAdd())

# Python3.6
loop = asyncio.new_event_loop()
loop.run_until_complete(testAdd())

使用 Pytest 运行异步测试代码

Pytest 是一个广为流行的 Python 测试框架，借助pytest-asyncio插件，我们可以更方便地编写异步测试代码。

# testasync.py
import pytest


async def add(a: int, b: int):
    return a + b


@pytest.mark.asyncio
async def testAdd():
    assert await add(1, 2) == 3

 pytest .\code\testasync.py
=============================================================================== test session starts ===============================================================================
platform win32 -- Python 3.8.6, pytest-6.2.1, py-1.10.0, pluggy-1.0.0.dev0
rootdir: C:\Users\duyix\Documents\md
plugins: asyncio-0.14.0
collected 1 item

code\testasync.py .                                                                                                                                                          [100%]

================================================================================ 1 passed in 0.04s ================================================================================

我们可以修改一下测试代码，让单元测试运行失败。

pytest .\code\testasync.py
=============================================================================== test session starts ===============================================================================
platform win32 -- Python 3.8.6, pytest-6.2.1, py-1.10.0, pluggy-1.0.0.dev0
rootdir: C:\Users\duyix\Documents\md
plugins: asyncio-0.14.0
collected 1 item

code\testasync.py F                                                                                                                                                          [100%]

==================================================================================== FAILURES =====================================================================================
_____________________________________________________________________________________ testAdd _____________________________________________________________________________________

    @pytest.mark.asyncio
    async def testAdd():
>       assert await add(1, 2) == 4
E       assert 3 == 4

code\testasync.py:10: AssertionError
============================================================================= short test summary info =============================================================================
FAILED code/testasync.py::testAdd - assert 3 == 4
================================================================================ 1 failed in 0.13s ================================================================================

mock 对象与异步测试

单元测试测试的是当前函数的行为，函数内部对于其他模块和组件的调用一般通过 mock 对象来模拟。

例如我们需要测试一个getIP函数，函数内通过向https://httpbin.org/ip接口发送请求来获取当前机器的 ip。

为了避免单元测试访问外部网络，同时消除在不同机器或者网络环境下getIP函数每次返回结果会不一样的影响，我们可以mock调网络请求部分的函数调用。

先看一下使用requests库的同步版本。

import requests
from unittest import mock


def getIP():
    resp = requests.get("https://httpbin.org/ip")
    return resp.json()["origin"]


@mock.patch("requests.get")
def testGetIP(mock_get):
    mock_response = mock.Mock()
    mock_response.json.return_value = {"origin": "127.0.0.1"}
    mock_get.return_value = mock_response
    assert getIP() == "127.0.0.1"
    mock_get.assert_called_once_with("https://httpbin.org/ip")

如果换一个asyncio的HTTP库的话，简单的mock就会失败。。

# getip.py
import httpx


client = httpx.AsyncClient()


async def getIP():
    resp = await client.get("http://httpbin.org/ip")
    return resp.json()["origin"]

#testhttpx.py
from getip import getIP
import pytest
from unittest import mock


@pytest.mark.asyncio
@mock.patch("getip.client")
async def testGetIP(mock_client):
    await getIP()

我们先把client对象mock掉来简单的调用一下getIP函数。

===================================================================================== FAILURES ======================================================================================
_____________________________________________________________________________________ testGetIP _____________________________________________________________________________________

mock_client = <MagicMock name='client' id='2180140905136'>

    @pytest.mark.asyncio
    @mock.patch("getip.client")
    async def testGetIP(mock_client):
>       await getIP()

testhttpx.py:9:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

    async def getIP():
>       resp = await client.get("http://httpbin.org/ip")
E       TypeError: object MagicMock can't be used in 'await' expression

getip.py:8: TypeError
============================================================================== short test summary info ==============================================================================
FAILED testhttpx.py::testGetIP - TypeError: object MagicMock can't be used in 'await' expression
================================================================================= 1 failed in 0.36s =================================================================================

可以看到默认的 mock 对象并不支持在await语句中使用。

解决方法也很简单，我们只需要指定需要mock的函数或方法的返回值为一个asyncio.Future对象。

A Future represents an eventual result of an asynchronous operation. Not thread-safe.
Future is an awaitable object. Coroutines can await on Future objects until they either have a result or an exception set, or until they are cancelled.

import asyncio
from unittest import mock

import pytest

from getip import getIP


@pytest.mark.asyncio
@mock.patch("getip.client")
async def testGetIP(mock_client):
    mock_response = mock.Mock()
    mock_response.json.return_value = {"origin": "127.0.0.1"}
    future = asyncio.Future()
    future.set_result(mock_response)
    mock_client.get.return_value = future
    assert await getIP() == "127.0.0.1"

我们也可以通过set_exception方法来指定asyncio.Future对象抛出的异常。

@pytest.mark.asyncio
@mock.patch("getip.client")
async def testGetIPFailed(mock_client):
    future = asyncio.Future()
    future.set_exception(Exception())
    mock_client.get.return_value = future
    with pytest.raises(Exception):
        await getIP()

值得注意的是如果不调用asyncio.Future对象的set_result方法或者set_exception方法的话，await语句会一直拿不到返回，程序会阻塞住。

总结

在这里总结一下异步 Python 代码的单元测试的要点：

测试代码也需要是异步代码
可以通过pytest-asyncio插件配合pytest简化异步测试代码的编写
对于需要mock的异步对象，可以指定mock方法或者函数返回一个asyncio.Future对象