legal-ai/web/tests/test_agent_health.py

"""Tests for #222 — agent health taxonomy (pure classifier + Paperclip fetch).

Pure-classifier tests need no deps. The fetch test fakes an asyncpg connection
(routing the three queries by content) so it never touches the live DB.
"""

from __future__ import annotations

import asyncio
import os
import sys
from pathlib import Path

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))  # web/

from agent_health import (  # noqa: E402
    STALL_WAKEUP_THRESHOLD,
    classify_issue_health,
    is_recovery_reason,
    zombie_escalation,
)
from agent_health import ZOMBIE_ESCALATE_MIN_RECOVERY  # noqa: E402


# ── pure classifier ────────────────────────────────────────────────────────
def test_live_run_is_working_regardless_of_wakeups():
    assert classify_issue_health(has_live_run=True, recovery_wakeups=9, total_wakeups=9) == "working"


def test_recovery_wakeup_without_live_run_is_zombie():
    # the stranded-child signature: assigned+open, nothing running, recovery firing
    assert classify_issue_health(has_live_run=False, recovery_wakeups=1, total_wakeups=1) == "zombie"


def test_repeated_wakeups_no_recovery_is_stalled():
    assert (
        classify_issue_health(
            has_live_run=False, recovery_wakeups=0, total_wakeups=STALL_WAKEUP_THRESHOLD
        )
        == "stalled"
    )


def test_quiet_open_issue_is_idle():
    assert classify_issue_health(has_live_run=False, recovery_wakeups=0, total_wakeups=1) == "idle"


# ── watchdog escalation policy (pure) ───────────────────────────────────────
def _item(health, recovery):
    return {"health": health, "recovery_wakeups": recovery, "identifier": "CMP-141", "agent_name": "writer"}


def test_zombie_escalation_medium_for_fresh_loop():
    decision = zombie_escalation(_item("zombie", ZOMBIE_ESCALATE_MIN_RECOVERY))
    assert decision is not None
    severity, reason = decision
    assert severity == "medium"
    assert "CMP-141" in reason and "recovery" in reason.lower()


def test_zombie_escalation_high_for_long_loop():
    severity, _ = zombie_escalation(_item("zombie", 5))
    assert severity == "high"


def test_zombie_below_threshold_is_not_escalated():
    assert zombie_escalation(_item("zombie", ZOMBIE_ESCALATE_MIN_RECOVERY - 1)) is None


def test_non_zombie_never_escalates():
    assert zombie_escalation(_item("stalled", 9)) is None
    assert zombie_escalation(_item("idle", 0)) is None
    assert zombie_escalation(_item("working", 9)) is None


def test_is_recovery_reason_markers():
    assert is_recovery_reason("source_scoped_recovery_action")
    assert is_recovery_reason("issue_continuation_needed")
    assert is_recovery_reason("issue_reopened_via_comment")
    assert not is_recovery_reason("start_workflow_8125-09-24")
    assert not is_recovery_reason(None)


# ── Paperclip fetch (fake DB) ───────────────────────────────────────────────
os.environ.setdefault("PAPERCLIP_DB_URL", "postgres://x:x@127.0.0.1:54329/paperclip")
pc = pytest.importorskip("paperclip_client", reason="web deps unavailable")


class _FakeConn:
    def __init__(self, issues, live, wakeups):
        self._issues, self._live, self._wakeups = issues, live, wakeups

    async def fetch(self, query, *args):
        if "FROM issues" in query:
            return self._issues
        if "heartbeat_runs" in query:
            return self._live
        if "agent_wakeup_requests" in query:
            return self._wakeups
        raise AssertionError(f"unexpected query: {query[:40]}")

    async def close(self):
        pass


def test_get_agent_health_classifies_and_sorts(monkeypatch):
    # zombie agent A (recovery wake, no live run); working agent B (live run);
    # idle agent C (quiet).
    issues = [
        {"id": "iA", "identifier": "CMP-1", "status": "blocked", "assignee_agent_id": "agA", "agent_name": "writer"},
        {"id": "iB", "identifier": "CMP-2", "status": "in_progress", "assignee_agent_id": "agB", "agent_name": "analyst"},
        {"id": "iC", "identifier": "CMP-3", "status": "todo", "assignee_agent_id": "agC", "agent_name": "qa"},
    ]
    live = [{"agent_id": "agB"}]
    wakeups = [
        {"agent_id": "agA", "issue_id": "iA", "reason": "source_scoped_recovery_action"},
        {"agent_id": "agC", "issue_id": "iC", "reason": "start_workflow_x"},
    ]

    async def _connect(_url):
        return _FakeConn(issues, live, wakeups)

    monkeypatch.setattr(pc.asyncpg, "connect", _connect)

    result = asyncio.run(pc.get_agent_health())
    assert result["ok"] is True
    by_issue = {it["issue_id"]: it["health"] for it in result["items"]}
    assert by_issue == {"iA": "zombie", "iB": "working", "iC": "idle"}
    assert result["counts"]["zombie"] == 1
    assert result["counts"]["working"] == 1
    assert result["counts"]["idle"] == 1
    # worst-first: zombie leads
    assert result["items"][0]["health"] == "zombie"
    assert result["items"][0]["recovery_wakeups"] == 1


def test_get_agent_health_empty(monkeypatch):
    async def _connect(_url):
        return _FakeConn([], [], [])

    monkeypatch.setattr(pc.asyncpg, "connect", _connect)
    result = asyncio.run(pc.get_agent_health())
    assert result == {"ok": True, "items": [], "counts": {"zombie": 0, "stalled": 0, "working": 0, "idle": 0}}