ModelForge-Public/backend/tests/test_self_check.py

from __future__ import annotations

import hashlib
import json
import uuid
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any

from fastapi.testclient import TestClient
from pydantic import SecretStr
from sqlalchemy import create_engine, event, text
from sqlalchemy.orm import Session
from sqlalchemy.pool import StaticPool

from modelforge_api.db import get_session
from modelforge_api.domain.release import TARGET_SCHEMA_REVISION
from modelforge_api.domain.self_check import SelfCheckStatus
from modelforge_api.main import app
from modelforge_api.persistence.models import (
    Accelerator,
    Base,
    ComputeNode,
    ManagedOperatorCredential,
    ManagedOperatorCredentialActivation,
    ManagedOperatorIdentity,
)
from modelforge_api.services.manifest_registry import get_manifest_registry
from modelforge_api.services.self_check import (
    _configuration_check,
    _credentials_check,
    _database_check,
    _manifest_check,
    _node_check,
    _redis_check,
    _scheduler_check,
    _storage_check,
    run_self_check,
)
from modelforge_api.settings import Settings, get_settings

OPERATOR_TOKEN = "self-check-test-operator-key-long-enough"  # noqa: S105 - synthetic test credential
DATABASE_SECRET = "postgresql+psycopg://runtime:db-do-not-leak@db:5432/modelforge"  # noqa: S105
REDIS_SECRET = "redis-do-not-leak"  # noqa: S105 - synthetic test credential


class FakeCandidates:
    def __init__(self) -> None:
        self.candidates = [object()]


class FakePolicies:
    def model_dump(self) -> dict[str, object]:
        return {
            "scheduling": object(),
            "security": object(),
            "promotion": object(),
            "retention": object(),
        }


class FakeRegistry:
    def capabilities(self) -> dict[str, object]:
        return {"a": object()}

    def projects(self) -> dict[str, object]:
        return {"p": object()}

    def candidates(self) -> FakeCandidates:
        return FakeCandidates()

    def benchmarks(self) -> dict[str, object]:
        return {"b": object()}

    def policies(self) -> FakePolicies:
        return FakePolicies()


def _settings(tmp_path: Path, *, env: str = "test") -> Settings:
    artifact = tmp_path / "artifacts"
    quarantine = tmp_path / "quarantine"
    backup = tmp_path / "backups"
    for path in (artifact, quarantine, backup):
        path.mkdir(exist_ok=True)
    return Settings(
        _env_file=None,
        env=env,
        operator_api_key=SecretStr(OPERATOR_TOKEN),
        database_url=DATABASE_SECRET,
        redis_url=f"redis://cache:{REDIS_SECRET}@redis.internal:6379/0",
        artifact_root=str(artifact),
        quarantine_root=str(quarantine),
        backup_root=backup,
    )


def _session() -> Session:
    engine = create_engine("sqlite+pysqlite:///:memory:", connect_args={"check_same_thread": False}, poolclass=StaticPool)
    Base.metadata.create_all(engine)
    with engine.begin() as connection:
        connection.execute(text("CREATE TABLE alembic_version (version_num VARCHAR(32) NOT NULL)"))
        connection.execute(
            text("INSERT INTO alembic_version(version_num) VALUES (:revision)"),
            {"revision": TARGET_SCHEMA_REVISION},
        )
    return Session(engine)


class RedisOk:
    def ping(self) -> bool:
        return True

    def close(self) -> None:
        return None


class RedisDown:
    def ping(self) -> bool:
        raise ConnectionError("synthetic Redis outage")

    def close(self) -> None:
        return None


