如何Mock SQLAlchemy异步上下文管理器以测试数据库查询代码？

我在测试基于SQLAlchemy的自定义异步上下文管理器连接数据库的代码时遇到问题，运行测试时报错：

TypeError: object MagicMock can't be used in 'await' expression

相关代码 #

my_module.py

from contextlib import asynccontextmanager
from typing import Any, AsyncGenerator

from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.ext.asyncio.engine import AsyncConnection


@asynccontextmanager
async def adbcontext(url):
    engine = create_async_engine(url)
    conn = await engine.connect()
    try:
        async with conn.begin():
            yield conn
    finally:
        await conn.close()
        await engine.dispose()

async def query(url, sql):
    async with adbcontext(url) as conn:
        await conn.execute(sql)

test_async.py（原错误版本）

from unittest.mock import MagicMock
from asynctest import patch
import pytest
from my_module import query


@patch("sqlalchemy.ext.asyncio.create_async_engine")
@pytest.mark.asyncio
async def test_async_query(mock_engine):
    async def async_func(): pass
    mock_engine.return_value.__aenter__.connect = MagicMock(async_func)
    await query()

解决方案 #

错误根源在于mock设置不符合异步方法要求，且测试调用query时缺少必要参数，修复步骤如下：

1. 使用异步Mock对象：普通MagicMock不支持await语法，需用asynctest.MagicMock（或Python 3.8+自带的unittest.mock.AsyncMock）模拟异步方法。
2. 正确模拟SQLAlchemy异步对象链：create_async_engine返回的AsyncEngine的connect()是异步方法，需返回模拟的AsyncConnection；AsyncConnection的begin()需返回异步上下文管理器，execute()需是异步方法。
3. 传递query所需参数：调用query时必须传入url和sql参数。

修复后的测试代码：

from unittest.mock import MagicMock
from asynctest import patch, MagicMock as AsyncMagicMock  # 使用asynctest的MagicMock支持异步
import pytest
from my_module import query


@patch("sqlalchemy.ext.asyncio.create_async_engine")
@pytest.mark.asyncio
async def test_async_query(mock_engine):
    # 模拟AsyncConnection对象
    mock_conn = AsyncMagicMock()
    # 模拟begin()返回的异步上下文管理器
    mock_begin = AsyncMagicMock()
    mock_conn.begin.return_value.__aenter__.return_value = mock_conn
    mock_conn.begin.return_value.__aexit__.return_value = None
    
    # 模拟AsyncEngine的connect()异步方法，返回mock_conn
    mock_engine.return_value.connect.return_value = mock_conn

    # 调用query时传入必要参数
    test_url = "sqlite+aiosqlite:///:memory:"
    test_sql = "SELECT * FROM test_table"
    await query(test_url, test_sql)

    # 验证关键方法是否被调用
    mock_engine.assert_called_once_with(test_url)
    mock_engine.return_value.connect.assert_awaited_once()
    mock_conn.begin.assert_awaited_once()
    mock_conn.execute.assert_awaited_once_with(test_sql)

关键说明 #

• AsyncMagicMock（或AsyncMock）会自动处理异步方法的await逻辑，无需手动包装异步函数。
• 模拟begin()的异步上下文管理器时，需要同时实现__aenter__和__aexit__方法，确保async with conn.begin()能正常执行。
• 增加断言可以验证代码逻辑是否符合预期，比如确认connect()、begin()、execute()是否被正确调用。

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