如何使用pytest模拟asyncpg的executemany和fetch方法进行测试

前置依赖 #

先安装测试所需依赖：

pip install pytest pytest-asyncio pytest-mock

提示：你提供的待测试代码里config字典的键未加字符串引号，运行会触发NameError，需先修正为：

config = {"host": "localhost", "port": 5432, "database": "postgres"}

核心Mock思路 #

两个被测函数的调用链是asyncpg.create_pool() -> 异步上下文返回连接池pool -> pool.acquire()异步上下文返回连接conn -> 调用conn.fetch/conn.executemany，不能直接mock底层Connection类的方法，必须顺着异步上下文的调用链逐层打桩，所有异步方法统一用AsyncMock模拟，避免await时报类型错误。

完整测试代码 #

假设你的待测试代码存放在db_ops.py文件中，测试用例写法如下：

import pytest
import asyncpg
from db_ops import read_data, write_data

@pytest.mark.asyncio
async def test_read_data(mocker):
    # 构造模拟的查询返回结果
    mock_feed_data = [
        {"feed_id": 1, "url": "https://demo.com/rss1", "etag": "w/abc", "last_modified": "Thu, 01 Feb 2024 00:00:00 GMT"},
        {"feed_id": 2, "url": "https://demo.com/rss2", "etag": None, "last_modified": None}
    ]

    # 逐层构造mock对象
    mock_conn = mocker.AsyncMock()
    mock_conn.fetch.return_value = mock_feed_data

    mock_pool = mocker.AsyncMock()
    mock_pool.acquire.return_value.__aenter__.return_value = mock_conn

    mock_create_pool = mocker.patch.object(asyncpg, "create_pool", return_value=mocker.AsyncMock())
    mock_create_pool.return_value.__aenter__.return_value = mock_pool

    # 执行被测函数
    res = await read_data()

    # 断言逻辑
    assert res == mock_feed_data
    mock_create_pool.assert_called_once_with(host="localhost", port=5432, database="postgres")
    mock_conn.fetch.assert_called_once_with(
        "select feed_id, url, etag, last_modified from feeds WHERE enabled=TRUE"
    )


@pytest.mark.asyncio
async def test_write_data(mocker):
    # 构造测试入参
    test_items = [
        (1, "2024-02-01", "https://demo.com/item1", "guid-1", "测试标题1", "摘要1", "正文1", "作者1", 1),
        (2, "2024-02-02", "https://demo.com/item2", "guid-2", "测试标题2", "摘要2", "正文2", "作者2", 1)
    ]
    expect_sql = '''
                INSERT INTO
                    feed_items (feed_item_id, pubdate, link, guid, title, summary, content, author, feed_id) 
                VALUES
                    (
                        $1, $2, $3, $4, $5, $6, $7, $8, $9
                    )
                    ON CONFLICT DO NOTHING
            '''

    # 逐层构造mock对象
    mock_conn = mocker.AsyncMock()
    mock_pool = mocker.AsyncMock()
    mock_pool.acquire.return_value.__aenter__.return_value = mock_conn
    mock_create_pool = mocker.patch.object(asyncpg, "create_pool", return_value=mocker.AsyncMock())
    mock_create_pool.return_value.__aenter__.return_value = mock_pool

    # 执行被测函数
    await write_data(test_items)

    # 断言逻辑
    mock_create_pool.assert_called_once_with(host="localhost", port=5432, database="postgres")
    mock_conn.executemany.assert_called_once_with(expect_sql, test_items)

常见避坑点 #

• 不要直接mock asyncpg.Connection.fetch/asyncpg.Connection.executemany：代码中连接对象是通过连接池acquire动态生成的，直接打桩底层类不会拦截到调用
• 异步上下文管理器必须给__aenter__属性设置返回值：否则async with语法拿到的是默认MagicMock对象，不会走到你预设的mock逻辑
• 所有被await的方法必须用AsyncMock模拟：如果用普通MagicMock，运行时会抛出object MagicMock can't be used in 'await' expression错误
• 不需要额外引入asynctest等第三方mock库：pytest-mock 3.8.0以上版本内置的AsyncMock已经完全覆盖异步mock需求

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