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 arrangessomethingbefore 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 athow 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 databaseif exists before the test
• create an empty databasebefore the test
• optionallyapply migrations or create a test databefore the test
• provide a connection to the test databaseto the test
• drop test databaseafter 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 ofPostgreSQL 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 useYoyo 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 toapply migrations to every test database, require yoyo fixture for test_db fixture:

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

Toapply 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!