140 lines
4.4 KiB
Python
140 lines
4.4 KiB
Python
"""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.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.fixture(scope="session")
|
|
def _seed_snapshot():
|
|
"""Capture migration-seeded rows (singleton config like import_settings
|
|
/ ml_settings) once, so the per-test truncation can restore them."""
|
|
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()
|
|
return snap
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_db_after_integration(request, _seed_snapshot):
|
|
"""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. After each integration-marked test, truncate
|
|
every model table, then restore the migration-seeded rows. Scoped to
|
|
the integration marker so the DB-less fast unit job never connects.
|
|
"""
|
|
yield
|
|
if request.node.get_closest_marker("integration") is None:
|
|
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.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()
|