homelab-codex-ws/services/kb-query/tests/test_fallback.py

"""Unit tests for the embed fallback state machine (app/fallback.py) -- module 5 phase 4 plan
§2 decision 2 / §5. No real HTTP, no real Ollama -- same mocking style as
packages/kb-retrieval/tests/test_embed.py."""
from __future__ import annotations

import pathlib
import sys

import aiohttp
import pytest

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

from app.fallback import (  # noqa: E402
    EMBED_TIMEOUT_S,
    SolCircuitBreaker,
    embed_with_fallback,
    resolve_sol_status,
)

SOLARIA_URL = "http://solaria:11434"
PIHA_URL = "http://piha:11434"


class _FakeClock:
    def __init__(self, start: float = 0.0):
        self.t = start

    def __call__(self) -> float:
        return self.t

    def advance(self, dt: float) -> None:
        self.t += dt


class _FakeGetResp:
    def __init__(self, status: int):
        self.status = status

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc):
        return False


class _FakePostResp:
    def __init__(self, embedding=None):
        self._embedding = embedding if embedding is not None else [0.01] * 1024

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc):
        return False

    def raise_for_status(self):
        pass

    async def json(self):
        return {"embedding": self._embedding}


class _FakeFallbackSession:
    """Routes GET -> health probe, POST -> embed. `solaria_get_status=None` simulates the probe
    itself being unreachable (raises); `*_post_behavior` in {"ok", "timeout", "error"}."""

    def __init__(self, solaria_get_status=200, solaria_post_behavior="ok", piha_post_behavior="ok"):
        self.solaria_get_status = solaria_get_status
        self.solaria_post_behavior = solaria_post_behavior
        self.piha_post_behavior = piha_post_behavior
        self.get_calls: list[str] = []
        self.post_calls: list[dict] = []

    def get(self, url, timeout=None):
        self.get_calls.append(url)
        if self.solaria_get_status is None:
            raise aiohttp.ClientConnectionError("refused")
        return _FakeGetResp(self.solaria_get_status)

    def post(self, url, json, timeout=None):
        self.post_calls.append({"url": url, "json": json, "timeout": timeout})
        behavior = self.solaria_post_behavior if url.startswith(SOLARIA_URL) else self.piha_post_behavior
        if behavior == "timeout":
            raise TimeoutError()
        if behavior == "error":
            raise aiohttp.ClientConnectionError("refused")
        return _FakePostResp()


class TestSolCircuitBreaker:
    def test_status_none_when_never_set(self):
        assert SolCircuitBreaker().status is None

    def test_status_returns_cached_value_within_ttl(self):
        clock = _FakeClock()
        breaker = SolCircuitBreaker(cache_ttl_s=30, clock=clock)
        breaker.set("up")
        clock.advance(29)
        assert breaker.status == "up"

    def test_status_expires_exactly_at_ttl(self):
        clock = _FakeClock()
        breaker = SolCircuitBreaker(cache_ttl_s=30, clock=clock)
        breaker.set("down")
        clock.advance(30)
        assert breaker.status is None


class TestResolveSolStatus:
    async def test_uses_fresh_cache_without_probing(self):
        breaker = SolCircuitBreaker()
        breaker.set("up")
        session = _FakeFallbackSession()
        status = await resolve_sol_status(breaker, session, SOLARIA_URL)
        assert status == "up"
        assert session.get_calls == []

    async def test_probes_and_caches_up(self):
        breaker = SolCircuitBreaker()
        session = _FakeFallbackSession(solaria_get_status=200)
        status = await resolve_sol_status(breaker, session, SOLARIA_URL)
        assert status == "up"
        assert breaker.status == "up"
        assert session.get_calls == [f"{SOLARIA_URL}/api/tags"]

    async def test_probes_and_caches_down_on_unreachable(self):
        breaker = SolCircuitBreaker()
        session = _FakeFallbackSession(solaria_get_status=None)
        status = await resolve_sol_status(breaker, session, SOLARIA_URL)
        assert status == "down"
        assert breaker.status == "down"

    async def test_reprobes_once_ttl_expires(self):
        clock = _FakeClock()
        breaker = SolCircuitBreaker(cache_ttl_s=30, clock=clock)
        session = _FakeFallbackSession(solaria_get_status=200)
        await resolve_sol_status(breaker, session, SOLARIA_URL)
        clock.advance(30)
        await resolve_sol_status(breaker, session, SOLARIA_URL)
        assert len(session.get_calls) == 2


