107 lines
3.1 KiB
Python
107 lines
3.1 KiB
Python
"""Tests for Alembic migration setup."""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
from sqlalchemy import create_engine, inspect
|
|
|
|
# Resolve the repo root regardless of where pytest is invoked from.
|
|
_REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
|
|
|
|
@pytest.fixture()
|
|
def alembic_cfg(tmp_path):
|
|
"""Create an Alembic config pointing at a temporary SQLite database."""
|
|
db_path = tmp_path / "test.db"
|
|
db_url = f"sqlite:///{db_path}"
|
|
|
|
cfg = Config(str(_REPO_ROOT / "alembic.ini"))
|
|
cfg.set_main_option("script_location", str(_REPO_ROOT / "alembic"))
|
|
cfg.set_main_option("sqlalchemy.url", db_url)
|
|
return cfg, db_url
|
|
|
|
|
|
def test_upgrade_head_creates_all_tables(alembic_cfg):
|
|
"""Running 'upgrade head' should create all expected tables."""
|
|
cfg, db_url = alembic_cfg
|
|
command.upgrade(cfg, "head")
|
|
|
|
engine = create_engine(db_url)
|
|
inspector = inspect(engine)
|
|
tables = set(inspector.get_table_names())
|
|
|
|
expected = {
|
|
"alembic_version",
|
|
"users",
|
|
"projects",
|
|
"experiments",
|
|
"runs",
|
|
"stage_results",
|
|
"scores",
|
|
"response_cache",
|
|
"webhook_configs",
|
|
}
|
|
assert expected == tables
|
|
|
|
|
|
def test_downgrade_base_removes_all_tables(alembic_cfg):
|
|
"""Running 'downgrade base' should remove all application tables."""
|
|
cfg, db_url = alembic_cfg
|
|
command.upgrade(cfg, "head")
|
|
command.downgrade(cfg, "base")
|
|
|
|
engine = create_engine(db_url)
|
|
inspector = inspect(engine)
|
|
tables = set(inspector.get_table_names())
|
|
|
|
# Only alembic_version should remain
|
|
assert tables == {"alembic_version"}
|
|
|
|
|
|
def test_runs_table_has_expected_columns(alembic_cfg):
|
|
"""Spot-check that the runs table has key columns."""
|
|
cfg, db_url = alembic_cfg
|
|
command.upgrade(cfg, "head")
|
|
|
|
engine = create_engine(db_url)
|
|
inspector = inspect(engine)
|
|
columns = {c["name"] for c in inspector.get_columns("runs")}
|
|
|
|
assert "id" in columns
|
|
assert "experiment_id" in columns
|
|
assert "config_hash" in columns
|
|
assert "status" in columns
|
|
assert "cost_estimate" in columns
|
|
|
|
|
|
def test_indexes_created(alembic_cfg):
|
|
"""Verify key indexes exist after migration."""
|
|
cfg, db_url = alembic_cfg
|
|
command.upgrade(cfg, "head")
|
|
|
|
engine = create_engine(db_url)
|
|
inspector = inspect(engine)
|
|
|
|
run_indexes = {idx["name"] for idx in inspector.get_indexes("runs")}
|
|
assert "ix_runs_config_hash" in run_indexes
|
|
assert "ix_runs_experiment_id" in run_indexes
|
|
|
|
score_indexes = {idx["name"] for idx in inspector.get_indexes("scores")}
|
|
assert "ix_scores_run_id" in score_indexes
|
|
assert "ix_scores_scorer_name" in score_indexes
|
|
|
|
|
|
def test_foreign_keys_on_experiments(alembic_cfg):
|
|
"""Verify experiments table has FK to projects."""
|
|
cfg, db_url = alembic_cfg
|
|
command.upgrade(cfg, "head")
|
|
|
|
engine = create_engine(db_url)
|
|
inspector = inspect(engine)
|
|
fks = inspector.get_foreign_keys("experiments")
|
|
|
|
referred_tables = {fk["referred_table"] for fk in fks}
|
|
assert "projects" in referred_tables
|