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)
I have tried pytest-postgresql which promises the same. I have tried it before writing my own fixture but I was not able to make it work for me. Maybe because their docs were very confusing to me. Another, pytest-dbt-postgres, I haven't tried at all.Motivation & alternatives
It looks like there are some Pytest plugins which promise PostgreSQL fixtures for tests that rely on databases. They might work well for you.
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!
Top comments (0)