I believe that a hacker/script-kiddie evolves into a software engineer at the point when they realize that they need to write tests for their own benefit and not because someone else forced them to complete 100% test coverage in the code. That realization took me four years down the line after writing my first professional code, but today, anything more than a pure exploratory script that I write will have some sort of automated testing associated with it.
Today I want to take a break from all the AI craziness and spend a little bit of time explaining a small technique to improve the way you write tests in Python.
Using Data in Tests
Any reasonably sized product and system will use a database to store/access the data. At a really small scale, it can even be something as simple as a SQLite/Files/Memory, but the moment you hit production, at any scale, you will want to use a real database. The standard choice today for 90% of cases is PostgreSQL, because the industry seems to have settled on it.
Using PostgreSQL as the DB, writing a test goes something like this.
- Add data into the DB
- Verify that the functionality works
For example, let’s say we have a User model in our application and want to test user permissions.
def test_user_permissions():
# Create a user and role directly in the database (setup)
conn = get_db_connection()
cursor = conn.cursor()
cursor.execute(
"INSERT INTO users (username, email) VALUES (%s, %s) RETURNING id",
("testuser", "[email protected]")
)
user_id = cursor.fetchone()[0]
cursor.execute("INSERT INTO roles (name) VALUES (%s) RETURNING id", ("admin",))
role_id = cursor.fetchone()[0]
cursor.execute(
"INSERT INTO user_roles (user_id, role_id) VALUES (%s, %s)",
(user_id, role_id)
)
conn.commit()
# Test permission functionality
has_admin = check_user_permission(user_id, "admin_access")
assert has_admin == True
The above setup works fine and if there are some issues with the application logic in check_user_permission, the test should find it correctly. However, the code above, as is, can cause insidious bugs.
Test isolation with Setup and Teardown
The reason is that once there are multiple tests, they could secretly affect each other. If you create data in test A, it could still be hanging around when test B runs. This causes flaky tests that pass sometimes and fail other times depending on the order they run in. There is already lots of material available online on test isolation, so I am not going to go into that here.
The traditional solution around the problem is setup and teardown. Before each test, you set up the data you need. After each test, you clean up everything you created.
def test_user_permissions():
conn = get_db_connection()
cursor = conn.cursor()
# Setup - create test data
cursor.execute(
"INSERT INTO users (username, email) VALUES (%s, %s) RETURNING id",
("testuser", "[email protected]")
)
user_id = cursor.fetchone()[0]
cursor.execute("INSERT INTO roles (name) VALUES (%s) RETURNING id", ("admin",))
role_id = cursor.fetchone()[0]
cursor.execute(
"INSERT INTO user_roles (user_id, role_id) VALUES (%s, %s)",
(user_id, role_id)
)
conn.commit()
# Test - actual functionality we care about
has_admin = check_user_permission(user_id, "admin_access")
has_super = check_user_permission(user_id, "super_admin_access")
assert has_admin == True
assert has_super == False
# Teardown - clean up everything we created
cursor.execute("DELETE FROM user_roles WHERE user_id = %s", (user_id,))
cursor.execute("DELETE FROM users WHERE id = %s", (user_id,))
cursor.execute("DELETE FROM roles WHERE id = %s", (role_id,))
conn.commit()
This seems to have fixed the problem to an extent, but we have more issues now.
- We have a lot of boilerplate now which we don’t care about.
- Tests stopping in the middle for any reason (error in the code or even a power failure) can cause the data state to reach an invalid state which can cause flakiness when running the test again.
MVCC
PostgreSQL has this brilliant feature called Multi-Version Concurrency Control (MVCC). Without getting too deep into the database internals, it can be summarized as: PostgreSQL can handle multiple versions of the same data simultaneously.
When you start a transaction, PostgreSQL creates a snapshot of the database at that moment. Any changes you make within that transaction are only visible to you on that connection. Other connections see the original data. If you rollback the transaction, it’s like your changes never happened - PostgreSQL just discards your version and keeps the original. What is beautiful is that even if you don’t manually rollback and the transaction gets cancelled for any reason, it is happening within a “sandbox” and other queries are not affected by these changes.
This is what makes PostgreSQL so good at handling concurrent queries without locks everywhere. But it also gives us a superpower for testing.
MVCC + Tests
The idea is simple: what if we run each test inside its own transaction? Any data we create during the test automatically becomes invisible to other tests. We don’t even need to worry about teardown!
# conftest.py
import pytest
import psycopg2
from typing import Generator, Any
@pytest.fixture(scope="function")
def db_connection() -> Generator[Any, Any, None]:
# Create a new connection for each test
conn = psycopg2.connect("your_test_db_connection_string")
conn.autocommit = False # Start transaction
try:
yield conn
finally:
conn.rollback()
conn.close()
# test_something.py
def test_user_permissions(db_connection):
cursor = db_connection.cursor()
# Create test data - same as before but now within transaction
cursor.execute(
"INSERT INTO users (username, email) VALUES (%s, %s) RETURNING id",
("testuser", "[email protected]")
)
user_id = cursor.fetchone()[0]
cursor.execute("INSERT INTO roles (name) VALUES (%s) RETURNING id", ("admin",))
role_id = cursor.fetchone()[0]
cursor.execute(
"INSERT INTO user_roles (user_id, role_id) VALUES (%s, %s)",
(user_id, role_id)
)
# Test permission functionality - same as before
has_admin = check_user_permission(user_id, "admin_access")
has_super = check_user_permission(user_id, "super_admin_access")
assert has_admin == True
assert has_super == False
# No teardown needed! MVCC handles isolation automatically
This seems like a very small change but the benefits can add up quickly since we are no longer worried about test isolation.
The connection object is best defined at the main conftest.py (with function/class scope) and be shared by all of the tests.
This pattern has saved me countless hours of writing teardown code and hunting down flaky tests.
One thing to note is that this works great for database logic, but if your tests interact with external services or run E2E tests or files, you’ll still need traditional teardown for those. For the vast majority of backend tests this shouldn’t be a problem since you’re mostly testing application logic and using the DB to store/retrieve data.
Give it a try in your next Python project and let me know how it goes!
You can use the same MVCC concepts within an API to keep it idempotent as well as handle failures cleanly! I’ll explain this in another post.
If you found this useful, please consider sharing it on Hacker News!