90 lines
4.0 KiB
Python
90 lines
4.0 KiB
Python
from pathlib import Path
|
|
from tempfile import TemporaryDirectory
|
|
|
|
from alembic.config import Config
|
|
from sqlalchemy import create_engine, inspect
|
|
|
|
from alembic import command
|
|
from modelforge_api.persistence.models import Base
|
|
|
|
|
|
def test_operator_benchmark_migration_upgrades_and_downgrades_sqlite() -> None:
|
|
with TemporaryDirectory() as temp:
|
|
database_path = Path(temp) / "benchmark-migration.db"
|
|
url = f"sqlite:///{database_path.as_posix()}"
|
|
engine = create_engine(url)
|
|
Base.metadata.create_all(engine)
|
|
with engine.begin() as connection:
|
|
connection.exec_driver_sql("DROP TABLE benchmark_runs")
|
|
connection.exec_driver_sql(
|
|
"""
|
|
CREATE TABLE benchmark_runs (
|
|
id CHAR(32) NOT NULL,
|
|
deployment_id CHAR(32) NOT NULL,
|
|
suite_id CHAR(32) NOT NULL,
|
|
model_revision_sha VARCHAR(64) NOT NULL,
|
|
artifact_sha256 VARCHAR(64) NOT NULL,
|
|
runtime_type VARCHAR(32) NOT NULL,
|
|
runtime_version VARCHAR(128) NOT NULL,
|
|
runtime_image_digest VARCHAR(71),
|
|
launch_arguments JSON NOT NULL,
|
|
suite_revision VARCHAR(128) NOT NULL,
|
|
dataset_revision VARCHAR(255) NOT NULL,
|
|
environment_fingerprint JSON NOT NULL,
|
|
environment_digest VARCHAR(64) NOT NULL,
|
|
seed INTEGER,
|
|
results JSON NOT NULL,
|
|
status VARCHAR(32) NOT NULL,
|
|
comparable BOOLEAN NOT NULL,
|
|
stale_reason TEXT,
|
|
started_at DATETIME NOT NULL,
|
|
completed_at DATETIME,
|
|
PRIMARY KEY (id),
|
|
FOREIGN KEY(deployment_id) REFERENCES deployments (id) ON DELETE RESTRICT,
|
|
FOREIGN KEY(suite_id) REFERENCES benchmark_suites (id) ON DELETE RESTRICT
|
|
)
|
|
"""
|
|
)
|
|
connection.exec_driver_sql(
|
|
"CREATE INDEX ix_benchmark_runs_deployment_id ON benchmark_runs (deployment_id)"
|
|
)
|
|
connection.exec_driver_sql(
|
|
"CREATE INDEX ix_benchmark_runs_suite_id ON benchmark_runs (suite_id)"
|
|
)
|
|
connection.exec_driver_sql(
|
|
"CREATE INDEX ix_benchmark_runs_artifact_sha256 ON benchmark_runs (artifact_sha256)"
|
|
)
|
|
connection.exec_driver_sql(
|
|
"CREATE INDEX ix_benchmark_runs_environment_digest ON benchmark_runs (environment_digest)"
|
|
)
|
|
connection.exec_driver_sql(
|
|
"CREATE INDEX ix_benchmark_runs_deployment_completed "
|
|
"ON benchmark_runs (deployment_id, completed_at)"
|
|
)
|
|
|
|
config = Config(str(Path(__file__).parents[1] / "alembic.ini"))
|
|
config.set_main_option("sqlalchemy.url", url)
|
|
command.stamp(config, "20260924_0038")
|
|
command.upgrade(config, "head")
|
|
|
|
upgraded = inspect(engine)
|
|
columns = {column["name"]: column for column in upgraded.get_columns("benchmark_runs")}
|
|
assert columns["deployment_id"]["nullable"] is True
|
|
assert "capability_deployment_id" in columns
|
|
indexes = {index["name"] for index in upgraded.get_indexes("benchmark_runs")}
|
|
assert "ix_benchmark_runs_actor_subject" in indexes
|
|
assert "ix_benchmark_runs_capability_deployment_completed" in indexes
|
|
|
|
command.downgrade(config, "20260924_0038")
|
|
downgraded = {
|
|
column["name"]: column for column in inspect(engine).get_columns("benchmark_runs")
|
|
}
|
|
assert downgraded["deployment_id"]["nullable"] is False
|
|
assert "capability_deployment_id" not in downgraded
|
|
|
|
command.upgrade(config, "head")
|
|
assert "capability_deployment_id" in {
|
|
column["name"] for column in inspect(engine).get_columns("benchmark_runs")
|
|
}
|
|
engine.dispose()
|