Newer
Older
navi-1 / tests / unit / llm / test_health.py
"""Target health: why it failed, when to retry, and how the probe switches back.

The old engine's blacklist was one-way — nothing ever re-checked a dead host, so
recovery meant a restart. These tests pin the two halves of the replacement:
``report_failure`` records a *reason* with a retry time, and ``probe_once`` is
what brings a target back.
"""

import asyncio
import time

import pytest

from navi.exceptions import (
    LLMAuthError,
    LLMBackendError,
    LLMConnectionError,
    LLMModelNotFoundError,
    LLMQuotaExceededError,
    LLMServerUnavailableError,
)
from navi.llm.config import LLMConfig, Route, RouteEntry, Target
from navi.llm.health import ModelHealth, RouteHealth


def _config(interval=60.0, timeout=5.0, chain=None):
    cloud = Target(name="cloud", host="https://cloud.test")
    home = Target(name="home", host="http://home.test")
    chain = chain or [(("glm", cloud)), ("qwen", home), ("qwen", cloud)]
    routes = {"default": Route("default", tuple(RouteEntry(m, t) for m, t in chain))}
    return LLMConfig(
        targets={"cloud": cloud, "home": home},
        routes=routes,
        probe_interval_sec=interval,
        probe_timeout_sec=timeout,
    )


def _entry(config, model, target):
    return next(
        e
        for e in config.routes["default"].entries
        if e.model == model and e.target.name == target
    )


# ── recorded failures ────────────────────────────────────────────────────────

def test_a_connection_failure_takes_the_target_down_and_opens_a_retry_window():
    config = _config()
    health = RouteHealth(config)
    entry = _entry(config, "glm", "cloud")

    health.report_failure(entry, LLMConnectionError("refused"))

    state = health.snapshot()["targets"]["cloud"]
    assert state["state"] == "down"
    assert state["reason"] == "network"
    assert state["retry_in_sec"] > 0
    assert health.entry_usable(entry) is False


def test_a_target_becomes_usable_again_once_its_retry_time_passes():
    config = _config(interval=10)
    health = RouteHealth(config)
    entry = _entry(config, "glm", "cloud")

    health.report_failure(entry, LLMConnectionError("refused"))

    assert health.target_usable(entry.target, now=time.monotonic() + 11) is True


def test_quota_cools_the_target_and_honours_retry_after():
    config = _config(interval=10)
    health = RouteHealth(config)
    entry = _entry(config, "glm", "cloud")

    health.report_failure(entry, LLMQuotaExceededError("limit", retry_after=120))

    state = health._target_state(entry.target)
    assert (state.state, state.reason) == ("cooling", "quota")
    assert state.next_probe - time.monotonic() > 100


def test_quota_never_opens_a_window_shorter_than_the_probe_interval():
    config = _config(interval=60)
    health = RouteHealth(config)
    entry = _entry(config, "glm", "cloud")

    health.report_failure(entry, LLMQuotaExceededError("limit", retry_after=5))

    state = health._target_state(entry.target)
    assert state.next_probe - time.monotonic() > 50


def test_auth_is_rechecked_slowly_because_a_key_can_be_rotated():
    config = _config(interval=10)
    health = RouteHealth(config)
    entry = _entry(config, "glm", "cloud")

    health.report_failure(entry, LLMAuthError("401"))

    state = health._target_state(entry.target)
    assert (state.state, state.reason) == ("auth", "auth")
    assert state.next_probe - time.monotonic() > 200  # 300s, not the 10s interval


def test_a_5xx_cools_the_target_without_calling_it_down():
    config = _config()
    health = RouteHealth(config)
    entry = _entry(config, "glm", "cloud")

    health.report_failure(entry, LLMServerUnavailableError("502"))

    state = health._target_state(entry.target)
    assert (state.state, state.reason) == ("cooling", "server")


def test_a_generic_backend_error_does_not_blacklist_the_host():
    """A malformed request would fail on every entry — no point in failing over."""
    config = _config()
    health = RouteHealth(config)
    entry = _entry(config, "glm", "cloud")

    health.report_failure(entry, LLMBackendError("bad request"))

    assert health.snapshot()["targets"]["cloud"]["state"] == "ok"
    assert health.entry_usable(entry) is True


