awoooi/apps/api/tests/test_ai_automation_asset_capability_matrix.py

from __future__ import annotations

import asyncio
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any

import src.services.ai_automation_asset_capability_matrix as matrix_service
from src.jobs.asset_capability_reconciliation_job import (
    build_reconciliation_summary_payload,
    build_reconciliation_work_item_input,
)
from src.services.ai_automation_asset_capability_matrix import (
    REQUIRED_CONTROL_IDS,
    STAGE_IDS,
    _load_live_rows,
    build_asset_capability_matrix,
    build_asset_capability_matrix_with_live_readback,
    build_reconciliation_candidates,
)
from src.services.awoooi_priority_work_order_readback import (
    apply_ai_automation_asset_capability_matrix,
    get_ai_automation_program_scope_classes,
)

_REPO_ROOT = Path(__file__).resolve().parents[3]


def _single_host_scope() -> list[dict[str, object]]:
    return [
        {
            "id": "hosts",
            "label": "Hosts",
            "asset_count": 1,
            "asset_ids": ["99"],
            "inventory_sources": ["host-inventory.json"],
        }
    ]


def _fully_covered_host(now: datetime) -> dict[str, object]:
    dimensions = (
        "auto_monitoring",
        "auto_alerting",
        "auto_rule_creation",
        "auto_rule_matching",
        "auto_playbook",
        "auto_remediation",
        "auto_km_creation",
    )
    return {
        "asset_id": 42,
        "asset_key": "host:99",
        "asset_type": "host",
        "environment": "prod",
        "namespace": "",
        "name": "99",
        "owner_team": "platform_sre",
        "source_repo": "wooo/awoooi",
        "lifecycle_state": "active",
        "last_seen_at": now,
        "coverage": {
            dimension: {
                "coverage_status": "green",
                "created_at": now,
            }
            for dimension in dimensions
        },
        "compliance": {
            "backup_tested": {
                "status": "compliant",
                "detected_at": now,
            }
        },
        "change_types": [],
    }


def test_declared_catalog_expands_to_193_unique_assets() -> None:
    matrix = build_asset_capability_matrix(
        get_ai_automation_program_scope_classes(),
        repo_root=_REPO_ROOT,
    )

    assert matrix["summary"]["declared_asset_count"] == 193
    assert matrix["summary"]["matrix_asset_count"] == 193
    assert len({asset["canonical_id"] for asset in matrix["assets"]}) == 193
    assert {
        item["scope_id"]: item["declared_asset_count"]
        for item in matrix["scope_rollups"]
    } == {
        "hosts": 7,
        "services": 31,
        "products": 12,
        "websites": 31,
        "tools": 23,
        "packages": 61,
        "data_and_backup": 18,
        "observability_and_security": 10,
    }
    assert matrix["controls"]["full_scope_class_coverage"] is True
    assert matrix["operation_boundaries"]["raw_asset_key_exposed"] is False


def test_public_projection_counts_live_only_assets_without_embedding_them() -> None:
    now = datetime(2026, 7, 15, 1, 0, tzinfo=UTC)
    live_assets = [
        {
            "asset_id": index + 1,
            "asset_key": f"zxq{index:x}v",
            "asset_type": "host",
            "name": f"zxq{index:x}v",
            "last_seen_at": now,
            "coverage": {},
            "compliance": {},
            "change_types": [],
        }
        for index in range(1_000)
    ]

    public_matrix = build_asset_capability_matrix(
        _single_host_scope(),
        live_assets=live_assets,
        repo_root=_REPO_ROOT,
        generated_at=now,
        include_live_only_assets=False,
    )
    reconciliation_matrix = build_asset_capability_matrix(
        _single_host_scope(),
        live_assets=live_assets,
        repo_root=_REPO_ROOT,
        generated_at=now,
        include_live_only_assets=True,
    )

    assert len(public_matrix["assets"]) == 1
    public_summary = public_matrix["summary"]
    assert public_summary["live_inventory_asset_count"] == 1_000
    assert public_summary["matrix_asset_count"] == 1
    assert public_summary["reconciled_asset_count"] == 1_001
    assert public_summary["live_only_assets_embedded"] is False
    assert public_summary["embedded_live_only_asset_count"] == 0
    assert public_summary["discovered_not_declared_count"] == 1_000
    assert len(reconciliation_matrix["assets"]) == 1_001
    assert reconciliation_matrix["summary"]["live_only_assets_embedded"] is True
    assert reconciliation_matrix["summary"]["embedded_live_only_asset_count"] == 1_000
    assert (
        public_matrix["matrix_fingerprint"]
        == (reconciliation_matrix["matrix_fingerprint"])
    )