def test_storage_probe_is_passive_and_reports_capacity(tmp_path: Path) -> None:
    settings = _settings(tmp_path)
    before = {
        path: set(Path(path).iterdir())
        for path in (settings.artifact_root, settings.quarantine_root, settings.backup_root)
    }

    result = _storage_check(settings)

    assert result.status is SelfCheckStatus.PASS
    assert all(
        item["writability"] == "not_probed_read_only" for item in result.evidence.values()
    )
    assert all(item["free_bytes"] >= 0 for item in result.evidence.values())
    after = {path: set(Path(path).iterdir()) for path in before}
    assert after == before


def test_storage_capacity_read_failure_remains_unknown(
    tmp_path: Path, monkeypatch: Any
) -> None:
    settings = _settings(tmp_path)

    def unavailable(_path: Path) -> object:
        raise PermissionError("mount details must not be returned")

    monkeypatch.setattr("modelforge_api.services.self_check.shutil.disk_usage", unavailable)
    result = _storage_check(settings)

    assert result.status is SelfCheckStatus.UNKNOWN
    assert "mount details must not be returned" not in result.model_dump_json()


def test_storage_stat_permission_failure_remains_unknown(tmp_path: Path, monkeypatch: Any) -> None:
    settings = _settings(tmp_path)

    class UnreadablePath:
        def __init__(self, value: object) -> None:
            self.value = value

        def stat(self) -> object:
            raise PermissionError("private storage path details must not be returned")

        def __str__(self) -> str:
            return str(self.value)

    monkeypatch.setattr("modelforge_api.services.self_check.Path", UnreadablePath)
    result = _storage_check(settings)

    assert result.status is SelfCheckStatus.UNKNOWN
    assert "private storage path details must not be returned" not in result.model_dump_json()


def test_database_check_fails_closed_on_unknown_schema() -> None:
    session = _session()
    try:
        session.execute(text("UPDATE alembic_version SET version_num='20991231_9999'"))
        session.commit()
        result = _database_check(session)
    finally:
        session.close()

    assert result.status is SelfCheckStatus.FAIL
    assert result.evidence["compatibility"] == "TOO_NEW"


def test_database_check_keeps_unreadable_state_unknown() -> None:
    class UnavailableSession:
        def execute(self, _statement: object) -> object:
            raise ConnectionError("database URL must not be returned")

        def rollback(self) -> None:
            return None

    result = _database_check(UnavailableSession())  # type: ignore[arg-type]

    assert result.status is SelfCheckStatus.UNKNOWN
    assert result.detail == "ConnectionError"
    assert "database URL must not be returned" not in result.model_dump_json()


def test_manifest_access_failure_is_unknown_but_missing_manifest_is_failure() -> None:
    class UnreadableRegistry:
        def candidates(self) -> object:
            raise PermissionError("private manifest path details must not be returned")

    class MissingRegistry:
        def candidates(self) -> object:
            raise FileNotFoundError("missing manifest")

    unreadable = _manifest_check(UnreadableRegistry())  # type: ignore[arg-type]
    missing = _manifest_check(MissingRegistry())  # type: ignore[arg-type]

    assert unreadable.status is SelfCheckStatus.UNKNOWN
    assert "private manifest path details must not be returned" not in unreadable.model_dump_json()
    assert missing.status is SelfCheckStatus.FAIL


def test_scheduler_inputs_remain_unknown_when_database_cannot_be_read() -> None:
    class UnavailableSession:
        def scalar(self, _statement: object) -> object:
            raise ConnectionError("scheduler database details must not be returned")

        def rollback(self) -> None:
            return None

    settings = Settings(_env_file=None)
    result = _scheduler_check(
        UnavailableSession(), settings, datetime.now(UTC)  # type: ignore[arg-type]
    )

    assert result.status is SelfCheckStatus.UNKNOWN
    assert result.detail == "ConnectionError"
    assert "scheduler database details must not be returned" not in result.model_dump_json()


