119 lines
5.2 KiB
Python
119 lines
5.2 KiB
Python
import copy
|
|
from datetime import UTC, datetime, timedelta
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import Session
|
|
|
|
from modelforge_api.persistence.models import (
|
|
Base,
|
|
Model,
|
|
ModelRevision,
|
|
ModelUpdateNotice,
|
|
ModelUpdateWatch,
|
|
)
|
|
from modelforge_api.services.audit import AuditWriter
|
|
from modelforge_api.services.recovery_fingerprint import (
|
|
CREATIVE_TABLES,
|
|
MODEL_UPDATE_TABLES,
|
|
SECRET_SCHEMA,
|
|
SELECTED_COMPARISON_SCHEMA,
|
|
SELECTED_COMPARISON_TABLES,
|
|
_without_creative,
|
|
_without_lab_execution,
|
|
_without_model_changes,
|
|
_without_model_updates,
|
|
_without_portfolio,
|
|
_without_selected_comparisons,
|
|
audit_fingerprint_differences_are_compatible,
|
|
fingerprint_compatibility,
|
|
)
|
|
from modelforge_api.services.recovery_fingerprint import (
|
|
fingerprint_session as _current_fingerprint_session,
|
|
)
|
|
|
|
|
|
def fingerprint_session(session: Session):
|
|
return _without_lab_execution(
|
|
_without_portfolio(_without_model_changes(_current_fingerprint_session(session)))
|
|
)
|
|
|
|
|
|
def test_previous_schema_allows_only_exact_empty_update_introduction() -> None:
|
|
engine = create_engine("sqlite://")
|
|
Base.metadata.create_all(engine)
|
|
with Session(engine) as session:
|
|
restored = fingerprint_session(session)
|
|
source = _without_creative(_without_selected_comparisons(_without_model_updates(restored)))
|
|
assert not (CREATIVE_TABLES | SELECTED_COMPARISON_TABLES | MODEL_UPDATE_TABLES) & source["tables"].keys()
|
|
contract = fingerprint_compatibility(source, restored, source_schema_revision=SECRET_SCHEMA)
|
|
assert contract.compatible
|
|
assert audit_fingerprint_differences_are_compatible(source, restored, contract)
|
|
for table in MODEL_UPDATE_TABLES:
|
|
for field, value in (("row_count", 1), ("digest", "0" * 64), ("status", "BOUNDED")):
|
|
changed = copy.deepcopy(restored)
|
|
changed["tables"][table][field] = value
|
|
assert not fingerprint_compatibility(source, changed, source_schema_revision=SECRET_SCHEMA).compatible
|
|
assert not audit_fingerprint_differences_are_compatible(source, changed, contract)
|
|
missing = copy.deepcopy(restored)
|
|
del missing["tables"][table]
|
|
assert not fingerprint_compatibility(source, missing, source_schema_revision=SECRET_SCHEMA).compatible
|
|
engine.dispose()
|
|
|
|
|
|
def test_watch_claim_is_transient_but_baseline_and_notice_history_are_authoritative() -> None:
|
|
engine = create_engine("sqlite://")
|
|
Base.metadata.create_all(engine)
|
|
with Session(engine) as session:
|
|
model = Model(key="watched", display_name="Watched", upstream_provider="huggingface", upstream_source="org/watched")
|
|
session.add(model)
|
|
session.flush()
|
|
baseline = ModelRevision(model_id=model.id, upstream_revision="main", resolved_commit_sha="a" * 40)
|
|
candidate = ModelRevision(model_id=model.id, upstream_revision="main", resolved_commit_sha="b" * 40)
|
|
session.add_all([baseline, candidate])
|
|
session.flush()
|
|
watch = ModelUpdateWatch(model_id=model.id, baseline_revision_id=baseline.id,
|
|
baseline_commit_sha=baseline.resolved_commit_sha, upstream_ref="main",
|
|
check_token=uuid4(), check_until=datetime.now(UTC) + timedelta(minutes=5))
|
|
session.add(watch)
|
|
session.flush()
|
|
notice = ModelUpdateNotice(watch_id=watch.id, model_id=model.id,
|
|
baseline_commit_sha=baseline.resolved_commit_sha, candidate_revision_id=candidate.id,
|
|
candidate_commit_sha=candidate.resolved_commit_sha)
|
|
session.add(notice)
|
|
session.commit()
|
|
before = fingerprint_session(session)
|
|
watch.check_token = None
|
|
watch.check_until = None
|
|
session.commit()
|
|
cleared = fingerprint_session(session)
|
|
assert before["tables"] == cleared["tables"]
|
|
AuditWriter(session, "operator", "restore-test").write("RECOVERY_RECONCILIATION_COMPLETED", "restore", None, {})
|
|
session.commit()
|
|
restored = fingerprint_session(session)
|
|
contract = fingerprint_compatibility(before, restored, source_schema_revision=SELECTED_COMPARISON_SCHEMA)
|
|
assert contract.compatible
|
|
assert audit_fingerprint_differences_are_compatible(before, restored, contract)
|
|
notice.status = "dismissed"
|
|
session.commit()
|
|
changed = fingerprint_session(session)
|
|
assert not audit_fingerprint_differences_are_compatible(before, changed, contract)
|
|
notice.status = "new"
|
|
watch.baseline_commit_sha = "c" * 40
|
|
session.commit()
|
|
assert not audit_fingerprint_differences_are_compatible(before, fingerprint_session(session), contract)
|
|
engine.dispose()
|
|
|
|
|
|
@pytest.mark.parametrize("table", sorted(MODEL_UPDATE_TABLES))
|
|
def test_current_schema_never_projects_away_missing_update_state(table: str) -> None:
|
|
engine = create_engine("sqlite://")
|
|
Base.metadata.create_all(engine)
|
|
with Session(engine) as session:
|
|
before = fingerprint_session(session)
|
|
restored = copy.deepcopy(before)
|
|
del restored["tables"][table]
|
|
assert not fingerprint_compatibility(before, restored, source_schema_revision=SELECTED_COMPARISON_SCHEMA).compatible
|
|
engine.dispose()
|