def test_indexed_matching_preserves_score_cap_and_original_live_order() -> None:
    now = datetime(2026, 7, 15, 1, 0, tzinfo=UTC)
    scope = [
        {
            "id": "services",
            "label": "Services",
            "asset_count": 1,
            "asset_ids": ["one two three four five six seven"],
            "inventory_sources": ["service-inventory.json"],
        }
    ]
    live_assets = [
        {
            "asset_id": 1,
            "asset_key": "candidate:first",
            "asset_type": "container",
            "name": "one two three four five six alpha",
            "last_seen_at": now,
        },
        {
            "asset_id": 2,
            "asset_key": "candidate:second",
            "asset_type": "container",
            "name": "one two three four five six seven beta",
            "last_seen_at": now,
        },
    ]

    matrix = build_asset_capability_matrix(
        scope,
        live_assets=live_assets,
        repo_root=_REPO_ROOT,
        generated_at=now,
        include_live_only_assets=False,
    )

    assert matrix["assets"][0]["runtime_identity"]["asset_inventory_id"] == 1


def test_live_asset_has_per_stage_covered_status_and_closes_after_receipt() -> None:
    now = datetime(2026, 7, 11, 12, 0, tzinfo=UTC)
    initial = build_asset_capability_matrix(
        _single_host_scope(),
        live_assets=[_fully_covered_host(now)],
        latest_run={"run_id": "run-1", "status": "success", "ended_at": now},
        repo_root=_REPO_ROOT,
        generated_at=now,
    )
    receipt = {
        "op_id": "receipt-1",
        "status": "success",
        "created_at": now,
        "input": {
            "trace_id": "trace-1",
            "automation_run_id": "automation-run-1",
            "work_item_id": "AIA-P0-006",
        },
        "output": {
            "matrix_fingerprint": initial["matrix_fingerprint"],
            "candidate_count": 0,
            "repository_readback_count": 0,
            "repository_readback_verified": True,
        },
    }

    matrix = build_asset_capability_matrix(
        _single_host_scope(),
        live_assets=[_fully_covered_host(now)],
        latest_run={"run_id": "run-1", "status": "success", "ended_at": now},
        reconciliation_receipt=receipt,
        repo_root=_REPO_ROOT,
        generated_at=now,
    )

    asset = matrix["assets"][0]
    assert set(asset["stages"]) == set(STAGE_IDS)
    assert {stage["status"] for stage in asset["stages"].values()} == {"covered"}
    assert all(matrix["controls"][control_id] for control_id in REQUIRED_CONTROL_IDS)
    assert matrix["closed"] is True
    assert matrix["summary"]["gap_asset_count"] == 0


def test_outdated_live_receipts_are_classified_stale_per_stage() -> None:
    now = datetime(2026, 7, 11, 12, 0, tzinfo=UTC)
    old = now - timedelta(hours=40)
    live = _fully_covered_host(old)

    matrix = build_asset_capability_matrix(
        _single_host_scope(),
        live_assets=[live],
        latest_run={"run_id": "run-old", "status": "success", "ended_at": old},
        repo_root=_REPO_ROOT,
        generated_at=now,
    )

    asset = matrix["assets"][0]
    assert asset["stages"]["runtime_identity"]["status"] == "stale"
    assert asset["stages"]["sensor"]["status"] == "stale"
    assert asset["stages"]["policy"]["status"] == "stale"
    assert asset["stages"]["executor"]["status"] == "stale"
    assert asset["stages"]["verifier"]["status"] == "stale"
    assert asset["stages"]["backup_restore"]["status"] == "stale"
    assert asset["stages"]["learning_writeback"]["status"] == "stale"
    assert asset["stale_stage_ids"] == [
        "runtime_identity",
        "sensor",
        "policy",
        "executor",
        "verifier",
        "backup_restore",
        "learning_writeback",
    ]