def test_scheduler_inputs_remain_unknown_without_fresh_accelerator_telemetry(
    tmp_path: Path,
) -> None:
    settings = _settings(tmp_path)
    session = _session()
    now = datetime(2026, 9, 6, tzinfo=UTC)
    try:
        node = ComputeNode(
            key="scheduler-node",
            hostname="scheduler-node",
            display_name="Scheduler node",
            identity_source="test",
        )
        session.add(node)
        session.flush()
        session.add(
            Accelerator(
                compute_node_id=node.id,
                device_index=0,
                device_uuid="GPU-scheduler-test",
                name="Test accelerator",
                total_vram_bytes=8_000_000_000,
                status="active",
            )
        )
        session.commit()

        result = _scheduler_check(session, settings, now)
    finally:
        session.close()

    assert result.status is SelfCheckStatus.UNKNOWN
    assert result.evidence["active_accelerators"] == 1
    assert result.evidence["telemetry_missing"] == 1


def test_redis_probe_is_typed_and_never_echoes_the_connection_url(
    tmp_path: Path, monkeypatch: Any
) -> None:
    settings = _settings(tmp_path)
    monkeypatch.setattr(
        "modelforge_api.services.self_check.redis.Redis.from_url",
        lambda *_args, **_kwargs: RedisDown(),
    )

    result = _redis_check(settings)

    assert result.status is SelfCheckStatus.FAIL
    assert REDIS_SECRET not in result.model_dump_json()
    assert result.detail == "ConnectionError"


