Publish curated ITWorx-ModelForge source from 990401737f96c8d46a9f7b0d84994516cadcb6a0

backend/tests/test_project_onboarding_postgres.py

@@ -0,0 +1,246 @@
+"""Isolated PostgreSQL 0025 upgrade/backup restore and credential race qualification."""
+from __future__ import annotations
+
+import hashlib
+import importlib
+import os
+import threading
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from datetime import UTC, datetime, timedelta
+from pathlib import Path
+
+import pytest
+from alembic.config import Config
+from sqlalchemy import create_engine, select, text
+from sqlalchemy.engine import make_url
+from sqlalchemy.orm import Session
+
+from modelforge_api.domain.serving import (
+    ServiceClientCreated,
+    ServiceClientResponse,
+    ServiceCredentialRevoke,
+    ServiceCredentialRotate,
+)
+from modelforge_api.persistence.models import AuditEvent, ServiceCredential
+from modelforge_api.services.audit import AuditWriter
+from modelforge_api.services.manifest_registry import ManifestRegistry
+from modelforge_api.services.recovery_fingerprint import (
+    audit_fingerprint_differences_are_compatible,
+    fingerprint_compatibility,
+    semantic_fingerprint,
+)
+from modelforge_api.services.serving import ServingError, ServingService
+from modelforge_api.settings import Settings
+from tests.test_configuration_postgres import _database_tool
+
+
+@pytest.mark.parametrize(
+    "competitors",
+    [("rotate", "rotate"), ("rotate", "revoke"), ("revoke", "revoke")],
+)
+def test_project_upgrade_restore_and_credential_rotation_races(
+    tmp_path: Path,
+    competitors: tuple[str, str],
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    admin_url = os.getenv("MODELFORGE_TEST_POSTGRES_URL")
+    if not admin_url:
+        pytest.skip("requires managed MODELFORGE_TEST_POSTGRES_URL")
+    owner_url = os.getenv("MODELFORGE_TEST_POSTGRES_OWNER_URL")
+    runtime_url = os.getenv("MODELFORGE_TEST_POSTGRES_RUNTIME_URL")
+    assert owner_url and runtime_url
+    database = "test_onboarding_" + uuid.uuid4().hex
+    restored = "test_onboarding_restore_" + uuid.uuid4().hex
+    admin = create_engine(admin_url, isolation_level="AUTOCOMMIT")
+    owner = make_url(owner_url).set(database=database)
+    runtime = create_engine(make_url(runtime_url).set(database=database))
+    restored_runtime = create_engine(make_url(runtime_url).set(database=restored))
+    config = Config()
+    config.set_main_option("script_location", str(Path(__file__).resolve().parents[1] / "alembic"))
+    config.set_main_option("sqlalchemy.url", owner.render_as_string(hide_password=False).replace("%", "%%"))
+    alembic = importlib.import_module("alembic.command")
+    client_id, credential_id = uuid.uuid4(), uuid.uuid4()
+    secret = "mf_" + "p" * 48
+    with admin.connect() as connection:
+        for name in (database, restored):
+            connection.execute(text(f'CREATE DATABASE "{name}" OWNER modelforge'))
+    try:
+        alembic.upgrade(config, "20260907_0025")
+        with Session(runtime) as session:
+            session.execute(text("INSERT INTO service_clients(id,name,status,allowed_capabilities,requests_per_minute,max_concurrent_requests,workload_priority) VALUES (:id,'prior-client','active','[\"rag.embedding@1\"]',60,1,'interactive')"), {"id": client_id})
+            session.execute(text("INSERT INTO service_credentials(id,service_client_id,secret_hash,secret_prefix,expires_at) VALUES (:id,:client,:digest,'mf_prior',:expiry)"), {"id": credential_id, "client": client_id, "digest": hashlib.sha256(secret.encode()).hexdigest(), "expiry": datetime.now(UTC) + timedelta(days=30)})
+            AuditWriter(session, "operator", "prior-release").write("qualification.baseline", "service_client", str(client_id), {})
+            session.commit()
+            source = semantic_fingerprint(session.connection())
+        archive = tmp_path / "before-onboarding.dump"
+        args = ["--host", "127.0.0.1" if os.getenv("MODELFORGE_TEST_POSTGRES_CONTAINER") else owner.host or "localhost", "--port", str(owner.port or 5432), "--username", owner.username or "modelforge", "--dbname", database]
+        _database_tool("pg_dump", [*args, "--format=custom"], owner.password, archive)
+        assert archive.stat().st_size > 0
+        alembic.upgrade(config, "head")
+        with Session(runtime) as session:
+            current = semantic_fingerprint(session.connection())
+            comparison = fingerprint_compatibility(source, current, source_schema_revision="20260907_0025")
+            assert comparison.compatible
+            assert audit_fingerprint_differences_are_compatible(source, current, comparison)
+            service = ServingService(session, Settings(_env_file=None), ManifestRegistry(), None)
+            assert service.authenticate("Bearer " + secret, "rag.embedding@1").id == client_id
+        alembic.downgrade(config, "20260907_0025")
+        alembic.upgrade(config, "head")
+        with Session(runtime) as session:
+            service = ServingService(session, Settings(_env_file=None), ManifestRegistry(), None)
+            prior_overlap = service.rotate_client_credential(
+                client_id,
+                ServiceCredentialRotate(
+                    expected_credential_id=credential_id,
+                    overlap_seconds=1800,
+                    expires_at=datetime.now(UTC) + timedelta(days=14),
+                ),
+            )
+            assert prior_overlap.credential_id
+            prior_overlap_id = prior_overlap.credential_id
+            prior_overlap_secret = prior_overlap.credential
+            prior_overlap_expires_at = prior_overlap.credential_expires_at
+            assert prior_overlap_expires_at
+            assert service.authenticate("Bearer " + secret, "rag.embedding@1").id == client_id
+        with Session(runtime) as session:
+            existing_event_ids = set(session.scalars(select(AuditEvent.id)))
+
+        # Hold both operations after their current-credential read and before the conditional
+        # claim UPDATE, so the loser must observe the winner's committed CAS change.
+        claim_barrier = threading.Barrier(2)
+        original_claim = ServingService._claim_credential
+
+        def claim_at_cas(service: ServingService, credential: ServiceCredential) -> None:
+            claim_barrier.wait(timeout=30)
+            original_claim(service, credential)
+
+        monkeypatch.setattr(ServingService, "_claim_credential", claim_at_cas)
+        race_credential_id = uuid.UUID(str(prior_overlap_id))
+        expires = datetime.now(UTC) + timedelta(days=14)
+
+        def contender(
+            action: str,
+        ) -> tuple[str, str, ServiceClientCreated | ServiceClientResponse | None]:
+            with Session(runtime) as session:
+                service = ServingService(session, Settings(_env_file=None), ManifestRegistry(), None)
+                try:
+                    if action == "rotate":
+                        result = service.rotate_client_credential(
+                            client_id,
+                            ServiceCredentialRotate(
+                                expected_credential_id=race_credential_id,
+                                overlap_seconds=900,
+                                expires_at=expires,
+                            ),
+                        )
+                    else:
+                        result = service.revoke_client_credential(
+                            client_id,
+                            ServiceCredentialRevoke(
+                                expected_credential_id=race_credential_id,
+                                confirmation="prior-client",
+                            ),
+                        )
+                    return action, "applied", result
+                except ServingError as error:
+                    assert error.code == "CREDENTIAL_CHANGED"
+                    return action, "conflict", None
+
+        with ThreadPoolExecutor(max_workers=2) as pool:
+            outcomes = list(pool.map(contender, competitors))
+        assert sorted(outcome for _action, outcome, _result in outcomes) == [
+            "applied",
+            "conflict",
+        ]
+        winner_action, _, winner_result = next(
+            outcome for outcome in outcomes if outcome[1] == "applied"
+        )
+        with Session(runtime) as session:
+            new_events = list(
+                session.scalars(
+                    select(AuditEvent).where(
+                        AuditEvent.action.in_(
+                            ["CLIENT_CREDENTIAL_ROTATED", "CLIENT_CREDENTIAL_REVOKED"]
+                        ),
+                        AuditEvent.id.not_in(existing_event_ids),
+                    )
+                )
+            )
+            assert len(new_events) == 1
+            assert new_events[0].action == (
+                "CLIENT_CREDENTIAL_ROTATED"
+                if winner_action == "rotate"
+                else "CLIENT_CREDENTIAL_REVOKED"
+            )
+            credentials = list(
+                session.scalars(
+                    select(ServiceCredential).where(
+                        ServiceCredential.service_client_id == client_id
+                    )
+                )
+            )
+            current = session.get(ServiceCredential, race_credential_id)
+            original = session.get(ServiceCredential, credential_id)
+            assert current and original
+            assert current.superseded_at is not None
+            assert original.superseded_at is not None and original.revoked_at is not None
+            active = [
+                item
+                for item in credentials
+                if item.superseded_at is None and item.revoked_at is None
+            ]
+            if winner_action == "rotate":
+                assert isinstance(winner_result, ServiceClientCreated)
+                winner_credential_id = uuid.UUID(str(winner_result.credential_id))
+                assert len(active) == 1 and active[0].id == winner_credential_id
+                assert current.revoked_at is None
+                assert current.expires_at is not None
+                assert datetime.now(UTC) < current.expires_at <= prior_overlap_expires_at
+                assert current.superseded_at is not None
+                assert current.expires_at <= current.superseded_at + timedelta(seconds=900)
+                assert new_events[0].details["previous_credential_id"] == str(
+                    race_credential_id
+                )
+                assert new_events[0].details["credential_id"] == str(winner_credential_id)
+                assert new_events[0].details["overlap_seconds"] == 900
+            else:
+                assert len(active) == 0
+                assert current.revoked_at is not None
+                assert all(item.revoked_at is not None for item in credentials)
+
+        with Session(runtime) as session:
+            service = ServingService(session, Settings(_env_file=None), ManifestRegistry(), None)
+            if winner_action == "rotate":
+                assert isinstance(winner_result, ServiceClientCreated)
+                assert service.authenticate(
+                    "Bearer " + prior_overlap_secret, "rag.embedding@1"
+                ).id == client_id
+                assert service.authenticate(
+                    "Bearer " + winner_result.credential, "rag.embedding@1"
+                ).id == client_id
+                with pytest.raises(ServingError, match="revoked"):
+                    service.authenticate("Bearer " + secret, "rag.embedding@1")
+            else:
+                for retired_secret in (secret, prior_overlap_secret):
+                    with pytest.raises(ServingError, match="revoked"):
+                        service.authenticate("Bearer " + retired_secret, "rag.embedding@1")
+        with pytest.raises(RuntimeError, match="history exists"):
+            alembic.downgrade(config, "20260907_0025")
+        args[-1] = restored
+        _database_tool("pg_restore", [*args, "--exit-on-error"], owner.password, archive)
+        with restored_runtime.connect() as connection:
+            recovered = semantic_fingerprint(connection)
+            assert recovered["tables"] == source["tables"]
+        config.set_main_option("sqlalchemy.url", owner.set(database=restored).render_as_string(hide_password=False).replace("%", "%%"))
+        alembic.upgrade(config, "head")
+        with restored_runtime.connect() as connection:
+            recovered = semantic_fingerprint(connection)
+            comparison = fingerprint_compatibility(source, recovered, source_schema_revision="20260907_0025")
+            assert comparison.compatible and audit_fingerprint_differences_are_compatible(source, recovered, comparison)
+    finally:
+        runtime.dispose(); restored_runtime.dispose()
+        with admin.connect() as connection:
+            for name in (database, restored):
+                connection.execute(text(f'DROP DATABASE IF EXISTS "{name}" WITH (FORCE)'))
+        admin.dispose()