[FIXED] How can I mock asyncpg executemany and fetch methods using pytest?

Issue

I am using pytest and asyncpg to load and save items to the database. How do I test the 2 functions below?

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)

Solution

We can mock the database calls, however there isn't much to test when doing that. Your functions simply read/write some data, there isn't much logic or conditionals to verify. My recommendation would be to spin up a Postgres instance in a Docker container alongside your tests and test by hitting that.

With that being said, below you can find how to mock your database calls and test some basic functionality. You will need to install pytest-asyncio to get this to work.

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 ================================================================

Answered By - gold_cy