如何用Pytest测试含无限循环与await的Python异步函数？

问题背景 #

我正在对一段包含asyncio循环的Python代码进行单元测试，代码中的两个核心异步函数均带有无限循环，非异步逻辑部分（sync_receive_data_packets和sync_data_analysis）已完成测试。现需要测试receive_data和analyse_data这两个异步函数，尤其是它们内部的await queue.put()和await queue.get()操作——这两个函数处于无限循环中，无返回值，仅分别执行队列的存入与读取操作。

对应的业务代码：

import asyncio

def main():
    loop = asyncio.get_event_loop()
    loop.run_until_complete(run_async_tasks())
    .
    .
    

async def run_async_tasks():
    # initialize loop and its variables
    loop = asyncio.get_event_loop()
    queue = asyncio.Queue(maxsize=100)

    # Creating Asyncio tasks
    task1 = loop.create_task(receive_data(queue), name="receive_data")
    task2 = loop.create_task(analyse_data(queue), name="analyse_data")

    await asyncio.gather(task1, task2)

async def receive_data(queue):
    # Data packets are yielded from an non-ending stream by sync_receive_data_packets
    for data_packet in sync_receive_data_packets():
        await queue.put(data_packet)

async def analyse_data(queue):
    while not termination_signal_received():
        data_packet = await queue.get()
        sync_data_analysis(data_packet)
        queue.task_done()

测试方案 #

核心思路是通过Mock控制循环终止条件，同时验证队列的交互行为，推荐用pytest+pytest-asyncio编写异步测试用例。

1. 测试receive_data函数 #

验证它能否正确从同步生成器取数据并写入队列：

import pytest
import asyncio
from unittest.mock import patch

@pytest.mark.asyncio
async def test_receive_data():
    # 准备有限的测试数据包，避免无限循环
    test_packets = [{"data": 1}, {"data": 2}, {"data": 3}]
    
    # Mock同步生成器，返回测试数据
    with patch("__main__.sync_receive_data_packets", return_value=test_packets):
        queue = asyncio.Queue()
        # 创建并运行任务
        task = asyncio.create_task(receive_data(queue))
        await task
        
        # 验证队列数据正确性
        assert queue.qsize() == len(test_packets)
        for expected in test_packets:
            assert await queue.get() == expected

2. 测试analyse_data函数 #

验证它能否正确从队列取数据、调用分析函数并标记任务完成：

from unittest.mock import Mock, patch

@pytest.mark.asyncio
async def test_analyse_data():
    # Mock终止信号：调用3次后返回True，终止循环
    terminate_count = 0
    def mock_terminate():
        nonlocal terminate_count
        terminate_count += 1
        return terminate_count >= 3
    
    # Mock分析函数，记录调用情况
    mock_analyse = Mock()
    
    with patch("__main__.termination_signal_received", mock_terminate):
        with patch("__main__.sync_data_analysis", mock_analyse):
            queue = asyncio.Queue()
            # 预先放入测试数据
            test_packets = [{"data": 1}, {"data": 2}]
            for pkt in test_packets:
                await queue.put(pkt)
            
            # 运行任务
            task = asyncio.create_task(analyse_data(queue))
            await task
            
            # 验证分析函数的调用次数和参数
            assert mock_analyse.call_count == len(test_packets)
            for pkt in test_packets:
                mock_analyse.assert_any_call(pkt)
            # 验证队列任务全部完成
            assert queue.empty()
            await queue.join()  # 无阻塞，说明所有task_done已调用

3. 测试两个函数的协同工作 #

验证数据从receive_data写入后，analyse_data能正确读取并处理：

@pytest.mark.asyncio
async def test_receive_and_analyse_cooperation():
    test_packets = [{"data": 1}, {"data": 2}]
    mock_analyse = Mock()
    
    with patch("__main__.sync_receive_data_packets", return_value=test_packets):
        with patch("__main__.termination_signal_received", side_effect=[False, False, True]):
            with patch("__main__.sync_data_analysis", mock_analyse):
                queue = asyncio.Queue()
                # 启动两个任务
                task_receive = asyncio.create_task(receive_data(queue))
                task_analyse = asyncio.create_task(analyse_data(queue))
                
                await asyncio.gather(task_receive, task_analyse)
                
                # 验证数据被完整处理
                assert mock_analyse.call_count == len(test_packets)
                for pkt in test_packets:
                    mock_analyse.assert_any_call(pkt)

关键技巧总结 #

• 控制循环终止：通过Mock同步依赖（生成器或终止信号），把无限循环转为有限执行，避免测试卡死
• 验证队列交互：检查队列的大小、内容，以及task_done的执行情况（通过queue.join()验证）
• 异步测试框架：pytest-asyncio自动管理事件循环，简化异步测试的运行逻辑

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