Publish curated ITWorx-ModelForge source from 990401737f96c8d46a9f7b0d84994516cadcb6a0

backend/tests/test_benchmark_runs_migration.py

@@ -0,0 +1,89 @@
+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()