def test_node_freshness_uses_existing_stale_and_offline_thresholds(tmp_path: Path) -> None:
    settings = _settings(tmp_path)
    session = _session()
    now = datetime(2026, 9, 6, tzinfo=UTC)
    try:
        fresh = ComputeNode(
            key="fresh",
            hostname="fresh",
            display_name="Fresh",
            identity_source="test",
            last_heartbeat_at=now
            - timedelta(seconds=max(1, settings.node_stale_after_seconds // 2)),
        )
        stale = ComputeNode(
            key="stale",
            hostname="stale",
            display_name="Stale",
            identity_source="test",
            last_heartbeat_at=now - timedelta(seconds=settings.node_stale_after_seconds + 1),
        )
        session.add_all([fresh, stale])
        session.commit()

        result = _node_check(session, settings, now)
    finally:
        session.close()

    assert result.status is SelfCheckStatus.WARN
    assert result.evidence["stale_nodes"] == ["Stale"]
    assert result.evidence["offline_nodes"] == []


def test_node_without_any_heartbeat_evidence_is_unknown(tmp_path: Path) -> None:
    settings = _settings(tmp_path)
    session = _session()
    try:
        session.add(
            ComputeNode(
                key="unobserved",
                hostname="unobserved",
                display_name="Unobserved",
                identity_source="test",
            )
        )
        session.commit()
        result = _node_check(session, settings, datetime(2026, 9, 6, tzinfo=UTC))
    finally:
        session.close()

    assert result.status is SelfCheckStatus.UNKNOWN
    assert result.evidence["heartbeat_unknown"] == ["Unobserved"]


def test_configuration_check_reports_missing_names_not_secret_values(tmp_path: Path) -> None:
    settings = _settings(tmp_path, env="production")

    result = _configuration_check(settings)

    assert result.status is SelfCheckStatus.FAIL
    serialized = result.model_dump_json()
    assert OPERATOR_TOKEN not in serialized
    assert DATABASE_SECRET not in serialized
    assert REDIS_SECRET not in serialized
    assert "missing_settings" in result.evidence


def test_credentials_check_reports_only_safe_metadata(tmp_path: Path) -> None:
    settings = _settings(tmp_path)
    settings.operator_api_key = None
    session = _session()
    now = datetime(2026, 9, 6, tzinfo=UTC)
    credential_id = uuid.uuid4()
    try:
        session.add(
            ManagedOperatorIdentity(
                slug="admin",
                role="admin",
                enabled=True,
                version=1,
                created_at=now,
            )
        )
        session.add(
            ManagedOperatorCredential(
                id=credential_id,
                identity_slug="admin",
                token_sha256="a" * 64,
                expires_at=now + timedelta(days=30),
                created_at=now,
            )
        )
        session.add(
            ManagedOperatorCredentialActivation(
                credential_id=credential_id,
                activated_at=now,
            )
        )
        session.commit()
        result = _credentials_check(session, settings, now)
    finally:
        session.close()

    assert result.status is SelfCheckStatus.PASS
    assert result.evidence["managed_usable_admin_credentials"] == 1
    assert result.evidence["administrator_metadata_present"] is True
    serialized = result.model_dump_json()
    assert "a" * 64 not in serialized
    assert "token_sha256" not in serialized
    assert OPERATOR_TOKEN not in serialized
    assert DATABASE_SECRET not in serialized


def test_credentials_check_warns_for_expiring_managed_metadata(tmp_path: Path) -> None:
    settings = _settings(tmp_path)
    session = _session()
    now = datetime(2026, 9, 6, tzinfo=UTC)
    try:
        session.add(
            ManagedOperatorIdentity(
                slug="admin",
                role="admin",
                enabled=True,
                version=1,
                created_at=now,
            )
        )
        session.add(
            ManagedOperatorCredential(
                id=uuid.uuid4(),
                identity_slug="admin",
                token_sha256="b" * 64,
                expires_at=now + timedelta(days=3),
                created_at=now,
            )
        )
        session.commit()
        result = _credentials_check(session, settings, now)
    finally:
        session.close()

    assert result.status is SelfCheckStatus.WARN
    assert result.evidence["managed_expiring_within_7_days"] == 1
    assert "b" * 64 not in result.model_dump_json()


def test_credential_metadata_read_failure_remains_unknown(tmp_path: Path) -> None:
    class UnavailableSession:
        def execute(self, _statement: object) -> object:
            raise ConnectionError("credential database details must not be returned")

        def rollback(self) -> None:
            return None

    settings = _settings(tmp_path)
    result = _credentials_check(
        UnavailableSession(), settings, datetime.now(UTC)  # type: ignore[arg-type]
    )

    assert result.status is SelfCheckStatus.UNKNOWN
    assert result.detail == "ConnectionError"
    serialized = result.model_dump_json()
    assert "credential database details must not be returned" not in serialized
    assert OPERATOR_TOKEN not in serialized
    assert DATABASE_SECRET not in serialized


def test_credential_metadata_without_an_admin_source_fails_closed(tmp_path: Path) -> None:
    settings = _settings(tmp_path)
    settings.operator_api_key = None
    viewer_token = "mfop_" + "v" * 43
    operator_token = "mfop_" + "o" * 43
    settings.operator_identities = SecretStr(
        json.dumps(
            [
                {
                    "credential_id": role,
                    "subject": role,
                    "role": role,
                    "token_sha256": hashlib.sha256(token.encode()).hexdigest(),
                    "expires_at": (datetime.now(UTC) + timedelta(days=30)).isoformat(),
                }
                for role, token in (
                    ("viewer", viewer_token),
                    ("operator", operator_token),
                )
            ]
        )
    )
    session = _session()
    try:
        result = _credentials_check(session, settings, datetime.now(UTC))
    finally:
        session.close()

    assert result.status is SelfCheckStatus.FAIL
    assert result.evidence["administrator_metadata_present"] is False
    serialized = result.model_dump_json()
    assert viewer_token not in serialized
    assert operator_token not in serialized


def test_self_check_endpoint_is_admin_authenticated_read_only_and_secret_safe(
    tmp_path: Path, monkeypatch: Any
) -> None:
    settings = _settings(tmp_path)
    session = _session()
    monkeypatch.setattr(
        "modelforge_api.services.self_check.redis.Redis.from_url",
        lambda *_args, **_kwargs: RedisOk(),
    )
    app.dependency_overrides[get_settings] = lambda: settings
    app.dependency_overrides[get_session] = lambda: session
    app.dependency_overrides[get_manifest_registry] = lambda: FakeRegistry()
    engine = session.get_bind()
    mutations: list[str] = []

    def reject_mutation(
        _connection: Any,
        _cursor: Any,
        statement: str,
        _parameters: Any,
        _context: Any,
        _executemany: bool,
    ) -> None:
        operation = statement.lstrip().split(maxsplit=1)[0].upper()
        if operation in {"INSERT", "UPDATE", "DELETE", "REPLACE", "CREATE", "DROP", "ALTER"}:
            mutations.append(operation)

    event.listen(engine, "before_cursor_execute", reject_mutation)
    client = TestClient(app)
    try:
        anonymous = client.get("/api/v1/admin/operations/self-check")
        assert anonymous.status_code == 401

        response = client.get(
            "/api/v1/admin/operations/self-check",
            headers={"X-ModelForge-Admin-Token": OPERATOR_TOKEN},
        )
        assert response.status_code == 200
        payload = response.json()
        assert payload["read_only"] is True
        assert payload["overall"] == "degraded"
        assert {check["key"] for check in payload["checks"]} == {
            "api",
            "database",
            "redis",
            "manifests",
            "storage",
            "nodes",
            "scheduler",
            "deployments",
            "backup",
            "credentials",
            "configuration",
            "release",
        }
        assert all(isinstance(check["next_action"], dict) for check in payload["checks"])
        assert all(check["next_action"]["code"] for check in payload["checks"])
        assert (
            next(check for check in payload["checks"] if check["key"] == "scheduler")["status"]
            == "unknown"
        )
        api_check = next(check for check in payload["checks"] if check["key"] == "api")
        assert api_check["status"] == "pass"
        assert api_check["next_action"]["code"] == "none_required"
        serialized = response.text
        assert OPERATOR_TOKEN not in serialized
        assert DATABASE_SECRET not in serialized
        assert REDIS_SECRET not in serialized

        settings.operator_identities = SecretStr(
            json.dumps(
                [
                    {
                        "credential_id": f"self-check-{role}",
                        "subject": f"self-check-{role}",
                        "role": role,
                        "token_sha256": hashlib.sha256(token.encode()).hexdigest(),
                        "expires_at": (datetime.now(UTC) + timedelta(days=30)).isoformat(),
                    }
                        for role, token in (
                            ("viewer", "mfop_" + "v" * 43),
                            ("operator", "mfop_" + "o" * 43),
                    )
                ]
            )
        )
        for token in ("mfop_" + "v" * 43, "mfop_" + "o" * 43):
            denied = client.get(
                "/api/v1/admin/operations/self-check",
                headers={"X-ModelForge-Admin-Token": token},
            )
            assert denied.status_code == 403
        for token in ("mfsvc_project-scoped-token", "mfnode_node-agent-token"):
            denied = client.get(
                "/api/v1/admin/operations/self-check",
                headers={"X-ModelForge-Admin-Token": token},
            )
            assert denied.status_code == 401
    finally:
        client.close()
        app.dependency_overrides.clear()
        event.remove(engine, "before_cursor_execute", reject_mutation)
        session.close()
    assert mutations == []


def test_aggregate_becomes_unhealthy_when_a_required_dependency_fails(
    tmp_path: Path, monkeypatch: Any
) -> None:
    settings = _settings(tmp_path)
    session = _session()
    monkeypatch.setattr(
        "modelforge_api.services.self_check.redis.Redis.from_url",
        lambda *_args, **_kwargs: RedisDown(),
    )
    try:
        report = run_self_check(session, settings, FakeRegistry())  # type: ignore[arg-type]
    finally:
        session.close()

    assert report.overall == "unhealthy"
    redis_check = next(check for check in report.checks if check.key == "redis")
    assert redis_check.status is SelfCheckStatus.FAIL