FabledCurator/tests/conftest.py

"""Shared pytest fixtures.

The async db fixture provides an AsyncSession bound to a transaction that
gets rolled back after each test. CI provisions a real Postgres + pgvector
(see .forgejo/workflows/ci.yml), so tests exercise the actual schema and
migration code paths.
"""

import os

import pytest
import pytest_asyncio
from sqlalchemy import create_engine
from sqlalchemy.ext.asyncio import (
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy.orm import sessionmaker

from backend.app import create_app
from backend.app.models import Base


def _async_database_url() -> str:
    user = os.environ.get("DB_USER", "fabledcurator")
    password = os.environ["DB_PASSWORD"]
    host = os.environ.get("DB_HOST", "postgres")
    port = os.environ.get("DB_PORT", "5432")
    name = os.environ.get("DB_NAME", "fabledcurator_test")
    return f"postgresql+asyncpg://{user}:{password}@{host}:{port}/{name}"


def _sync_database_url() -> str:
    user = os.environ.get("DB_USER", "fabledcurator")
    password = os.environ["DB_PASSWORD"]
    host = os.environ.get("DB_HOST", "postgres")
    port = os.environ.get("DB_PORT", "5432")
    name = os.environ.get("DB_NAME", "fabledcurator_test")
    return f"postgresql+psycopg://{user}:{password}@{host}:{port}/{name}"


@pytest_asyncio.fixture
async def engine():
    e = create_async_engine(_async_database_url(), future=True)
    yield e
    await e.dispose()


@pytest_asyncio.fixture
async def db(engine) -> AsyncSession:
    Session = async_sessionmaker(engine, expire_on_commit=False)
    async with Session() as session:
        await session.begin()
        try:
            yield session
        finally:
            await session.rollback()


@pytest.fixture
def sync_engine():
    e = create_engine(_sync_database_url(), future=True)
    yield e
    e.dispose()


@pytest.fixture
def db_sync(sync_engine):
    """Synchronous Session bound to a savepoint — used by Importer tests
    (the Importer is sync-only by design)."""
    SyncSession = sessionmaker(sync_engine, expire_on_commit=False)
    with SyncSession() as session:
        session.begin()
        try:
            yield session
        finally:
            session.rollback()


@pytest_asyncio.fixture
async def app():
    return create_app()


@pytest_asyncio.fixture
async def client(app):
    async with app.test_client() as c:
        yield c


# Migration-seeded baseline (singleton config like import_settings /
# ml_settings), captured lazily at the FIRST integration test's setup —
# when the CI integration job has just run `alembic upgrade head` so the DB
# is pristine. Cached module-wide. Never populated in the DB-less unit job
# because only integration-marked tests trigger the connection.
_SEED_SNAPSHOT: dict[str, list[dict]] | None = None


@pytest.fixture(autouse=True)
def _reset_db_after_integration(request):
    """Integration tests exercise app/celery code that commits on its own
    connection, which the rollback fixtures above can't undo — so data
    would leak between tests. Capture the seeded baseline at the first
    integration test, then after each integration-marked test truncate
    every model table and restore that baseline. Connects to the DB ONLY
    for integration-marked tests, so the no-DB fast unit job is untouched.
    """
    global _SEED_SNAPSHOT
    is_integration = request.node.get_closest_marker("integration") is not None

    if is_integration and _SEED_SNAPSHOT is None:
        eng = create_engine(_sync_database_url(), future=True)
        snap: dict[str, list[dict]] = {}
        try:
            with eng.connect() as conn:
                for t in Base.metadata.sorted_tables:
                    rows = [
                        dict(m)
                        for m in conn.execute(t.select()).mappings().all()
                    ]
                    if rows:
                        snap[t.name] = rows
        finally:
            eng.dispose()
        _SEED_SNAPSHOT = snap

    yield

    if not is_integration:
        return
    tables = ", ".join(t.name for t in Base.metadata.sorted_tables)
    if not tables:
        return
    eng = create_engine(_sync_database_url(), future=True)
    try:
        with eng.begin() as conn:
            conn.exec_driver_sql(
                f"TRUNCATE {tables} RESTART IDENTITY CASCADE"
            )
            for t in Base.metadata.sorted_tables:
                rows = (_SEED_SNAPSHOT or {}).get(t.name)
                if rows:
                    conn.execute(t.insert(), rows)
    finally:
        eng.dispose()


@pytest_asyncio.fixture(autouse=True)
async def _dispose_app_engine(request):
    """The app's shared async engine (extensions.get_engine) is a process
    singleton, but pytest-asyncio gives each test a fresh event loop;
    asyncpg connections are loop-bound, so reuse across tests raises
    'attached to a different loop'. Dispose after each integration test so
    the next test rebuilds the engine on its own loop.
    """
    yield
    if request.node.get_closest_marker("integration") is None:
        return
    from backend.app import extensions
    await extensions.dispose_engine()