我如何使用pytest模拟asyncpg executemany和fetch方法?

Question

我使用pytest和asyncpg来加载和保存项目到数据库。如何测试下面的两个函数?

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

async def read_data():
async with asyncpg.create_pool(**config) as pool:
async with pool.acquire() as conn:
feeds = await conn.fetch(
"select feed_id, url, etag, last_modified from feeds WHERE enabled=TRUE"
)
return feeds

async def write_data(items):
async with asyncpg.create_pool(**config) as pool:
async with pool.acquire() as conn:
query = '''
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
'''
await conn.executemany(query, items)

Answer 1

我们可以模拟数据库调用，但是这样做时没有太多要测试的。你的函数只是读/写一些数据，没有太多的逻辑或条件来验证。我的建议是在Docker容器中启动一个Postgres实例，然后点击它进行测试。

话虽如此，下面您可以找到如何模拟数据库调用和测试一些基本功能。您需要安装pytest-asyncio才能使其工作。

from unittest.mock import patch
import pytest
import pytest_asyncio
from src.get_events import read_data, write_data

# fixture which mocks database behavior
# take note that you have to mock __aenter__ on the calls as you are
# using with statements in your code
@pytest_asyncio.fixture
async def db_conn():
with patch("src.get_events.asyncpg") as mock_pool:
mock_pool.configure_mock(
**{
"create_pool.return_value.__aenter__.return_value": mock_pool,
"acquire.return_value.__aenter__.return_value": mock_pool
}
)
yield mock_pool

@pytest.mark.asyncio
async def test_read_data(db_conn):
expected = [1, 2, 3]
expected_query = "select feed_id, url, etag, last_modified from feeds WHERE enabled=TRUE"
async def fetch():
return expected

# async function calls need to return an awaitable
db_conn.fetch.return_value = fetch()
feeds = await read_data()

assert feeds == expected
# assert function was called with the query we expect
db_conn.fetch.assert_called_once_with(expected_query)

@pytest.mark.asyncio
async def test_write_data(db_conn):
async def sink(*args, **kwargs):
return
items = [3, 4]
db_conn.executemany.return_value = sink()
await write_data(items)
db_conn.executemany.assert_called_once()
kall = db_conn.executemany.call_args
# confirm that the items passed in to the function are the same ones
# being sent to the database
assert items in kall.args

---

=============================================================== test session starts ===============================================================
platform darwin -- Python 3.8.9, pytest-7.0.1, pluggy-1.0.0
cachedir: .pytest_cache
rootdir: **
asyncio: mode=strict
collected 2 items                                                                                                                                 
tests/test_script.py::test_read_data PASSED
tests/test_script.py::test_write_data PASSED
================================================================ 2 passed in 0.03s ================================================================