Pytest：如何为asyncio.create_task调用的函数正确打桩？

问题原因 #

核心问题在于：asyncio.create_task创建的任务是异步调度的，不会阻塞当前协程执行。你的测试函数在await a.run()执行完成后就会退出，此时@patch的上下文已经结束，补丁被自动撤销，后续执行的任务就会调用原函数。

解决方案 #

方案1：在测试函数中等待所有未完成任务

修改测试函数，在await a.run()后等待事件循环中所有未完成任务执行完毕，确保补丁上下文在任务执行期间始终生效：

@pytest.mark.asyncio
@patch("main.another_func")
async def test_otg_test_low_disk_space_handler_fired(patched):
    a = A()
    await a.run()
    # 获取当前任务以外的所有未完成任务
    pending_tasks = asyncio.all_tasks() - {asyncio.current_task()}
    # 等待所有任务执行完成
    await asyncio.gather(*pending_tasks)

方案2：保存并显式等待创建的任务

先修改A类的run方法，保存创建的任务对象：

class A:
    async def func(self):
        print(f"!!!!!!!!!!!{another_func}!!!!!!!!!!!")
        another_func()
        
    async def run(self):
        await self.func()
        # 保存创建的任务
        self.pending_task = asyncio.create_task(self.func())

再修改测试函数，显式等待该任务完成：

@pytest.mark.asyncio
@patch("main.another_func")
async def test_otg_test_low_disk_space_handler_fired(patched):
    a = A()
    await a.run()
    # 等待保存的任务执行完成
    await a.pending_task

验证结果 #

修改后再次运行测试，两次输出都会显示MagicMock对象，说明补丁在create_task创建的任务中已生效。

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