Jan Švejda
Jan Švejda

Reputation: 160

pytest PostgreSQL - different fixture for two environments to set up test DB either locally or on CI

I am trying to have pytest setup a test database automatically based on whether it runs locally or on our CI server. Please see below the code. For this I created a CLI option.

When this option is passed I want pytest-postgresql to use the postgresql_noproc_factory (GitHub README). So I created a fixture that selects the factory to be used. However, if I understand correctly, this creates a "fixture of a fixture" so when it is injected into setup_database it doesn't work and gives an error message:

    def dbcreator():
>       return postgresql_my.cursor().connection
E       AttributeError: 'function' object has no attribute 'cursor'

Do you possibly know how this could be done with pytest? How do I make pytest instantiate the nested fixture so that it configures the right PostgreSQL engine?

import os
import tempfile

from psycopg2.extensions import connection
import pytest
from pytest_postgresql import factories
from sqlalchemy import create_engine
from sqlalchemy.engine import Engine


def pytest_addoption(parser):
    parser.addoption(
        "--remote-psql", action="store_true", help="Run it in environment where PSQL runs on network."
    )


@pytest.fixture(scope="session")
def cmdopt_remote_psql(request):
    return request.config.getoption("--remote-psql")


# Using the postgresql factories to create a postgresql fixture instance
socket_dir = tempfile.TemporaryDirectory()
postgresql_proc_factory = factories.postgresql_proc(port=None, unixsocketdir=socket_dir.name)
postgresql_noproc_factory = factories.postgresql_noproc()


@pytest.fixture(scope="session")
def postgresql_my(cmdopt_remote_psql):
    if cmdopt_remote_psql:
        return factories.postgresql("postgresql_noproc_factory")
    else:
        return factories.postgresql("postgresql_proc_factory")


@pytest.fixture(scope="function", autouse=True)
def setup_database(postgresql_my: connection) -> Engine:
    """
    See https://medium.com/@geoffreykoh/fun-with-fixtures-for-database-applications-8253eaf1a6d
    """
    def dbcreator():
        return postgresql_my.cursor().connection

    engine = create_engine("postgresql+psycopg2://", creator=dbcreator)
    return engine

Upvotes: 1

Views: 4715

Answers (1)

Jan Švejda
Jan Švejda

Reputation: 160

So in the end I made it work by customizing a the factories.postgresql function from pytest-postgresql to accept the cmdopt_remote_psql fixture as argument. Here it is:


def postgresql(db_name: str = None, load: List[str] = None) -> Callable[[FixtureRequest, bool], connection]:
    """
    Copied from pytest-postgresql. Needed to customize the returned fixture
    based on the CLI option cmdopt_remote_psql.
    """

    @pytest.fixture(scope="function")
    def postgresql_factory(request: FixtureRequest, cmdopt_remote_psql: bool) -> connection:
        """
        Fixture factory for PostgreSQL.

        :param FixtureRequest request: fixture request object
        :param cmdopt_remote_psql: fixture for CLI arg to pick correct factory
        :returns: postgresql client
        """
        if not psycopg2:
            raise ImportError(
                "No module named psycopg2. Please install either "
                "psycopg2 or psycopg2-binary package for CPython "
                "or psycopg2cffi for Pypy."
            )
        if cmdopt_remote_psql is True:
            process_fixture_name = "postgresql_noproc_factory"
        else:
            process_fixture_name = "postgresql_proc_factory"
        config = get_config(request)
        proc_fixture: Union[PostgreSQLExecutor, NoopExecutor] = request.getfixturevalue(process_fixture_name)

        pg_host = proc_fixture.host
        pg_port = proc_fixture.port
        pg_user = proc_fixture.user
        pg_password = proc_fixture.password
        pg_options = proc_fixture.options
        pg_db = db_name or config["dbname"]
        pg_load = load or config["load"]

        with DatabaseJanitor(pg_user, pg_host, pg_port, pg_db, proc_fixture.version, pg_password):
            db_connection: connection = psycopg2.connect(
                dbname=pg_db,
                user=pg_user,
                password=pg_password,
                host=pg_host,
                port=pg_port,
                options=pg_options,
            )
            if pg_load:
                for filename in pg_load:
                    with open(filename, "r") as _fd:
                        with db_connection.cursor() as cur:
                            cur.execute(_fd.read())
                            db_connection.commit() # I added this as well since the tables weren't available otherwise.
            yield db_connection
            db_connection.close()

    return postgresql_factory

Then I replaced the fixture postgresql_my with this call:

postgresql_my = postgresql()

and it works!

Upvotes: 2

Related Questions