3

I am pretty new to testing in general. I am using postgreSQL and SQLAlchemy for databases, FastAPI backend and Alembic for migrations. I want to set up a test database for testing. I have custom migration scripts that I want to execute on my test database. Here is a sample test_users file.

import pytest
import alembic
from alembic.config import Config

from fastapi.testclient import TestClient
import sqlalchemy as sa
from sqlalchemy import create_engine
from sqlalchemy_utils import create_database, database_exists, drop_database
from sqlalchemy.orm import sessionmaker
`
TEST_DB_URL = f"{DATABASE_URL}_test"

if database_exists(TEST_DB_URL):
    drop_database(TEST_DB_URL)

create_database(TEST_DB_URL)

alembic_cfg = Config("alembic.ini")
alembic_cfg.set_main_option('sqlalchemy.url', str(TEST_DB_URL))
alembic.command.upgrade(alembic_cfg, "head")

engine = create_engine(TEST_DB_URL)
TestingSessionLocal = sessionmaker()


@pytest.fixture(scope='session')
def connection():
    connection = engine.connect()
    yield connection
    connection.close()

@pytest.fixture
def session(connection):
    transaction = connection.begin()
    session = TestingSessionLocal(bind=connection)
    yield session
    session.close()
    transaction.rollback()
    # drop_database(TEST_DB_URL)

@pytest.fixture()
def client(session):
    def override_get_db():
        yield session

    app.dependency_overrides[get_db] = override_get_db
    client = TestClient(app)
    yield client
    del app.dependency_overrides[get_db]

def test_user_register(client):
    response = client.post("/register", json={"email": "def@def.com", "master_pwd": "pass"})
    print(response)
    assert response.status_code == 200

But as I am checking pgadmin, the test database is created but the migrations are not being applied, and running pytest gives 404 error. What am I doing wrong?

Shiladitya Bose
  • 893
  • 2
  • 13
  • 31

2 Answers2

1

I know it's an older thread, but I faced the same issue and this popped up on my Google search.

In my case, the problem is that in my env.py I have the following line:

config.set_main_option("sqlalchemy.url", settings.DATABASE_URL)

Where settings.DATABASE_URL is a variable that contains the db connection read from an environment variable (to avoid hard-coding it in the alembic.ini file).

Now when you call alembic.command.upgrade(alembic_cfg, "head") the env.py is implicitly run as well and overrides your alembic_cfg.set_main_option('sqlalchemy.url', str(TEST_DB_URL)) back to the real DATABASE_URL.

I solved this by mocking my settings module via

from core import settings as mocked_settings
mocked_settings.DATABASE_URL = TEST_DB_URL
mocker.patch.dict("sys.modules", {"core.settings": mocked_settings})

where mocker refers to the fixture from the pytest-mock module (you can use the standard library mock as well though).

In case you read the database URL directly from an environment variable in your env.py it might be even easier:

mocker.patch.dict("os.environ", DATABASE_URL=TEST_DB_URL)

Hope that helps others! :)

SaturnFromTitan
  • 1,334
  • 14
  • 20
0

An alternative to the above approach is to setup the alembic_engine in your tests/conftest.py file:

@pytest.fixture(scope="session")
def alembic_engine():
    """Override this fixture to provide pytest-alembic powered tests with a database handle.
    """
    test_db = set_up_empty_test_db()   # Whatever this means in your code..
    try:
        sqlalchemy_url = f"..."  # whatever your test cx string would be
        engine = sqlalchemy.create_engine(sqlalchemy_url)
    except Exception as e:
        test_db.stop()  # whatever your cleanup logic requires
        raise e

    yield engine

    test_db.stop()  # whatever your cleanup logic requires

And then, in your env.py check for an existing connection and use it if present (pytest-alembic will set it at startup):

def run_migrations_online():
    # When running pytest-alembic the connection will be pre-populated
    connectable = context.config.attributes.get("connection", None)

    if connectable is None:
        connectable = engine_from_config(...  # regular alembic

    # The with statement on down is unchanged
    with connectable.connect() as connection:
       ...

See the setup docs for further details: https://pytest-alembic.readthedocs.io/en/latest/setup.html