In the previous post, we created the Pytest fixture which will create/drop Postgres database before/after the test method. In this part, I want to improve the fixture to be more flexible and configurable with the help of Pytest factory fixtures.
Limits of static fixture
For example, if you have more than one database to mock in the test
def test_create_user(test_db1, test_db2):
...
you must create almost two identical fixtures:
TEST_DB_URL = "postgresql://localhost"
TEST_DB1_NAME = "test_foo"
TEST_DB2_NAME = "test_bar"
@pytest.fixture
def test_db1():
with psycopg.connect(TEST_DB_URL, autocommit=True) as conn:
cur = conn.cursor()
cur.execute(f'DROP DATABASE IF EXISTS "{TEST_DB1_NAME}" WITH (FORCE)')
cur.execute(f'CREATE DATABASE "{TEST_DB1_NAME}"')
with psycopg.connect(TEST_DB_URL, dbname=TEST_DB1_NAME) as conn:
yield conn
cur.execute(f'DROP DATABASE IF EXISTS "{TEST_DB1_NAME}" WITH (FORCE)')
@pytest.fixture
def test_db2():
with psycopg.connect(TEST_DB_URL, autocommit=True) as conn:
cur = conn.cursor()
cur.execute(f'DROP DATABASE IF EXISTS "{TEST_DB2_NAME}" WITH (FORCE)')
cur.execute(f'CREATE DATABASE "{TEST_DB2_NAME}"')
with psycopg.connect(TEST_DB_URL, dbname=TEST_DB2_NAME) as conn:
yield conn
cur.execute(f'DROP DATABASE IF EXISTS "{TEST_DB2_NAME}" WITH (FORCE)')
Pytest fixture factories
"Static" fixtures are a bit limiting here. When needed almost the same with just a slight difference, you need to duplicate a code. Hopefully, the Pytest has a concept of factories as fixtures.
Factory fixture is a fixture which returns another fixture. Because, like every factory, it is a function, it can accept arguments to customize returned fixtures. By convention, you can prefix them with make_*, like make_test_db.
Specialized fixtures
The only argument to our fixture factory make_test_db will be a test database name to create/drop.
So, let's create two "specialized" fixtures based on the make_test_db factory fixture.
The usage will looks like:
@pytest.fixture
def test_db_foo(make_test_db):
yield from make_test_db("test_foo")
@pytest.fixture
def test_db_bar(make_test_db):
yield from make_test_db("test_bar")
Sidenote: yield from
Did you notice the yield from? There is a key difference between yield and yield from on how they handle the flow of data and control within generators.
In Python, both yield and yield from are used within generator functions to produce a sequence of values, but
-
yieldis used to pause a generator function's execution and return a single value to the caller. - while
yield fromis used to delegate the generation of values to another generator. It essentially "flattens" the nested generator, passing its yielded values directly to the caller of the outer generator.
I.e., we don't want to "yield" from a specialized fixture but from a fixture factory. Therefore yield from is required here.
Fixture factory to create/drop database
Changes required to our original fixture creating/dropping database are actually almost none except the wrapping the code to the inner function.
@pytest.fixture
def make_test_db():
def _(test_db_name: str):
with psycopg.connect(TEST_DB_URL, autocommit=True) as conn:
cur = conn.cursor()
cur.execute(f'DROP DATABASE IF EXISTS "{test_db_name}" WITH (FORCE)') # type: ignore
cur.execute(f'CREATE DATABASE "{test_db_name}"') # type: ignore
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)') # type: ignore
yield _
Bonus: Rewrite migrations fixture as factory fixture
In the previous part, I also had a fixture applying Yoyo migrations to just created empty database. It was also not very flexible. Let's do the same and wrap actual code to the inner function.
In this case, because the code doesn't need to do cleanup after return from test method (no yield in it), the
- factory fixture
returns (notyield) inner function - specialized fixture calls (not
yield from) factory fixture
@pytest.fixture
def make_yoyo():
"""Applies Yoyo migrations to test DB."""
def _(test_db_name: str, migrations_dir: str):
url = (
urlparse(TEST_DB_URL)
.
_replace(scheme="postgresql+psycopg")
.
_replace(path=test_db_name)
.geturl()
)
backend = get_backend(url)
migrations = read_migrations(migrations_dir)
if len(migrations) == 0:
raise ValueError(f"No Yoyo migrations found in '{migrations_dir}'")
with backend.lock():
backend.apply_migrations(backend.to_apply(migrations))
return _
@pytest.fixture
def yoyo_foo(make_yoyo):
migrations_dir = str(Path(__file__, "../../foo/migrations").resolve())
make_yoyo("test_foo", migrations_dir)
@pytest.fixture
def yoyo_bar(make_yoyo):
migrations_dir = str(Path(__file__, "../../bar/migrations").resolve())
make_yoyo("test_bar", migrations_dir)
A test method which need two databases and apply migrations to them:
from psycopg import Connection
def test_get_new_users_since_last_run(
test_db_foo: Connection,
test_db_bar: Connection,
yoyo_foo,
yoyo_bar):
test_db_foo.execute("...")
...
Conclusion
Building your own fixture factory creating and dropping databases for the Pytest method is actually a good exercise to practice Python generator and yield/yield from operators.
I hope this article helped you with your own database test suite. Feel free to leave me your question in the comments and happy coding!
Top comments (0)