sqlalchemy - SQLAlchemy 连接的模拟异步上下文管理器

Question

我在使用自己的异步上下文管理器通过 SQLAlchemy 连接到数据库时测试代码时遇到问题。

# 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()

我收到错误:TypeError:对象 MagicMock 不能在“await”表达式中使用

有人知道如何处理吗？

Answer 1

我按如下方式解决了这个问题:

# test_async.py

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


class AsyncContextManager:

    async def __aenter__(self):
        pass

    async def __aexit__(self, exc_type, exc, traceback):
        pass


@patch("sqlalchemy.ext.asyncio.create_async_engine")
@pytest.mark.asyncio
async def test_async_query(mock_engine):
    mock_engine.return_value.connect = AsyncMock()
    mock_engine.return_value.dispose = AsyncMock()
    mock_conn = mock_engine.return_value.connect
    mock_conn.return_value.begin = MagicMock((AsyncContextManager))
    await query()

使用了AsyncMock并创建了具有神奇方法的AsyncContextManager类。