模拟 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: object MagicMock can't be used in 'await' expression

有谁知道如何处理这个？

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 并创建了具有魔术方法的AsyncContextManagerclass。