如何在pytest中获取被测asyncio函数事件循环的所有任务？

我有如下fetcher.py模块代码，需要用pytest结合pytest-asyncio插件测试。尝试在tests.py中获取被测函数事件循环的所有任务，但断言失败。

测试代码（tests.py） #

import asyncio
import pytest
from fetcher import main

@pytest.mark.asyncio
async def test_01(self):
    await main(10)
    all_tasks = len(asyncio.all_tasks()) # 获取任务数量
    assert all_tasks == 10

断言错误 #

1 != 10

Expected :10
Actual   :1

我推测这是因为获取的是pytest事件循环的任务而非main(10)的事件循环任务，当前仅存在1个任务。

被测代码（fetcher.py） #

import asyncio


async def worker(queue: asyncio.Queue):
    while True:
        current_url = await queue.get()
        try:
            if current_url is None:
                break
            print(f"{current_url}: OK")
            print(len(asyncio.all_tasks())) # 输出10
        finally:
            queue.task_done()


async def main(max_requests):
    queue = asyncio.Queue(maxsize=max_requests)
    workers = {asyncio.create_task(worker(queue)) for _ in range(max_requests)}
    urls = ["https://example1.org/", "https://example2.org/", "https://example3.org/", "https://example4.org/"]
    for url in urls:
        await queue.put(url)
    for _ in range(len(workers)):
        await queue.put(None)
    await queue.join()

if __name__ == "__main__":
    asyncio.run(main(10))

问题核心是：你在await main(10)执行完毕后才调用asyncio.all_tasks()，此时main内部创建的10个worker任务已经全部完成并被事件循环销毁，只剩当前测试用例的任务，所以断言失败。

要验证main确实创建了10个任务，需要在任务运行期间（未被销毁时）检查，以下是几种可行方案：

方案1：修改main函数返回任务集合 #

直接让main返回创建的worker任务集合，测试时直接断言集合大小：

修改fetcher.py的main函数：

async def main(max_requests):
    queue = asyncio.Queue(maxsize=max_requests)
    workers = {asyncio.create_task(worker(queue)) for _ in range(max_requests)}
    urls = ["https://example1.org/", "https://example2.org/", "https://example3.org/", "https://example4.org/"]
    for url in urls:
        await queue.put(url)
    for _ in range(len(workers)):
        await queue.put(None)
    await queue.join()
    return workers  # 返回创建的任务集合

测试代码修改为：

import asyncio
import pytest
from fetcher import main

@pytest.mark.asyncio
async def test_01():
    workers = await main(10)
    assert len(workers) == 10

方案2：在任务运行时同步检查 #

如果不想修改原代码，可以用asyncio.Event同步测试代码和被测代码的执行时机，在worker任务活跃时检查数量：

修改fetcher.py的main函数，添加可选的检查事件参数：

async def main(max_requests, check_event=None):
    queue = asyncio.Queue(maxsize=max_requests)
    workers = {asyncio.create_task(worker(queue)) for _ in range(max_requests)}
    urls = ["https://example1.org/", "https://example2.org/", "https://example3.org/", "https://example4.org/"]
    for url in urls:
        await queue.put(url)
    
    # 触发测试检查点：所有worker已创建并运行
    if check_event:
        await check_event.wait()
    
    for _ in range(len(workers)):
        await queue.put(None)
    await queue.join()

测试代码：

import asyncio
import pytest
from fetcher import main

@pytest.mark.asyncio
async def test_01():
    check_event = asyncio.Event()
    # 异步启动main，避免阻塞测试任务
    main_task = asyncio.create_task(main(10, check_event))
    
    # 等待worker任务全部创建完成
    await asyncio.sleep(0.1)
    # 获取当前循环的所有任务，过滤掉测试任务和main任务
    all_tasks = asyncio.all_tasks(asyncio.get_running_loop())
    worker_tasks = [t for t in all_tasks if t not in (asyncio.current_task(), main_task)]
    
    assert len(worker_tasks) == 10
    
    # 允许main继续执行完成
    check_event.set()
    await main_task

方案3：直接在worker中验证任务数量 #

如果只是要确认任务创建数量，也可以在worker函数中添加断言（适合单元测试场景）：

async def worker(queue: asyncio.Queue, expected_task_count):
    while True:
        current_url = await queue.get()
        try:
            if current_url is None:
                break
            print(f"{current_url}: OK")
            # 验证当前任务数量是否符合预期
            assert len(asyncio.all_tasks()) == expected_task_count
        finally:
            queue.task_done()

然后修改main函数传递预期值，测试时只需确保main执行无异常即可。

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