def test_a_missing_model_degrades_only_that_entry():
    config = _config()
    health = RouteHealth(config)
    missing = _entry(config, "glm", "cloud")
    sibling = _entry(config, "qwen", "cloud")

    health.report_failure(missing, LLMModelNotFoundError("404"))

    assert health.snapshot()["targets"]["cloud"]["state"] == "ok"
    assert health.entry_usable(missing) is False
    assert health.entry_usable(sibling) is True
    assert health.snapshot()["missing_models"] == [
        {"target": "cloud", "model": "glm", "detail": "404"}
    ]


def test_a_missing_model_is_retried_after_the_interval():
    config = _config(interval=10)
    health = RouteHealth(config)
    missing = _entry(config, "glm", "cloud")
    health.report_failure(missing, LLMModelNotFoundError("404"))

    assert health.entry_usable(missing, now=time.monotonic() + 11) is True


def test_success_clears_the_failure_and_announces_recovery():
    config = _config()
    health = RouteHealth(config)
    entry = _entry(config, "glm", "cloud")
    changes = []
    health.on_change = lambda target, state: changes.append((target, state))

    health.report_failure(entry, LLMConnectionError("refused"))
    health.report_success(entry)

    assert health.snapshot()["targets"]["cloud"]["state"] == "ok"
    assert changes == [("cloud", "down"), ("cloud", "ok")]


def test_success_clears_a_missing_model():
    config = _config()
    health = RouteHealth(config)
    entry = _entry(config, "glm", "cloud")
    health.report_failure(entry, LLMModelNotFoundError("404"))

    health.report_success(entry)

    assert health.snapshot()["missing_models"] == []


def test_reset_forgets_everything():
    config = _config()
    health = RouteHealth(config)
    health.report_failure(_entry(config, "glm", "cloud"), LLMConnectionError("x"))
    health.report_failure(_entry(config, "qwen", "home"), LLMModelNotFoundError("404"))

    health.reset()

    snapshot = health.snapshot()
    assert all(t["state"] == "ok" for t in snapshot["targets"].values())
    assert snapshot["missing_models"] == []


def test_a_failing_listener_never_breaks_the_llm_path():
    config = _config()
    health = RouteHealth(config)
    health.on_change = lambda *_: (_ for _ in ()).throw(RuntimeError("boom"))

    health.report_failure(_entry(config, "glm", "cloud"), LLMConnectionError("x"))

    assert health.snapshot()["targets"]["cloud"]["state"] == "down"


# ── the background probe ────────────────────────────────────────────────────

class FakeClient:
    def __init__(self, list_raises=None, chat_raises=None, show_raises=None):
        self.list_raises = list_raises
        self.chat_raises = chat_raises
        self.show_raises = show_raises or {}
        self.calls = []

    async def list(self):
        self.calls.append("list")
        if self.list_raises:
            raise self.list_raises
        return {"models": []}

    async def show(self, model):
        self.calls.append(("show", model))
        if model in self.show_raises:
            raise self.show_raises[model]
        return {}

    async def chat(self, **kwargs):
        self.calls.append(("chat", kwargs.get("model")))
        if self.chat_raises:
            raise self.chat_raises
        return {}


def _health_with_clients(config, **clients):
    health = RouteHealth(config)
    health._clients.update(clients)
    return health


def test_the_probe_does_nothing_while_every_target_is_healthy():
    config = _config()
    client = FakeClient()
    health = _health_with_clients(config, cloud=client, home=FakeClient())

    asyncio.run(health.probe_once())

    assert client.calls == []


def test_the_probe_brings_a_down_target_back():
    config = _config(interval=0)
    client = FakeClient()
    health = _health_with_clients(config, cloud=client)
    entry = _entry(config, "glm", "cloud")
    health.report_failure(entry, LLMConnectionError("refused"))

    asyncio.run(health.probe_once())

    assert health.snapshot()["targets"]["cloud"]["state"] == "ok"
    assert health.entry_usable(entry) is True


def test_the_probe_does_not_run_before_the_retry_window_closes():
    config = _config(interval=600)
    client = FakeClient()
    health = _health_with_clients(config, cloud=client)
    health.report_failure(_entry(config, "glm", "cloud"), LLMConnectionError("refused"))

    asyncio.run(health.probe_once())

    assert client.calls == []