def test_reconciliation_candidates_are_stable_and_public_safe() -> None:
    matrix = build_asset_capability_matrix(
        _single_host_scope(),
        repo_root=_REPO_ROOT,
        generated_at=datetime(2026, 7, 11, 12, 0, tzinfo=UTC),
    )

    first = build_reconciliation_candidates(matrix)
    second = build_reconciliation_candidates(matrix)

    assert first == second
    assert len(first) == 1
    assert first[0]["work_item_id"].startswith("AIA-P0-006-WI-")
    assert first[0]["candidate_kind"] == "capability_gap"
    assert first[0]["controlled_apply_allowed"] is True
    assert "asset_key" not in first[0]
    assert "host" not in first[0]


def test_reconciliation_summary_requires_repository_readback() -> None:
    matrix = build_asset_capability_matrix(
        _single_host_scope(),
        latest_run={"run_id": "discovery-run-1", "status": "success"},
        repo_root=_REPO_ROOT,
    )
    candidates = build_reconciliation_candidates(matrix)

    failed = build_reconciliation_summary_payload(
        matrix,
        candidates,
        trace_id="trace-1",
        automation_run_id="run-1",
        work_item_id="AIA-P0-006",
        repository_readback_count=0,
        created_count=1,
        existing_count=0,
    )
    passed = build_reconciliation_summary_payload(
        matrix,
        candidates,
        trace_id="trace-1",
        automation_run_id="run-1",
        work_item_id="AIA-P0-006",
        repository_readback_count=1,
        created_count=1,
        existing_count=0,
    )

    assert failed["repository_readback_verified"] is False
    assert failed["closed"] is False
    assert passed["repository_readback_verified"] is True
    assert passed["closed"] is True


def test_reconciliation_work_item_uses_same_trace_and_run_as_summary() -> None:
    matrix = build_asset_capability_matrix(
        _single_host_scope(),
        repo_root=_REPO_ROOT,
    )
    candidate = build_reconciliation_candidates(matrix)[0]

    receipt_input = build_reconciliation_work_item_input(
        candidate,
        trace_id="trace-1",
        automation_run_id="reconciliation-run-1",
        program_work_item_id="AIA-P0-006",
        matrix_fingerprint=matrix["matrix_fingerprint"],
    )

    assert receipt_input["trace_id"] == "trace-1"
    assert receipt_input["automation_run_id"] == "reconciliation-run-1"
    assert receipt_input["program_work_item_id"] == "AIA-P0-006"
    assert (
        receipt_input["source_diff"]["matrix_fingerprint"]
        == (matrix["matrix_fingerprint"])
    )


def test_priority_overlay_projects_aia_p0_006_runtime_controls() -> None:
    matrix = build_asset_capability_matrix(
        _single_host_scope(),
        repo_root=_REPO_ROOT,
    )
    payload: dict[str, object] = {"summary": {}, "rollups": {}}

    apply_ai_automation_asset_capability_matrix(payload, matrix)

    assert payload["ai_automation_asset_capability_matrix"]["schema_version"] == (
        "ai_automation_asset_capability_matrix_v1"
    )
    assert payload["ai_automation_asset_capability_matrix"]["assets"] == []
    assert payload["ai_automation_asset_capability_matrix"]["assets_embedded"] is False
    assert payload["ai_automation_asset_capability_matrix"]["assets_endpoint"] == (
        "/api/v1/agents/ai-automation-asset-capability-matrix"
    )
    assert (
        payload["summary"]["ai_automation_asset_capability_declared_asset_count"] == 1
    )
    work_item = next(
        item
        for item in payload["ai_automation_program_ledger"]["work_items"]
        if item["id"] == "AIA-P0-006"
    )
    assert work_item["status"] == "in_progress"
    assert work_item["runtime_progress"]["required_control_count"] == len(
        REQUIRED_CONTROL_IDS
    )


class _FakeMappings:
    def __init__(self, rows: list[dict[str, Any]]) -> None:
        self._rows = rows

    def first(self) -> dict[str, Any] | None:
        return self._rows[0] if self._rows else None

    def all(self) -> list[dict[str, Any]]:
        return self._rows


class _FakeResult:
    def __init__(self, rows: list[dict[str, Any]]) -> None:
        self._rows = rows

    def mappings(self) -> _FakeMappings:
        return _FakeMappings(self._rows)


