Pytest and PostgreSQL: Fresh database for every test

In Pytest, everyone's favorite Python testing framework, a fixture is a reusable piece of code that arranges something before the test enters, and cleanup after it quits. For example, a temporary file or folder, setups environment, starting a web server, etc. In this post, we will look at how to create a Pytest fixture which creates a test database (empty or with known state) that gets cleaned up, allowing each test to run on a completely clean database.

The goals

We will create a Pytest fixture using Psycopg 3 to prepare and clean up the test database. Because an empty database is rarely helpful for testing, we will optionally apply Yoyo migrations (at the time of writing website is down, go to archive.org snapshot) to fill it up.

So, requirements for the Pytest fixture named test_db created in this blogpost are:

• drop test database if exists before the test
• create an empty database before the test
• optionally apply migrations or create a test data before the test
• provide a connection to the test database to the test
• drop test database after the test (even in the case of failure)

Any test method which requests it by listing it a test method argument:

def test_create_admin_table(test_db):
    ...

Will receive a regular Psycopg Connection instance connected to the test DB. Test can do whatever it need as with plain Psycopg common usage, e.g.:

def test_create_admin_table(test_db):
    # Open a cursor to perform database operations
    cur = test_db.cursor()

    # Pass data to fill a query placeholders and let Psycopg perform
    # the correct conversion (no SQL injections!)
    cur.execute(
        "INSERT INTO test (num, data) VALUES (%s, %s)",
        (100, "abc'def"))

    # Query the database and obtain data as Python objects.
    cur.execute("SELECT * FROM test")
    cur.fetchone()
    # will return (1, 100, "abc'def")

    # You can use `cur.fetchmany()`, `cur.fetchall()` to return a list
    # of several records, or even iterate on the cursor
    for record in cur:
        print(record)

Project file layout

In classic Python project, the sources lives in src/ and tests in tests/:

├── src
│   └── tuvok
│       ├── __init__.py
│       └── sales
│           └── new_user.py
├── tests
│   ├── conftest.py
│   └── sales
│       └── test_new_user.py
├── requirements.txt
└── yoyo.ini

If you use migrations library like fantastical Yoyo, migration scripts are likely in migrations/:

├── migrations
    ├── 20240816_01_Yn3Ca-sales-user-user-add-last-run-table.py
    ├── ...

Configuration

Our test DB fixture will need a very little configuration:

• connection URL - (without database)
• test database name - will be recreated for every test
• (optionally) migrations folder - migrations scripts to apply for every test

Pytest has a natural place conftest.py for sharing fixtures across multiple files. The fixture configuration will go there too:

# Without DB name!
TEST_DB_URL = "postgresql://localhost"
TEST_DB_NAME = "test_tuvok"
TEST_DB_MIGRATIONS_DIR = str(Path(__file__, "../../migrations").resolve())

You can set these values from the environment variable or whatever suits your case.

Create test_db fixture

With knowledge of PostgreSQL and Psycopg library, write the fixture in conftest.py:

@pytest.fixture
def test_db():
    # autocommit=True start no transaction because CREATE/DROP DATABASE
    # cannot be executed in a transaction block.
    with psycopg.connect(TEST_DB_URL, autocommit=True) as conn:
        cur = conn.cursor()

        # create test DB, drop before
        cur.execute(f'DROP DATABASE IF EXISTS "{TEST_DB_NAME}" WITH (FORCE)')
        cur.execute(f'CREATE DATABASE "{TEST_DB_NAME}"')

        # Return (a new) connection to just created test DB
        # Unfortunately, you cannot directly change the database for an existing Psycopg connection. Once a connection is established to a specific database, it's tied to that database.
        with psycopg.connect(TEST_DB_URL, dbname=TEST_DB_NAME) as conn:
            yield conn

        cur.execute(f'DROP DATABASE IF EXISTS "{TEST_DB_NAME}" WITH (FORCE)')

Create migrations fixture

In our case, we use Yoyo migrations. Write apply migrations as another fixture called yoyo:

@pytest.fixture
def yoyo():
    # Yoyo expect `driver://user:pass@host:port/database_name?param=value`.
    # In passed URL we need to
    url = (
        urlparse(TEST_DB_URL)
        .
        # 1) Change driver (schema part) with `postgresql+psycopg` to use
        # psycopg 3 (not 2 which is `postgresql+psycopg2`)
        _replace(scheme="postgresql+psycopg")
        .
        # 2) Change database to test db (in which migrations will apply)
        _replace(path=TEST_DB_NAME)
        .geturl()
    )

    backend = get_backend(url)
    migrations = read_migrations(TEST_DB_MIGRATIONS_DIR)

    if len(migrations) == 0:
        raise ValueError(f"No Yoyo migrations found in '{TEST_DB_MIGRATIONS_DIR}'")

    with backend.lock():
        backend.apply_migrations(backend.to_apply(migrations))

If you want to apply migrations to every test database, require yoyo fixture for test_db fixture:

@pytest.fixture
def test_db(yoyo):
    ...

To apply migration to some test only, require yoyo individually:

def test_create_admin_table(test_db, yoyo):
    ...

Conclusion

Building own fixture to give your tests a clean database was a rewarding experience for me allowing me to delve deeper into both Pytest and Postgres.

I hope this article helped you with your own database test suite. Feel free to leave me the your question in comments and happy coding!

The above is the detailed content of Pytest and PostgreSQL: Fresh database for every test. For more information, please follow other related articles on the PHP Chinese website!