Python中如何利用pytest-mock的高级特性模拟复杂的上下文管理器?
异步上下文管理器需用mocker.async_mock()模拟__aenter__(返回资源)和__aexit__(返回True/False),不可用普通patch;__aexit__必须显式设返回值,否则报TypeError。
如何模拟异步上下文管理器的 __aenter__ 和 __aexit__
异步上下文管理器(如 aiohttp.ClientSession 或自定义 async with 块)不能用普通 mocker.patch 直接替换,因为其 __aenter__ 和 __aexit__ 是协程方法,需返回 Awaitable。直接 mock 普通方法会导致 TypeError: object X is not async iterable 或 RuntimeWarning: coroutine 'X' was never awaited。
正确做法是使用 mocker.async_mock() 创建支持异步协议的模拟对象,并显式设置 __aenter__ 和 __aexit__ 的返回值:
-
__aenter__.return_value设为你要暴露给as变量的对象(比如字符串、字典或另一个 mock) -
__aexit__.return_value必须设为False(表示不压制异常)或True(压制),不能留空或设为None,否则会报TypeError: __aexit__ must return a bool - 若需验证异常处理逻辑,可在测试中主动抛出异常,并断言
__aexit__被调用且参数匹配
@pytest.mark.asyncio
async def test_async_context_manager(mocker):
async_mock = mocker.async_mock()
async_mock.__aenter__.return_value = {"data": "ok"}
async_mock.__aexit__.return_value = False # 不压制异常
async with async_mock as resource:
assert resource == {"data": "ok"}
async_mock.__aenter__.assert_called_once()
async_mock.__aexit__.assert_called_once()
模拟带状态的同步上下文管理器(如数据库连接)
很多上下文管理器内部维护状态(如连接是否已关闭、事务是否已提交),单纯 patch __enter__ 返回固定值无法覆盖这类行为。常见错误是忽略 __exit__ 的副作用,导致后续断言失败或资源泄漏误判。
推荐用 mocker.Mock 手动构造,并通过 side_effect 控制状态流转:
- 让
__enter__返回一个可变对象(如 dict 或自定义类实例),供后续操作修改 - 让
__exit__接收exc_type, exc_val, exc_tb并根据需要改变内部状态(例如设置.closed = True) - 避免用
mocker.patch.object(obj, '__enter__', ...)—— 它无法同时控制__exit__行为,容易漏掉清理逻辑
def test_db_connection_context(mocker):
mock_conn = mocker.Mock()
mock_conn.closed = False
def exit_side_effect(exc_type, exc_val, exc_tb):
mock_conn.closed = True
return False # 不压制异常
mock_conn.__enter__.return_value = mock_conn
mock_conn.__exit__.side_effect = exit_side_effect
with mock_conn:
assert not mock_conn.closed
assert mock_conn.closed
链式调用 + 上下文管理器混合场景(如 session.get(...).json())
当被测代码写成 async with session.get(url) as resp: return await resp.json() 这类嵌套结构时,mock 需要同时满足:1)session.get 返回一个支持 __aenter__ 的响应 mock;2)该响应 mock 的 json 方法是协程。
Python 3.14.3
微软官方的 Python 扩展,是 VS Code 安装量最高的扩展(209M+)。集成 IntelliSense(通过 Pylance)、调试(通过 Python Debugger)、代码检查、格式化、重构和单元测试等功能。支持 Jupyter Notebook、虚拟环境管理和多 Python 版本切换。
关键点在于分层构造,不能只 mock 最外层:
- 先用
mocker.async_mock()创建响应对象,设置json.return_value为期望数据 - 再用另一个
mocker.async_mock()创建 session mock,并让其get.return_value指向上一步的响应 mock - 确保所有中间对象都支持异步协议,否则
await resp.json()会报TypeError: object MagicMock is not awaitable
@pytest.mark.asyncio
async def test_session_get_json_chain(mocker):
# 模拟响应对象
mock_resp = mocker.async_mock()
mock_resp.json.return_value = {"status": "ok"}
# 模拟 session 对象
mock_session = mocker.async_mock()
mock_session.get.return_value = mock_resp
# 替换模块中的 session 实例
mocker.patch("my_module.get_session", return_value=mock_session)
result = await my_module.fetch_data()
assert result == {"status": "ok"}
mock_session.get.assert_called_once()
mock_resp.json.assert_awaited_once()
为什么 mocker.spy() 不适用于上下文管理器?
mocker.spy() 的设计目标是“监控真实方法调用而不改变行为”,但它对 __enter__/__exit__ 无效——因为这两个方法在 with 语句中由 Python 解释器隐式调用,不走常规属性访问路径。尝试对它们 spy 会触发 AttributeError: __enter__ 或静默失败。
真正需要监控上下文行为时,应改用完整 mock 替换整个上下文管理器对象,再通过 assert_called_once() 等验证其协议方法是否被触发。spy 更适合用于普通实例方法或函数,比如 db.commit() 或 logger.info()。
最容易被忽略的是:异步上下文管理器的 __aexit__ 即使没发生异常也会被调用,且必须返回布尔值;同步版本的 __exit__ 同理——漏掉这个返回值,mock 就不是合法的上下文管理器。