Mocking asynchronous context manager for SQLAlchemy connection

Question

I have a problem testing the code where I connect to the database via SQLAlchemy using my own asynchronous context manager.

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

I get the error: TypeError: object MagicMock can't be used in 'await' expression

Does anyone know how to deal with this?

score 1 · Answer 1 · answered Oct 28 '22 at 10:17

I solved this problem as follows:

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

used AsyncMock and created AsyncContextManagerclass that have magic methods.

Mocking asynchronous context manager for SQLAlchemy connection

1 Answers1