class _FakeLiveReadbackDb:
    def __init__(self, now: datetime) -> None:
        self.now = now
        self.calls: list[tuple[str, dict[str, Any]]] = []

    async def execute(
        self,
        statement: Any,
        params: dict[str, Any] | None = None,
    ) -> _FakeResult:
        sql = str(statement)
        bound = params or {}
        self.calls.append((sql, bound))
        if "set_config('statement_timeout'" in sql:
            return _FakeResult([])
        if "FROM asset_discovery_run" in sql:
            return _FakeResult(
                [
                    {
                        "run_id": "00000000-0000-0000-0000-000000000001",
                        "status": "success",
                        "ended_at": self.now,
                    }
                ]
            )
        if "FROM asset_inventory" in sql:
            return _FakeResult(
                [
                    {
                        "asset_id": 42,
                        "asset_key": "host:99",
                        "asset_type": "host",
                        "environment": "prod",
                        "namespace": "",
                        "name": "99",
                        "owner_team": "platform_sre",
                        "source_repo": "wooo/awoooi",
                        "source_commit_sha": "abc123",
                        "lifecycle_state": "active",
                        "last_seen_at": self.now,
                        "dimension": "auto_monitoring",
                        "coverage_status": "green",
                        "coverage_created_at": self.now,
                    }
                ]
            )
        if "FROM asset_compliance_snapshot" in sql:
            return _FakeResult([])
        if "FROM asset_change_event" in sql:
            return _FakeResult([{"asset_id": 42, "change_type": "asset_modified"}])
        if "FROM automation_operation_log" in sql:
            return _FakeResult([])
        raise AssertionError(f"unexpected SQL: {sql}")


class _FakeDbContext:
    def __init__(self, db: _FakeLiveReadbackDb) -> None:
        self.db = db

    async def __aenter__(self) -> _FakeLiveReadbackDb:
        return self.db

    async def __aexit__(self, *_args: object) -> None:
        return None


def test_live_change_event_readback_is_scoped_to_selected_asset_ids(
    monkeypatch: Any,
) -> None:
    now = datetime(2026, 7, 15, 1, 0, tzinfo=UTC)
    db = _FakeLiveReadbackDb(now)

    monkeypatch.setattr(
        "src.db.base.get_db_context",
        lambda _project_id: _FakeDbContext(db),
    )

    _latest_run, assets, _receipt = asyncio.run(_load_live_rows("awoooi"))

    timeout_sql, timeout_params = next(
        (sql, params)
        for sql, params in db.calls
        if "set_config('statement_timeout'" in sql
    )
    assert "statement_timeout" in timeout_sql
    assert timeout_params == {"statement_timeout": "6000ms"}
    change_sql, change_params = next(
        (sql, params) for sql, params in db.calls if "FROM asset_change_event" in sql
    )
    assert "asset_id = ANY(:asset_ids)" in change_sql
    assert change_params == {"asset_ids": [42]}
    assert assets[0]["change_types"] == ["asset_modified"]


def test_live_readback_timeout_fails_closed_without_hanging(
    monkeypatch: Any,
) -> None:
    async def _slow_live_readback(
        _project_id: str,
    ) -> tuple[dict[str, Any], list[dict[str, Any]], dict[str, Any]]:
        await asyncio.sleep(1)
        return {}, [], {}

    monkeypatch.setattr(matrix_service, "_load_live_rows", _slow_live_readback)

    payload = asyncio.run(
        build_asset_capability_matrix_with_live_readback(
            repo_root=_REPO_ROOT,
            timeout_seconds=0.01,
        )
    )

    assert payload["live_readback_status"] == "degraded"
    assert payload["live_readback_error_type"] == "TimeoutError"
    assert payload["live_readback_timeout_seconds"] == 0.01
    assert payload["summary"]["declared_asset_count"] == 193
    assert payload["operation_boundaries"]["host_write_performed"] is False


def test_matrix_build_phase_is_inside_the_total_timeout(monkeypatch: Any) -> None:
    async def _fast_live_readback(
        _project_id: str,
    ) -> tuple[dict[str, Any], list[dict[str, Any]], dict[str, Any]]:
        return {}, [], {}

    async def _slow_to_thread(*_args: Any, **_kwargs: Any) -> dict[str, Any]:
        await asyncio.sleep(1)
        return {}

    monkeypatch.setattr(matrix_service, "_load_live_rows", _fast_live_readback)
    monkeypatch.setattr(matrix_service.asyncio, "to_thread", _slow_to_thread)

    payload = asyncio.run(
        build_asset_capability_matrix_with_live_readback(
            repo_root=_REPO_ROOT,
            timeout_seconds=0.01,
        )
    )

    assert payload["live_readback_status"] == "degraded"
    assert payload["live_readback_error_type"] == "TimeoutError"
    assert payload["live_readback_timeout_seconds"] == 0.01