Pytest and PostgreSQL: Fresh database for every test (part II)

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

• yield is used to pause a generator function's execution and return a single value to the caller.
• while yield from is 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 (not yield) 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!