157 lines
7.7 KiB
Python
157 lines
7.7 KiB
Python
"""Actual owner/runtime PostgreSQL upgrade and dump/restore qualification."""
|
|
|
|
import importlib
|
|
import os
|
|
import uuid
|
|
from datetime import UTC, datetime, timedelta
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
from sqlalchemy import inspect, select, text
|
|
from sqlalchemy.exc import DBAPIError
|
|
from sqlalchemy.orm import Session
|
|
|
|
from modelforge_api.domain.model_updates import ModelWatchCreate, ModelWatchUpdate
|
|
from modelforge_api.domain.release import TARGET_SCHEMA_REVISION
|
|
from modelforge_api.persistence.models import Model, ModelRevision, ModelUpdateWatch
|
|
from modelforge_api.services.acquisition import AcquisitionService
|
|
from modelforge_api.services.audit import AuditContext, AuditWriter
|
|
from modelforge_api.services.model_updates import ModelUpdateService
|
|
from modelforge_api.services.recovery import RecoveryError, RecoveryService
|
|
from modelforge_api.services.recovery_fingerprint import (
|
|
audit_fingerprint_differences_are_compatible,
|
|
fingerprint_compatibility,
|
|
semantic_fingerprint,
|
|
)
|
|
from modelforge_api.settings import Settings
|
|
from tests.test_managed_identities_postgres import (
|
|
IdentityDatabase,
|
|
ManagedPostgresTools,
|
|
identity_database,
|
|
)
|
|
|
|
__all__ = ["identity_database"]
|
|
PRIOR = "20260909_0029"
|
|
CURRENT = "20260910_0034"
|
|
|
|
|
|
def test_model_watch_upgrade_prior_restore_and_runtime_permissions(
|
|
identity_database: IdentityDatabase, tmp_path: Path,
|
|
) -> None:
|
|
database = identity_database
|
|
alembic = importlib.import_module("alembic.command")
|
|
alembic.upgrade(database.migration(), PRIOR)
|
|
with database.runtime.connect() as connection:
|
|
prior_tables = set(inspect(connection).get_table_names())
|
|
assert not {"model_update_watches", "model_update_notices"} & prior_tables
|
|
with Session(database.runtime) as session:
|
|
model = Model(key="selected-model", display_name="Selected model",
|
|
upstream_provider="huggingface", upstream_source="fixture/selected")
|
|
session.add(model)
|
|
session.flush()
|
|
revision = ModelRevision(model_id=model.id, upstream_revision="main",
|
|
resolved_commit_sha="a" * 40)
|
|
session.add(revision)
|
|
AuditWriter(session, "operator", "prior-release").write(
|
|
"qualification.baseline", "model", str(model.id), {})
|
|
session.commit()
|
|
request = ModelWatchCreate(model_id=model.id, baseline_revision_id=revision.id)
|
|
before = semantic_fingerprint(session.connection())
|
|
archive = tmp_path / "pre-model-updates.dump"
|
|
database.archive(archive)
|
|
alembic.upgrade(database.migration(), CURRENT)
|
|
with database.runtime.connect() as connection:
|
|
after = semantic_fingerprint(connection)
|
|
comparison = fingerprint_compatibility(before, after, source_schema_revision=PRIOR)
|
|
assert comparison.compatible
|
|
assert audit_fingerprint_differences_are_compatible(before, after, comparison)
|
|
# Empty tables are mechanically reversible; populated history is deliberately retained.
|
|
alembic.downgrade(database.migration(), PRIOR)
|
|
alembic.upgrade(database.migration(), CURRENT)
|
|
with Session(database.runtime) as session:
|
|
service = ModelUpdateService(session, Mock(spec=AcquisitionService),
|
|
AuditContext("operator", "qualified", "model-watch-pg"))
|
|
watch = service.create(request)
|
|
assert session.scalar(select(ModelUpdateWatch.baseline_commit_sha)) == "a" * 40
|
|
assert watch.model_id == request.model_id
|
|
assert watch.enabled and watch.version == 1
|
|
paused = service.change(watch.id, ModelWatchUpdate(
|
|
expected_version=watch.version, expected_baseline_revision_id=watch.baseline_revision_id,
|
|
baseline_revision_id=watch.baseline_revision_id, enabled=False, reason="Qualified pause",
|
|
))
|
|
assert not paused.enabled and paused.version == 2
|
|
with database.runtime.connect() as connection:
|
|
for statement in (
|
|
"DELETE FROM model_update_watches",
|
|
"TRUNCATE model_update_notices CASCADE",
|
|
"ALTER TABLE model_update_watches ADD COLUMN forbidden integer",
|
|
):
|
|
with pytest.raises(DBAPIError):
|
|
connection.execute(text(statement))
|
|
connection.rollback()
|
|
with pytest.raises(RuntimeError, match="history exists"):
|
|
alembic.downgrade(database.migration(), PRIOR)
|
|
database.archive(archive, restore=True)
|
|
alembic.upgrade(database.migration(restored=True), CURRENT)
|
|
with database.restored_runtime.connect() as connection:
|
|
restored = semantic_fingerprint(connection)
|
|
comparison = fingerprint_compatibility(before, restored, source_schema_revision=PRIOR)
|
|
assert comparison.compatible
|
|
assert audit_fingerprint_differences_are_compatible(before, restored, comparison)
|
|
assert connection.scalar(text("SELECT count(*) FROM model_update_watches")) == 0
|
|
|
|
|
|
def test_model_watch_restore_clears_execution_claim_but_preserves_tracking(
|
|
identity_database: IdentityDatabase, tmp_path: Path,
|
|
) -> None:
|
|
database = identity_database
|
|
importlib.import_module("alembic.command").upgrade(database.migration(), TARGET_SCHEMA_REVISION)
|
|
with Session(database.runtime) as session:
|
|
model = Model(key="restore-watch", display_name="Restore watch",
|
|
upstream_provider="huggingface", upstream_source="fixture/restore")
|
|
session.add(model)
|
|
session.flush()
|
|
revision = ModelRevision(model_id=model.id, upstream_revision="main",
|
|
resolved_commit_sha="a" * 40)
|
|
session.add(revision)
|
|
session.commit()
|
|
watch = ModelUpdateService(session, Mock(spec=AcquisitionService),
|
|
AuditContext("operator", "restore-test")).create(
|
|
ModelWatchCreate(model_id=model.id, baseline_revision_id=revision.id))
|
|
row = session.get(ModelUpdateWatch, watch.id)
|
|
assert row is not None
|
|
row.check_token = uuid.uuid4()
|
|
row.check_until = datetime.now(UTC) + timedelta(minutes=5)
|
|
session.commit()
|
|
source = semantic_fingerprint(session.connection())
|
|
archive = tmp_path / "model-watch.dump"
|
|
database.archive(archive)
|
|
database.archive(archive, restore=True)
|
|
settings = Settings(_env_file=None, env="test",
|
|
alembic_directory=Path(__file__).resolve().parents[1] / "alembic")
|
|
with Session(database.runtime) as session:
|
|
recovery = RecoveryService(session, settings)
|
|
recovery.engine = ManagedPostgresTools(base_environment=dict(os.environ), timeout_seconds=60)
|
|
record = SimpleNamespace(id=uuid.uuid4(), source_fingerprint=source)
|
|
plan = SimpleNamespace(
|
|
database_destination=database.restored_owner.render_as_string(hide_password=False),
|
|
backup_set=SimpleNamespace(backup_id="watch-qualification", schema_revision=TARGET_SCHEMA_REVISION,
|
|
completed_at=datetime.now(UTC)), validation_requirements={},
|
|
)
|
|
with pytest.raises(RecoveryError):
|
|
recovery._validate_restored(record, plan)
|
|
session.rollback()
|
|
recovery._reconcile_restored(record, plan)
|
|
session.commit()
|
|
assert all(recovery._validate_restored(record, plan)["gate"].values())
|
|
session.commit()
|
|
with database.restored_runtime.connect() as connection:
|
|
assert connection.execute(text(
|
|
"SELECT check_token, check_until FROM model_update_watches"
|
|
)).one() == (None, None)
|
|
restored = semantic_fingerprint(connection)
|
|
for table in ("model_update_watches", "model_update_notices"):
|
|
assert restored["tables"][table] == source["tables"][table]
|