如何在pytest中Mock asyncio.ensure_future()与await asyncio.gather(*tasks)

首先先把你的代码整理成可测试的结构化形式（假设主逻辑封装在fetch_all_data函数中）：

import asyncio
import httpx

data = {}

async def get_info(client, url, key):
    resp = await client.get(url)
    if resp.status_code == 200:
        res = resp.json()
    else:
        res = None
    return key, res

async def fetch_all_data():
    lis = []
    async with httpx.AsyncClient() as client:
        lis.append(asyncio.ensure_future(get_info(client, "example/api1", "api_1")))
        lis.append(asyncio.ensure_future(get_info(client, "example/api2", "api_2")))
        lis.append(asyncio.ensure_future(get_info(client, "example/api3", "api_3")))
        infos = await asyncio.gather(*lis)
        for info in infos:
            data[info[0]] = info[1]
    return data

1. 准备依赖 #

先安装测试所需的包：

pip install pytest pytest-asyncio httpx

2. 测试get_info函数 #

核心是Mock httpx.AsyncClient.get 方法，模拟成功和失败两种响应场景：

import pytest
import asyncio
from unittest.mock import Mock
from your_module import get_info  # 替换成你的实际模块名

@pytest.mark.asyncio
async def test_get_info_success():
    # 构造Mock响应和客户端
    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"content": "api1_result"}
    
    mock_client = Mock()
    mock_client.get.return_value = asyncio.coroutine(lambda: mock_response)()
    
    # 执行测试
    key, result = await get_info(mock_client, "example/api1", "api_1")
    
    # 断言结果
    assert key == "api_1"
    assert result == {"content": "api1_result"}
    mock_client.get.assert_called_once_with("example/api1")

@pytest.mark.asyncio
async def test_get_info_failure():
    mock_response = Mock()
    mock_response.status_code = 404
    
    mock_client = Mock()
    mock_client.get.return_value = asyncio.coroutine(lambda: mock_response)()
    
    key, result = await get_info(mock_client, "example/api1", "api_1")
    
    assert key == "api_1"
    assert result is None
    mock_client.get.assert_called_once_with("example/api1")

3. 测试主函数fetch_all_data #

有两种Mock思路，按需选择：

方式一：Mock httpx.AsyncClient（贴近真实请求场景） #

使用pytest-mock的mocker fixture直接patch HTTP客户端的get方法：

import pytest
import httpx
from your_module import fetch_all_data  # 替换成你的实际模块名

@pytest.mark.asyncio
async def test_fetch_all_data(mocker):
    # 为不同URL设置不同的Mock响应
    mock_get = mocker.patch.object(httpx.AsyncClient, 'get')
    mock_get.side_effect = [
        Mock(status_code=200, json=lambda: {"content": "api1_data"}),
        Mock(status_code=200, json=lambda: {"content": "api2_data"}),
        Mock(status_code=500, json=lambda: {})
    ]
    
    # 执行主函数
    result_data = await fetch_all_data()
    
    # 断言结果
    assert result_data == {
        "api_1": {"content": "api1_data"},
        "api_2": {"content": "api2_data"},
        "api_3": None
    }
    
    # 验证调用次数和参数
    assert mock_get.call_count == 3
    mock_get.assert_any_call("example/api1")
    mock_get.assert_any_call("example/api2")
    mock_get.assert_any_call("example/api3")

方式二：直接Mock get_info（聚焦主逻辑正确性） #

如果只关心主函数的结果收集逻辑，可以直接Mock get_info函数：

import pytest
from your_module import fetch_all_data, get_info  # 替换成你的实际模块名

@pytest.mark.asyncio
async def test_fetch_all_data_mock_get_info(mocker):
    # 预设get_info的返回值
    mock_get_info = mocker.patch('your_module.get_info')
    mock_get_info.side_effect = [
        ("api_1", {"content": "mock1"}),
        ("api_2", {"content": "mock2"}),
        ("api_3", None)
    ]
    
    result_data = await fetch_all_data()
    
    assert result_data == {
        "api_1": {"content": "mock1"},
        "api_2": {"content": "mock2"},
        "api_3": None
    }
    
    # 验证get_info的调用参数
    assert mock_get_info.call_count == 3
    mock_get_info.assert_any_call(Mock(), "example/api1", "api_1")
    mock_get_info.assert_any_call(Mock(), "example/api2", "api_2")
    mock_get_info.assert_any_call(Mock(), "example/api3", "api_3")

关键注意事项 #

• 所有异步测试函数必须加上@pytest.mark.asyncio装饰器，pytest才会用异步运行器执行。
• Mock异步方法时，要确保返回的是协程对象（可以用asyncio.coroutine包装，或者直接定义异步Mock函数）。
• 使用side_effect可以为多次调用返回不同结果，适合模拟多URL的差异化响应。

内容的提问来源于stack exchange，提问作者Akash Kumar