43 lines
1.7 KiB
Python
43 lines
1.7 KiB
Python
"""Exercise the model configuration Alembic migration against a real SQLite file."""
|
|
|
|
from pathlib import Path
|
|
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
from sqlalchemy import create_engine, inspect, text
|
|
|
|
|
|
def test_model_config_migration_upgrades_existing_database(tmp_path: Path, monkeypatch):
|
|
from agenteval.storage import db as db_module
|
|
|
|
database_url = f"sqlite:///{tmp_path / 'migration.db'}"
|
|
monkeypatch.setattr(db_module, "DATABASE_URL", database_url)
|
|
config = Config(str(Path(__file__).resolve().parents[2] / "alembic.ini"))
|
|
|
|
legacy_engine = create_engine(database_url)
|
|
with legacy_engine.begin() as connection:
|
|
connection.execute(text("CREATE TABLE scenarios (id VARCHAR PRIMARY KEY, name VARCHAR NOT NULL)"))
|
|
connection.execute(text("CREATE TABLE eval_results (id VARCHAR PRIMARY KEY, run_id VARCHAR NOT NULL)"))
|
|
connection.execute(
|
|
text(
|
|
"CREATE TABLE eval_runs (id VARCHAR PRIMARY KEY, target_id VARCHAR NOT NULL, "
|
|
"scenario_id VARCHAR NOT NULL)"
|
|
)
|
|
)
|
|
connection.execute(text("CREATE TABLE turns (id VARCHAR PRIMARY KEY, run_id VARCHAR NOT NULL)"))
|
|
|
|
command.upgrade(config, "df320ecde84f")
|
|
before = inspect(create_engine(database_url)).get_table_names()
|
|
assert "model_configs" not in before
|
|
|
|
command.upgrade(config, "head")
|
|
inspector = inspect(create_engine(database_url))
|
|
assert "model_configs" in inspector.get_table_names()
|
|
assert "scenario_model_bindings" in inspector.get_table_names()
|
|
assert {column["name"] for column in inspector.get_columns("model_configs")} >= {
|
|
"name",
|
|
"capability",
|
|
"endpoint_url",
|
|
"api_key_encrypted",
|
|
}
|