如何用pytest在异步环境中Mock多个上下文管理器与函数调用？

业务代码（需要Mock的部分） #

async with await self._db.get_pool().getconn() as conn:
    async with await conn.cursor() as cursor:
        await cursor.execute(...)
        result = await cursor.fetchone()
        ...
        return Object(id=result[0]...)

注：self._db为自定义DB类型，execute为数据库查询操作，返回值由fetchone的结果元组填充。

当前Mock配置 #

mock_db = AsyncMock(spec=DB)
mock_pool = AsyncMock(spec=AsyncConnectionPool)
mock_conn = AsyncMock(spec=AsyncConnection)
mock_cursor = AsyncMock(spec=AsyncCursor)

mock_db.__aenter__.return_value = mock_db
mock_db.get_pool.return_value = mock_pool
mock_pool.getconn.return_value = mock_conn
mock_conn.__aenter__.return_value = mock_conn
mock_conn.cursor.return_value = mock_cursor
mock_cursor.__aenter__.return_value = mock_cursor
mock_cursor.fetchone.return_value = (uuid(), ....)
result = await test_func(...)

问题原因与修复方案 #

核心问题 #

使用spec参数时，如果原类的__aenter__是同步方法，AsyncMock会生成普通Mock实例而非异步版本，导致async with调用时出错。此外，异步方法的Mock配置需确保返回值正确关联协程结果。

修改后的Mock配置 #

mock_db = AsyncMock(spec=DB)
mock_pool = AsyncMock(spec=AsyncConnectionPool)
mock_conn = AsyncMock(spec=AsyncConnection)
mock_cursor = AsyncMock(spec=AsyncCursor)

# 显式将异步上下文管理器方法设为AsyncMock
mock_db.__aenter__ = AsyncMock(return_value=mock_db)
mock_db.get_pool.return_value = mock_pool

# getconn是异步方法，直接用AsyncMock定义并返回mock_conn
mock_pool.getconn = AsyncMock(return_value=mock_conn)
mock_conn.__aenter__ = AsyncMock(return_value=mock_conn)

# cursor是异步方法，用AsyncMock定义并返回mock_cursor
mock_conn.cursor = AsyncMock(return_value=mock_cursor)
mock_cursor.__aenter__ = AsyncMock(return_value=mock_cursor)

# 设置fetchone返回值
mock_cursor.fetchone.return_value = (uuid.uuid4(), ....)

result = await test_func(...)

额外检查点 #

1. Mock注入验证：确保测试函数使用的是mock_db而非真实DB实例，可通过unittest.mock.patch装饰器或手动注入实现。
2. execute方法Mock：若业务代码依赖execute的返回值，需添加mock_cursor.execute.return_value = ...配置；若需断言调用，后续可通过mock_cursor.execute.assert_called_once_with(...)验证。
3. uuid依赖处理：若业务代码中使用uuid生成值，需确保测试中返回的元组与业务逻辑预期一致，或Mockuuid.uuid4方法固定返回值。

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