如何对asyncio数据库上下文管理器进行Mock？

问题场景 #

在模拟aiomysql的典型异步数据库连接流程时，出现了AttributeError: 'AsyncMock' object has no attribute '__aenter__'错误。业务代码的三层异步嵌套结构如下：

async with aiomysql.create_pool(...) as pool:
    async with pool.acquire() as connection:
        async with connection.cursor() as cursor:
            await cursor.execute("BEGIN")
            ...

初次编写的测试代码尝试用AsyncMock模拟各层对象，但运行时触发了上述错误：

async def test_database(mocker: pytest_mock.MockerFixture):
    context = mocker.AsyncMock()
    pool = mocker.AsyncMock()
    connection = mocker.AsyncMock()
    cursor = mocker.AsyncMock()
    cursor.fetchall.return_value = [{'Database': 'information_schema'}]
    cursor.fetchone.return_value = {'COUNT(*)': 0}
    cursor.rowcount = 0
    connection.cursor.return_value.__aenter__.return_value = cursor
    pool.acquire.return_value.__aenter__.return_value = connection
    context.__aenter__.return_value = pool
    mocker.patch('aiomysql.create_pool', return_value=context)

    async with aiomysql.create_pool() as p:
        async with p.acquire() as c:
            async with c.cursor() as cur:
                await cur.execute("BEGIN")

错误原因 #

错误核心在于：pool.acquire()或connection.cursor()返回的AsyncMock对象未被正确配置为异步上下文管理器。虽然代码中设置了return_value.__aenter__.return_value，但如果返回的AsyncMock本身未生成__aenter__和__aexit__方法（或因版本兼容问题未自动生成），就会触发属性缺失错误。

解决方案 #

确保每个异步上下文方法的返回对象是合法的异步上下文管理器即可，以下是两种可靠写法：

写法一：明确创建上下文模拟对象 #

async def test_database(mocker: pytest_mock.MockerFixture):
    # 模拟cursor及返回值
    cursor = mocker.AsyncMock()
    cursor.fetchall.return_value = [{'Database': 'information_schema'}]
    cursor.fetchone.return_value = {'COUNT(*)': 0}
    cursor.rowcount = 0

    # 模拟connection：cursor()返回的上下文管理器进入后得到cursor
    connection = mocker.AsyncMock()
    cursor_context = mocker.AsyncMock()
    cursor_context.__aenter__.return_value = cursor
    connection.cursor.return_value = cursor_context

    # 模拟pool：acquire()返回的上下文管理器进入后得到connection
    pool = mocker.AsyncMock()
    conn_context = mocker.AsyncMock()
    conn_context.__aenter__.return_value = connection
    pool.acquire.return_value = conn_context

    # 模拟create_pool返回的上下文管理器进入后得到pool
    pool_context = mocker.AsyncMock()
    pool_context.__aenter__.return_value = pool
    mocker.patch('aiomysql.create_pool', return_value=pool_context)

    # 执行测试逻辑
    async with aiomysql.create_pool() as p:
        async with p.acquire() as c:
            async with c.cursor() as cur:
                await cur.execute("BEGIN")

写法二：简化链式设置 #

利用AsyncMock自动生成上下文方法的特性，直接链式配置返回值：

async def test_database(mocker: pytest_mock.MockerFixture):
    cursor = mocker.AsyncMock()
    cursor.fetchall.return_value = [{'Database': 'information_schema'}]
    cursor.fetchone.return_value = {'COUNT(*)': 0}
    cursor.rowcount = 0

    # 配置connection的cursor上下文返回cursor
    connection = mocker.AsyncMock()
    connection.cursor.return_value.__aenter__.return_value = cursor

    # 配置pool的acquire上下文返回connection
    pool = mocker.AsyncMock()
    pool.acquire.return_value.__aenter__.return_value = connection

    # 配置create_pool的上下文返回pool
    mocker.patch('aiomysql.create_pool').return_value.__aenter__.return_value = pool

    # 执行测试逻辑
    async with aiomysql.create_pool() as p:
        async with p.acquire() as c:
            async with c.cursor() as cur:
                await cur.execute("BEGIN")

核心逻辑是：让create_pool()、pool.acquire()、connection.cursor()的返回值，都成为带有__aenter__和__aexit__方法的异步上下文管理器，并正确设置__aenter__的返回对象。

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