Python AsyncMock 怎样断言异步迭代器被正确消费

用 AsyncMock 测试 async for 时，最稳妥的断言分成两层：先把测试数据放进 mock.__aiter__.return_value，再断言消费函数拿到的列表；同时用 mock.__aiter__.assert_called_once_with() 确认异步迭代入口确实被调用。不要用 assert_awaited() 代替它，因为后者检查的是某个异步 mock 是否被 await。

官方地址：https://docs.python.org/3/library/unittest.mock.html

一、先分清 AsyncMock 的调用和异步迭代

Python 官方文档把 AsyncMock 定义为异步版本的 MagicMock：调用它会得到可等待对象，适合测试 await service.fetch() 这类协程调用。但 async for item in stream 走的是异步迭代协议，入口是对象的 __aiter__，随后由迭代器提供元素。

因此，测试目标如果是“异步流被完整消费”，至少要检查两个结果：业务函数收集到的值是否正确，以及 __aiter__ 是否被访问。单独检查 mock.called 或 assert_awaited，都可能把错误的测试写法当成通过。

二、用 __aiter__.return_value 提供异步 for 数据

一个简单的消费函数可以直接接收异步可迭代对象。配置 return_value 时传入列表即可，mock 会把它适配为异步迭代所需的对象；文章只关心协议和断言，不需要启动真实网络或消息服务。

import asyncio
from unittest.mock import AsyncMock


async def consume(stream):
    values = []
    # async for 会读取 stream 的异步迭代协议
    async for value in stream:
        values.append(value)
    return values


async def main():
    stream = AsyncMock(name="event_stream")
    # 列表是测试数据；AsyncMock 会把它作为异步迭代内容
    stream.__aiter__.return_value = ["created", "updated"]

    values = await consume(stream)

    # 先断言业务结果，确认两个元素都被消费且顺序正确
    assert values == ["created", "updated"]
    # 再断言异步迭代入口被调用，而不是断言 stream 被 await
    stream.__aiter__.assert_called_once_with()


asyncio.run(main())

这里的 stream 没有被直接调用，也没有被 await；它作为异步可迭代对象交给 async for。所以正确的行为证据是收集结果和 __aiter__ 调用记录。

三、把“消费正确”拆成结果断言和协议断言

只断言列表相等，能发现过滤、顺序和丢数据问题，但不能覆盖“代码根本没有进入异步迭代”的场景。只断言 __aiter__ 被调用，又无法保证循环体处理了每一项。两者组合后，测试意图会更清楚。

from unittest.mock import AsyncMock


async def collect_names(stream):
    names = []
    # 只保留名称，便于精确判断循环体的处理结果
    async for record in stream:
        names.append(record["name"])
    return names


async def test_collect_names():
    stream = AsyncMock(name="records")
    # 两条记录覆盖顺序和字段读取这两个关注点
    stream.__aiter__.return_value = [
        {"name": "alpha"},
        {"name": "beta"},
    ]

    names = await collect_names(stream)

    # 结果断言覆盖循环体是否处理了每条记录
    assert names == ["alpha", "beta"]
    # 协议断言覆盖代码是否真的从 stream 开始异步迭代
    stream.__aiter__.assert_called_once_with()

如果循环体被改成直接返回空列表，第一条断言会失败；如果代码改为读取另一个对象，第二条断言会失败。这样失败信息能对应到“数据处理错了”或“迭代对象用错了”两类问题。

四、工厂方法是异步的时要分开断言

有些接口不是直接传入异步迭代器，而是先 await client.open_stream() 得到一个流，再用 async for 消费。此时有两个对象：工厂方法的 AsyncMock 负责等待断言，返回值上的 __aiter__ 负责迭代断言，不能混为一谈。

from unittest.mock import AsyncMock


async def load_all(client):
    stream = await client.open_stream()
    values = []
    # 工厂返回的对象才是 async for 的输入
    async for value in stream:
        values.append(value)
    return values


async def test_load_all():
    client = AsyncMock(name="client")
    stream = AsyncMock(name="returned_stream")
    # 先配置异步工厂的返回对象，再配置这个对象的迭代内容
    client.open_stream.return_value = stream
    stream.__aiter__.return_value = [1, 2, 3]

    values = await load_all(client)

    # 工厂方法被 await 一次，使用 AsyncMock 专用断言
    client.open_stream.assert_awaited_once_with()
    # 返回的流被真正遍历一次，使用 __aiter__ 的调用断言
    stream.__aiter__.assert_called_once_with()
    assert values == [1, 2, 3]

这个写法把两个边界固定下来：open_stream 是否被正确等待，和返回流是否被正确消费。若只写 client.open_stream.assert_awaited_once_with()，即使循环没有遍历返回对象，测试也可能通过。

相关问题

为什么不能只用 assert_awaited 断言 async for？

因为 assert_awaited 检查的是 AsyncMock 自身或其方法被 await 的记录；直接交给 async for 的对象应检查 __aiter__ 和最终消费结果。

__aiter__.return_value 一定要写成异步生成器吗？

不需要。这个场景用普通列表就足够构造稳定测试数据；只有要模拟延迟、异常或更复杂的异步状态时，才需要单独准备符合异步迭代协议的对象。