awoooi/apps/api/tests/test_health_ollama_provider_chain.py

from __future__ import annotations

import pytest

from src.api.v1 import health
from src.services.ollama_endpoint_circuit_breaker import (
    is_ollama_endpoint_blocked,
    record_ollama_endpoint_failure,
    reset_ollama_endpoint_cooldown_for_tests,
)


def setup_function() -> None:
    reset_ollama_endpoint_cooldown_for_tests()


def teardown_function() -> None:
    reset_ollama_endpoint_cooldown_for_tests()


def _set_ollama_settings(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(health.settings, "OLLAMA_URL", "http://gcp-a:11434")
    monkeypatch.setattr(health.settings, "OLLAMA_SECONDARY_URL", "http://gcp-b:11434")
    monkeypatch.setattr(health.settings, "OLLAMA_FALLBACK_URL", "http://local-111:11434")


@pytest.mark.asyncio
async def test_ollama_provider_chain_reports_fallback_when_primary_down(monkeypatch: pytest.MonkeyPatch) -> None:
    _set_ollama_settings(monkeypatch)

    async def fake_http_check(name: str, _url: str, _path: str) -> health.ComponentHealth:
        if name == "ollama_gcp_a":
            return health.ComponentHealth(status="down", error="timeout")
        if name == "ollama_gcp_b":
            return health.ComponentHealth(status="up", latency_ms=42.0)
        return health.ComponentHealth(status="up", latency_ms=9.0)

    monkeypatch.setattr(health, "_http_health_check", fake_http_check)

    aggregate, details = await health.check_ollama_provider_chain()

    assert aggregate.status == "degraded"
    assert aggregate.latency_ms == 42.0
    assert aggregate.error == "primary unavailable; fallback active: ollama_gcp_b"
    assert details["ollama_gcp_a"].status == "down"
    assert details["ollama_gcp_a"].provider_name == "ollama_gcp_a"
    assert details["ollama_gcp_a"].diagnosis_code == "endpoint_timeout"
    assert details["ollama_gcp_b"].status == "up"
    assert details["ollama_gcp_b"].diagnosis_code == "endpoint_reachable"
    assert details["ollama_local"].status == "up"


@pytest.mark.asyncio
async def test_ollama_provider_chain_reports_all_endpoints_when_down(monkeypatch: pytest.MonkeyPatch) -> None:
    _set_ollama_settings(monkeypatch)

    async def fake_http_check(name: str, _url: str, _path: str) -> health.ComponentHealth:
        return health.ComponentHealth(status="down", error=f"{name} timeout")

    monkeypatch.setattr(health, "_http_health_check", fake_http_check)

    aggregate, details = await health.check_ollama_provider_chain()

    assert aggregate.status == "down"
    assert "ollama_gcp_a=ollama_gcp_a timeout" in (aggregate.error or "")
    assert "ollama_gcp_b=ollama_gcp_b timeout" in (aggregate.error or "")
    assert "ollama_local=ollama_local timeout" in (aggregate.error or "")
    assert set(details) == {"ollama_gcp_a", "ollama_gcp_b", "ollama_local"}


def test_overall_status_uses_aggregate_ollama_not_endpoint_details() -> None:
    components = {
        "api": health.ComponentHealth(status="up"),
        "postgresql": health.ComponentHealth(status="up"),
        "redis": health.ComponentHealth(status="up"),
        "ollama": health.ComponentHealth(status="degraded"),
        "openclaw": health.ComponentHealth(status="up"),
        "signoz": health.ComponentHealth(status="up"),
        "ollama_gcp_a": health.ComponentHealth(status="down"),
        "ollama_gcp_b": health.ComponentHealth(status="up"),
        "ollama_local": health.ComponentHealth(status="up"),
    }

    assert health._determine_overall_status(components) == "degraded"


@pytest.mark.asyncio
async def test_ollama_provider_chain_uses_cooldown_after_failure(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _set_ollama_settings(monkeypatch)
    calls: list[str] = []

    async def fake_http_check(name: str, _url: str, _path: str) -> health.ComponentHealth:
        calls.append(name)
        if name == "ollama_gcp_a":
            return health.ComponentHealth(status="down", error="http 502")
        return health.ComponentHealth(status="up", latency_ms=9.0)

    monkeypatch.setattr(health, "_http_health_check", fake_http_check)

    aggregate, details = await health.check_ollama_provider_chain()

    assert aggregate.status == "degraded"
    assert details["ollama_gcp_a"].status == "down"
    assert is_ollama_endpoint_blocked("http://gcp-a:11434")

    calls.clear()
    aggregate, details = await health.check_ollama_provider_chain()

    assert aggregate.status == "degraded"
    assert details["ollama_gcp_a"].status == "down"
    assert "cooldown" in (details["ollama_gcp_a"].error or "")
    assert details["ollama_gcp_a"].diagnosis_code == "endpoint_cooldown"
    assert details["ollama_gcp_a"].is_cooldown is True
    assert details["ollama_gcp_a"].retry_after_seconds is not None
    assert "ollama_gcp_a" not in calls
    assert {"ollama_gcp_b", "ollama_local"} == set(calls)


@pytest.mark.asyncio
async def test_ollama_provider_chain_success_clears_cooldown(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _set_ollama_settings(monkeypatch)
    record_ollama_endpoint_failure("http://gcp-a:11434")
    monkeypatch.setattr(
        health,
        "get_ollama_endpoint_cooldown_remaining_seconds",
        lambda _url: 0.0,
    )

    async def fake_http_check(name: str, _url: str, _path: str) -> health.ComponentHealth:
        return health.ComponentHealth(status="up", latency_ms=8.0)

    monkeypatch.setattr(health, "_http_health_check", fake_http_check)

    await health.check_ollama_provider_chain()

    assert not is_ollama_endpoint_blocked("http://gcp-a:11434")


def test_ollama_failure_classifier_marks_local_proxy_502() -> None:
    assert (
        health._classify_ollama_endpoint_failure(
            "ollama_local",
            "Server error '502 Bad Gateway' for url",
        )
        == "local_proxy_upstream_unreachable"
    )


def test_ollama_failure_classifier_marks_network_unreachable() -> None:
    assert (
        health._classify_ollama_endpoint_failure(
            "ollama_local",
            "No route to host",
        )
        == "endpoint_network_unreachable"
    )