def test_a_probe_that_still_cannot_reach_the_target_keeps_it_down():
    config = _config(interval=0)
    client = FakeClient(list_raises=ConnectionError("refused"))
    health = _health_with_clients(config, cloud=client)
    health.report_failure(_entry(config, "glm", "cloud"), LLMConnectionError("refused"))

    asyncio.run(health.probe_once())

    state = health.snapshot()["targets"]["cloud"]
    assert state["state"] == "down"
    assert state["reason"] == "network"


def test_a_probe_that_times_out_marks_the_target_down():
    class SlowClient(FakeClient):
        async def list(self):
            await asyncio.sleep(10)

    config = _config(interval=0, timeout=0.01)
    health = _health_with_clients(config, cloud=SlowClient())
    health.report_failure(_entry(config, "glm", "cloud"), LLMConnectionError("refused"))

    asyncio.run(health.probe_once())

    assert health.snapshot()["targets"]["cloud"]["detail"].startswith("probe timed out")


def test_a_probe_that_gets_a_401_keeps_the_target_in_auth():
    config = _config(interval=0)

    class AuthClient(FakeClient):
        async def list(self):
            raise type("ResponseError", (Exception,), {"status_code": 401})("unauthorized")

    health = _health_with_clients(config, cloud=AuthClient())
    health.report_failure(_entry(config, "glm", "cloud"), LLMAuthError("401"))

    asyncio.run(health.probe_once())

    assert health.snapshot()["targets"]["cloud"]["state"] == "auth"


def test_a_cooling_target_is_checked_with_a_real_generation():
    """``list()`` answers while generation is blocked, so quota needs a chat call."""
    config = _config(interval=0)
    client = FakeClient(chat_raises=LLMQuotaExceededError("still limited"))
    health = _health_with_clients(config, cloud=client)
    health.report_failure(_entry(config, "glm", "cloud"), LLMQuotaExceededError("limit"))

    asyncio.run(health.probe_once())

    state = health.snapshot()["targets"]["cloud"]
    assert state["state"] == "cooling"
    assert ("chat", "glm") in client.calls


def test_a_cooling_target_returns_once_the_generation_succeeds():
    config = _config(interval=0)
    health = _health_with_clients(config, cloud=FakeClient())
    health.report_failure(_entry(config, "glm", "cloud"), LLMQuotaExceededError("limit"))

    asyncio.run(health.probe_once())

    assert health.snapshot()["targets"]["cloud"]["state"] == "ok"


def test_a_missing_model_is_rechecked_with_show_even_on_a_healthy_target():
    config = _config(interval=0)
    health = _health_with_clients(config, cloud=FakeClient())
    entry = _entry(config, "glm", "cloud")
    health.report_failure(entry, LLMModelNotFoundError("404"))

    asyncio.run(health.probe_once())

    assert health.snapshot()["missing_models"] == []
    assert health.entry_usable(entry) is True


def test_a_model_still_missing_keeps_its_entry_unusable():
    config = _config(interval=600)
    client = FakeClient(show_raises={"glm": LLMModelNotFoundError("404")})
    health = _health_with_clients(config, cloud=client)
    entry = _entry(config, "glm", "cloud")
    # Already due for a re-check, so the probe will ask this time.
    health._models[("cloud", "glm")] = ModelHealth(missing=True, next_probe=0.0)

    asyncio.run(health.probe_once())

    assert health.entry_usable(entry) is False
    assert health.snapshot()["targets"]["cloud"]["state"] == "ok"


def test_the_probe_loop_stops_on_cancel():
    config = _config(interval=3600)
    health = RouteHealth(config)

    async def run():
        health.start()
        task = health._task
        assert task is not None
        await asyncio.sleep(0)
        await health.stop()
        assert task.cancelled() or task.done()

    asyncio.run(run())
    assert health._task is None


# ── probe error classification ──────────────────────────────────────────────

@pytest.mark.parametrize(
    "status, state, reason",
    [(401, "auth", "auth"), (403, "auth", "auth"), (429, "cooling", "quota"),
     (500, "cooling", "server"), (503, "cooling", "server")],
)
def test_probe_errors_are_classified_by_status(status, state, reason):
    from navi.llm.health import _classify_probe_error

    err = type("ResponseError", (Exception,), {"status_code": status})("nope")

    assert _classify_probe_error(err) == (state, reason)


def test_a_network_error_is_a_down_target():
    from navi.llm.health import _classify_probe_error

    assert _classify_probe_error(ConnectionError("refused")) == ("down", "network")