class TestEmbedWithFallback:
    async def test_solaria_up_embeds_on_solaria(self):
        breaker = SolCircuitBreaker()
        session = _FakeFallbackSession(solaria_get_status=200, solaria_post_behavior="ok")
        embedding, status = await embed_with_fallback(
            breaker, session, SOLARIA_URL, PIHA_URL, "bge-m3", "q"
        )
        assert status == "up"
        assert len(embedding) == 1024
        assert session.post_calls == [
            {"url": f"{SOLARIA_URL}/api/embeddings", "json": {"model": "bge-m3", "prompt": "q"},
             "timeout": aiohttp.ClientTimeout(total=EMBED_TIMEOUT_S)}
        ]

    async def test_cached_down_skips_probe_and_solaria_entirely(self):
        breaker = SolCircuitBreaker()
        breaker.set("down")
        session = _FakeFallbackSession(piha_post_behavior="ok")
        embedding, status = await embed_with_fallback(
            breaker, session, SOLARIA_URL, PIHA_URL, "bge-m3", "q"
        )
        assert status == "down"
        assert len(embedding) == 1024
        assert session.get_calls == []
        assert session.post_calls == [
            {"url": f"{PIHA_URL}/api/embeddings", "json": {"model": "bge-m3", "prompt": "q"}, "timeout": None}
        ]

    async def test_solaria_timeout_mid_request_falls_through_to_piha_same_request(self):
        breaker = SolCircuitBreaker()
        breaker.set("up")  # cache says up; the real call below discovers it's actually stuck
        session = _FakeFallbackSession(solaria_post_behavior="timeout", piha_post_behavior="ok")
        embedding, status = await embed_with_fallback(
            breaker, session, SOLARIA_URL, PIHA_URL, "bge-m3", "q"
        )
        assert status == "down"
        assert len(embedding) == 1024
        assert breaker.status == "down"  # one-shot switch persists for the rest of the cache window
        assert [c["url"] for c in session.post_calls] == [
            f"{SOLARIA_URL}/api/embeddings", f"{PIHA_URL}/api/embeddings",
        ]

    async def test_solaria_connection_error_mid_request_falls_through(self):
        breaker = SolCircuitBreaker()
        breaker.set("up")
        session = _FakeFallbackSession(solaria_post_behavior="error", piha_post_behavior="ok")
        embedding, status = await embed_with_fallback(
            breaker, session, SOLARIA_URL, PIHA_URL, "bge-m3", "q"
        )
        assert status == "down"
        assert breaker.status == "down"

    async def test_both_legs_failing_raises_to_caller(self):
        breaker = SolCircuitBreaker()
        breaker.set("up")
        session = _FakeFallbackSession(solaria_post_behavior="timeout", piha_post_behavior="error")
        with pytest.raises(aiohttp.ClientError):
            await embed_with_fallback(breaker, session, SOLARIA_URL, PIHA_URL, "bge-m3", "q")

    async def test_both_legs_use_identical_embed_model(self):
        # Structural proof of the "no per-request DB check needed" reasoning (module docstring):
        # a single embed_model argument is threaded through both the failed SOLARIA attempt and
        # the successful PIHA attempt in the same request.
        breaker = SolCircuitBreaker()
        breaker.set("up")
        session = _FakeFallbackSession(solaria_post_behavior="timeout", piha_post_behavior="ok")
        await embed_with_fallback(breaker, session, SOLARIA_URL, PIHA_URL, "bge-m3", "q")
        models = {c["json"]["model"] for c in session.post_calls}
        assert models == {"bge-m3"}

    async def test_piha_leg_has_no_hard_timeout_override(self):
        # Only the SOLARIA leg gets the interactive-request hard timeout (plan §2 step 3) -- the
        # PIHA leg is the fallback of last resort, no shorter budget to enforce beyond it.
        breaker = SolCircuitBreaker()
        breaker.set("down")
        session = _FakeFallbackSession(piha_post_behavior="ok")
        await embed_with_fallback(breaker, session, SOLARIA_URL, PIHA_URL, "bge-m3", "q")
        assert session.post_calls[0]["timeout"] is None