How to Mock PostgreSQL with Pytest and pytest-postgresql

We mock PostgreSQL tables using Docker and pytest-postgresql, allowing unit testing for application code queries without a remote database. This allows our tests to cover both the SQL logic and the Python code at the same time, all without hitting network resources.

Is Mocking PostgreSQL a Good Idea?

It is best for simple, unchanging schemas. For a complex database-driven application, there might be dozens or hundreds of table definitions, which will quickly become impractical to maintain CREATE and INSERT statements for. If the schemas are changing frequently, it is harder to testing SQL statements up to date.

PostgreSQL License and copyright notice at bottom of article.

Minimum Working Example

Start a Docker container running PostgreSQL:

docker run -e POSTGRES_USER=postgresUser -e POSTGRES_PASSWORD=postgresPW -e POSTGRES_DB=postgresDB -d --rm postgres

The output is a 64-character ID. Start a shell in the new container with it and execute the following commands:

docker exec -it <container_id> bash

Install Python and the necessary dependencies, replacing vim with a terminal editor of your choice:

# Install dependencies.
apt update ; apt install -y python3 pip vim
pip3 install ipdb pytest-postgresql psycopg psycopg2-binary pandas SQLAlchemy
mkdir /app
cd /app

Copy contents into the following files using a text editor of your choice:

1. vim app.py

import pandas as pd
def get_drivers(db_url):
    """Return rows from the driver table."""
    sql = "SELECT * FROM car_schedule.drivers"
    results = pd.read_sql_query(sql=sql, con=db_url)
    if len(results) > 1:
        # Do logic here.
        pass
    return results

2. vim test.sql

CREATE SCHEMA car_schedule;
CREATE TABLE car_schedule.drivers (
  id SERIAL PRIMARY KEY, 
  name VARCHAR(100) NOT NULL, 
  license_number VARCHAR(100) NOT NULL
);
INSERT INTO car_schedule.drivers (name, license_number) 
VALUES 
  ('Jeff', 'abc123'), 
  ('Dan', 'def456');

3.vim test_app.py

import pandas as pd
import pytest
from app import get_drivers
@pytest.fixture()
def database(postgresql):
    """Set up the mock DB with the SQL flat file."""
    with open("test.sql") as f:
        setup_sql = f.read()
    with postgresql.cursor() as cursor:
        cursor.execute(setup_sql)
        postgresql.commit()
    yield postgresql
def test_example_postgres(database):
    drivers = get_drivers(db_url=database)
    assert len(drivers) == 2
    assert set(drivers["name"]) == {"Dan", "Jeff"}

Run test_app.py (as user postgres) using Pytest to see the test results:

su postgres -c pytest test_app.py

You should see output like the following:

root@f45ca72980ea:/app# su postgres -c pytest test_app.py
============ test session starts ============
platform linux -- Python 3.9.2, pytest-7.1.2, pluggy-1.0.0
rootdir: /app
plugins: postgresql-4.1.1
collected 1 item
test_app.py .                                                                                                                                                                                                                                          [100%]
============ 1 passed in 1.87s ============

What Did We Just Do?

get_drivers() contains the application SQL logic to test. It takes in a DB connection, executes hard-coded SQL, does some Python stuff and returns results. To test the inputs and outputs of that function, we:

1. define the database/table schema in test.sql
2. in the test file, spin up the Pytest fixture database
3. test the application function using database as the connection parameter

The get_drivers() function executes its SQL and Python logic without modification and returns a response, querying the mocked database instead of a real network one. The test can be run from any machine that has the three files we wrote and the Python and Postgres dependencies.

An Alternate Approach

When mocking PostgreSQL isn’t a good idea, a simpler workaround is to write a wrapper function for the query results:

def query_db():
    res = pd.read_sql_query(sql="SELECT * FROM drivers", con=DB_URL)
    return res

In Pytest, monkeypatch a hardcoded, expected response from the drivers table. The downside is logic within application SQL is not tested, but that might be worth it if maintaining definitions in test.sql is impractical. It is best to try mocking PostgreSQL schemas and data first, and if it becomes a burden in practice then switch to monkeypatching hardcoded query results.

In conclusion, if you want to unit test a Python function that contains a SQL query, this article shows how that can be done with pytest-postgresql.

PostgreSQL License (to repost image)

Portions Copyright © 1996–2022, The PostgreSQL Global Development Group

Portions Copyright © 1994, The Regents of the University of California

Permission to use, copy, modify, and distribute this software and its documentation for any purpose, without fee, and without a written agreement is hereby granted, provided that the above copyright notice and this paragraph and the following two paragraphs appear in all copies.

IN NO EVENT SHALL THE UNIVERSITY OF CALIFORNIA BE LIABLE TO ANY PARTY FOR DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, INCLUDING LOST PROFITS, ARISING OUT OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION, EVEN IF THE UNIVERSITY OF CALIFORNIA HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

THE UNIVERSITY OF CALIFORNIA SPECIFICALLY DISCLAIMS ANY WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. THE SOFTWARE PROVIDED HEREUNDER IS ON AN “AS IS” BASIS, AND THE UNIVERSITY OF CALIFORNIA HAS NO OBLIGATIONS TO PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.