Model/tests/conftest.py
Daniel Roth 82156fac8f Audit generator populates XLSX, uploads to S3, and records UploadedFile row 🟪
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-09 12:05:55 +00:00

84 lines
3 KiB
Python

"""Shared pytest fixtures for the ``tests/`` tree.
Provides an ephemeral PostgreSQL engine for tests that exercise SQLModel
repositories. PostgreSQL has no true in-memory mode; ``pytest-postgresql``
starts a real, throwaway server in a temp directory (the process is started
once per session and a fresh database is created/dropped per test). That is
the closest equivalent to "in-memory" and matches production behaviour far
better than SQLite (enums, JSONB, constraint semantics, etc.).
"""
from __future__ import annotations
import glob
from collections.abc import Iterator
from typing import Any
import pytest
from psycopg import Connection
from pytest_postgresql import factories
from sqlalchemy import Engine, text
from sqlmodel import SQLModel, create_engine
# Importing the SQLModel row modules registers their tables on
# SQLModel.metadata so ``create_all`` builds the full schema. Imports look
# unused; they aren't.
import infrastructure.postgres.uploaded_file_table as _uf_table # pyright: ignore[reportUnusedImport]
# pg_ctl ships under a versioned path and is not on PATH in the dev container.
_PG_CTL = next(iter(sorted(glob.glob("/usr/lib/postgresql/*/bin/pg_ctl"))), "pg_ctl")
postgresql_proc = factories.postgresql_proc(executable=_PG_CTL) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
postgresql = factories.postgresql("postgresql_proc")
def _create_pg_enum_types(engine: Engine) -> None:
"""Emit CREATE TYPE for PostgreSQL enum types used by UploadedFile.
SQLModel.metadata.create_all uses create_type=False for these enums
(they are normally created by Alembic migrations). Tests need them upfront.
A DO block swallows duplicate_object so the fixture is safe to call on a
pre-seeded database.
"""
from infrastructure.postgres.uploaded_file_table import FileSourceEnum, FileTypeEnum
ft_values = ", ".join(f"'{e.value}'" for e in FileTypeEnum)
fs_values = ", ".join(f"'{e.value}'" for e in FileSourceEnum)
with engine.connect() as conn:
conn.execute(
text(
f"""
DO $$ BEGIN
CREATE TYPE file_type AS ENUM ({ft_values});
EXCEPTION WHEN duplicate_object THEN NULL;
END $$;
"""
)
)
conn.execute(
text(
f"""
DO $$ BEGIN
CREATE TYPE file_source AS ENUM ({fs_values});
EXCEPTION WHEN duplicate_object THEN NULL;
END $$;
"""
)
)
conn.commit()
@pytest.fixture
def db_engine(postgresql: Connection[Any]) -> Iterator[Engine]:
"""A SQLModel engine bound to a fresh, ephemeral PostgreSQL database."""
info = postgresql.info
url = f"postgresql+psycopg://{info.user}:@{info.host}:{info.port}/{info.dbname}"
engine = create_engine(url)
_create_pg_enum_types(engine)
SQLModel.metadata.create_all(engine)
try:
yield engine
finally:
SQLModel.metadata.drop_all(engine)
engine.dispose()