如何模拟异步数据库上下文管理器？

Question

我一直在努力模拟典型的异步数据库连接设置：

async with aiomysql.create_pool(...) as pool:
    async with pool.acquire() as connection:
        async with connection.cursor() as cursor:
            await connection.begin()
            ...

我第一次尝试测试函数看起来像这样：

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 c.begin()

如果您因缺少__aenter__s 而收到AttributeErrors，那么这篇文章适合您。

Answer 1

需要注意的重要部分是async with 和函数调用之间没有await，因为create_pool、acquire 和cursor 是同步的。上面的测试函数将生成新的AsyncMock 对象，这些对象需要在acquire() 上等待以返回下一个准备好的AsyncMock。相反，我们希望acquire() 等立即返回。解决方案是混合Mock/MagicMock和AsyncMock。

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

    # calls create_pool synchronously and gets 'context',
    # which is an AsyncMock and facilitates the __aenter__ call,
    # which returns 'pool' as a regular Mock
    async with aiomysql.create_pool() as p:
        # calls 'acquire()' synchronously and gets an anonymous AsyncMock,
        # which facilitates the __aenter__ call,
        # which returns 'connection' as a regular Mock
        async with p.acquire() as c:
            # again, 'cursor()' synchronously, get AsyncMock,
            # __aenter__ to get 'cursor'
            async with c.cursor() as cur:
                # continue regular operations on AsyncMock object
                await c.begin()