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

"""Unit tests for the SOLARIA->PIHA embed fallback state machine (app.embed_router) -- no
real Ollama on either side. Fake backends are keyed by URL so one fake session can serve
both legs; the clock is injected so the 30 s TTL is tested without sleeping."""
from __future__ import annotations

import asyncio
import pathlib
import sys

import aiohttp
import pytest

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

from app.embed_router import EmbedBackendError, EmbedRouter, ModelMismatchError  # noqa: E402

PRIMARY = "http://solaria:11434"
FALLBACK = "http://192.168.31.5:11434"


class _FakeResponse:
    def __init__(self, payload=None, exc=None, status=200, delay=0.0):
        self._payload = payload
        self._exc = exc
        self.status = status  # check_ollama_health reads resp.status directly
        self._delay = delay

    async def __aenter__(self):
        if self._delay:
            await asyncio.sleep(self._delay)
        if self._exc is not None:
            raise self._exc
        return self

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

    def raise_for_status(self):
        pass

    async def json(self):
        return self._payload


class _FakeBackend:
    """One fake Ollama: `models` drives /api/tags, `down=True` refuses every connection,
    `embed_exc` makes /api/embeddings fail while /api/tags still answers (the plan's step-3b
    scenario: probe says up, the real embed then dies mid-query)."""

    def __init__(self, models=("bge-m3:latest",), down=False, embed_exc=None, vector_value=0.01,
                 embed_delay=0.0):
        self.models = list(models)
        self.down = down
        self.embed_exc = embed_exc
        self.vector_value = vector_value
        self.embed_delay = embed_delay  # seconds the /api/embeddings answer hangs before serving
        self.tags_calls = 0
        self.embed_calls = 0


class _FakeSession:
    def __init__(self, backends: dict):
        self._backends = backends

    def _backend(self, url):
        for base, backend in self._backends.items():
            if url.startswith(base):
                return backend
        raise AssertionError(f"unexpected url: {url}")

    def get(self, url, timeout=None):
        backend = self._backend(url)
        backend.tags_calls += 1
        if backend.down:
            return _FakeResponse(exc=aiohttp.ClientConnectionError("connection refused"))
        return _FakeResponse({"models": [{"name": n} for n in backend.models]})

    def post(self, url, json):
        backend = self._backend(url)
        backend.embed_calls += 1
        if backend.down:
            return _FakeResponse(exc=aiohttp.ClientConnectionError("connection refused"))
        if backend.embed_exc is not None:
            return _FakeResponse(exc=backend.embed_exc)
        return _FakeResponse({"embedding": [backend.vector_value] * 1024},
                             delay=backend.embed_delay)


class _FakeClock:
    def __init__(self):
        self.now = 1000.0

    def __call__(self):
        return self.now


def _router(fallback_url=FALLBACK, clock=None, **kwargs):
    return EmbedRouter(
        PRIMARY,
        fallback_url,
        embed_model="bge-m3",
        clock=clock or _FakeClock(),
        **kwargs,
    )


class TestPrimaryHealthy:
    async def test_embeds_on_primary_and_reports_its_name(self):
        solaria = _FakeBackend()
        piha = _FakeBackend()
        session = _FakeSession({PRIMARY: solaria, FALLBACK: piha})
        embedding, backend = await _router().embed(session, "q")
        assert backend == "solaria"
        assert len(embedding) == 1024
        assert solaria.embed_calls == 1
        assert piha.embed_calls == 0

    async def test_health_verdict_is_cached_within_ttl(self):
        solaria = _FakeBackend()
        clock = _FakeClock()
        router = _router(clock=clock)
        session = _FakeSession({PRIMARY: solaria, FALLBACK: _FakeBackend()})
        await router.embed(session, "q1")
        probes_after_first = solaria.tags_calls  # probe + first-use model verification
        clock.now += 10  # inside the 30 s TTL
        await router.embed(session, "q2")
        assert solaria.tags_calls == probes_after_first  # no new probe, model check done once

    async def test_model_tag_with_latest_suffix_satisfies_bare_model_name(self):
        solaria = _FakeBackend(models=["bge-m3:latest"])
        session = _FakeSession({PRIMARY: solaria, FALLBACK: _FakeBackend()})
        _, backend = await _router().embed(session, "q")
        assert backend == "solaria"


class TestFailover:
    async def test_primary_down_at_probe_falls_back_to_piha(self):
        solaria = _FakeBackend(down=True)
        piha = _FakeBackend(vector_value=0.02)
        session = _FakeSession({PRIMARY: solaria, FALLBACK: piha})
        embedding, backend = await _router().embed(session, "q")
        assert backend == "piha"
        assert embedding[0] == 0.02
        assert solaria.embed_calls == 0

    async def test_mid_embed_failure_serves_same_query_from_fallback(self):
        # Plan §2 D2 step 3b: probe said up, the real embed then dies -- the SAME request
        # must come back from the fallback, never as an error to the user.
        solaria = _FakeBackend(embed_exc=aiohttp.ClientConnectionError("died mid-embed"))
        piha = _FakeBackend(vector_value=0.02)
        session = _FakeSession({PRIMARY: solaria, FALLBACK: piha})
        router = _router()
        embedding, backend = await router.embed(session, "q")
        assert backend == "piha"
        assert embedding[0] == 0.02
        assert solaria.embed_calls == 1

    async def test_mid_embed_timeout_serves_same_query_from_fallback(self):
        # The step-3b guarantee for the OTHER failure shape: SOLARIA hangs instead of
        # refusing. This exercises the asyncio.wait_for path (a hang raises TimeoutError,
        # not aiohttp.ClientError) -- the hard primary timeout must cut the hang off and
        # the SAME request must still come back from the fallback.
        solaria = _FakeBackend(embed_delay=0.2)
        piha = _FakeBackend(vector_value=0.02)
        session = _FakeSession({PRIMARY: solaria, FALLBACK: piha})
        router = _router(primary_embed_timeout_s=0.05)
        embedding, backend = await router.embed(session, "q")
        assert backend == "piha"
        assert embedding[0] == 0.02
        assert solaria.embed_calls == 1

    async def test_mid_embed_failure_opens_circuit_for_subsequent_requests(self):
        # The one-shot flip must persist: after a mid-embed failure, requests inside the
        # same TTL window go straight to the fallback -- no re-probe, no primary retry.
        solaria = _FakeBackend(embed_exc=aiohttp.ClientConnectionError("died mid-embed"))
        piha = _FakeBackend()
        clock = _FakeClock()
        router = _router(clock=clock)
        session = _FakeSession({PRIMARY: solaria, FALLBACK: piha})
        await router.embed(session, "q1")  # probe says up -> embed dies -> flip to down
        tags_after_first = solaria.tags_calls
        clock.now += 10  # still inside the TTL window the flip opened
        _, backend = await router.embed(session, "q2")
        assert backend == "piha"
        assert solaria.embed_calls == 1  # primary never retried inside the window
        assert solaria.tags_calls == tags_after_first  # and never re-probed either

    async def test_fallback_embed_is_not_bounded_by_the_primary_timeout(self):
        # Last-resort semantics: a Pi-5 CPU embed plus a cold model load is legitimately
        # slow -- the fallback leg must NOT inherit the primary's hard timeout, or "slow
        # but alive" would turn back into "dead".
        solaria = _FakeBackend(down=True)
        piha = _FakeBackend(embed_delay=0.2, vector_value=0.02)
        session = _FakeSession({PRIMARY: solaria, FALLBACK: piha})
        router = _router(primary_embed_timeout_s=0.05)
        embedding, backend = await router.embed(session, "q")
        assert backend == "piha"
        assert embedding[0] == 0.02

    async def test_down_verdict_is_cached_and_skips_primary_until_ttl_expires(self):
        solaria = _FakeBackend(down=True)
        piha = _FakeBackend()
        clock = _FakeClock()
        router = _router(clock=clock)
        session = _FakeSession({PRIMARY: solaria, FALLBACK: piha})
        await router.embed(session, "q1")
        probes = solaria.tags_calls
        clock.now += 10  # still inside TTL -> no re-probe, straight to fallback
        _, backend = await router.embed(session, "q2")
        assert backend == "piha"
        assert solaria.tags_calls == probes

    async def test_traffic_returns_to_primary_after_ttl_expiry(self):
        # Test C from the task spec: SOLARIA comes back -> within one TTL window the
        # router re-probes and routes to the GPU again.
        solaria = _FakeBackend(down=True)
        piha = _FakeBackend()
        clock = _FakeClock()
        router = _router(clock=clock)
        session = _FakeSession({PRIMARY: solaria, FALLBACK: piha})
        _, backend = await router.embed(session, "q1")
        assert backend == "piha"
        solaria.down = False  # SOLARIA wakes up
        clock.now += 31  # TTL (30 s) expired
        _, backend = await router.embed(session, "q2")
        assert backend == "solaria"


class TestNoUsableBackend:
    async def test_primary_down_without_fallback_raises_backend_error(self):
        solaria = _FakeBackend(down=True)
        session = _FakeSession({PRIMARY: solaria})
        with pytest.raises(EmbedBackendError, match="no EMBED_FALLBACK_URL"):
            await _router(fallback_url=None).embed(session, "q")

    async def test_both_backends_down_raises_backend_error(self):
        session = _FakeSession({PRIMARY: _FakeBackend(down=True), FALLBACK: _FakeBackend(down=True)})
        with pytest.raises(EmbedBackendError, match="unreachable"):
            await _router().embed(session, "q")


class TestModelInvariant:
    async def test_fallback_without_bge_m3_is_a_loud_error_not_a_silent_embed(self):
        solaria = _FakeBackend(down=True)
        piha = _FakeBackend(models=["llama3:8b"])
        session = _FakeSession({PRIMARY: solaria, FALLBACK: piha})
        with pytest.raises(ModelMismatchError, match="bge-m3"):
            await _router().embed(session, "q")
        assert piha.embed_calls == 0  # never embedded in the wrong vector space

    async def test_primary_without_bge_m3_fails_over_to_verified_fallback(self):
        solaria = _FakeBackend(models=["llama3:8b"])
        piha = _FakeBackend()
        session = _FakeSession({PRIMARY: solaria, FALLBACK: piha})
        _, backend = await _router().embed(session, "q")
        assert backend == "piha"
        assert solaria.embed_calls == 0


class TestStatusReporting:
    async def test_primary_status_up_and_fallback_status_up(self):
        session = _FakeSession({PRIMARY: _FakeBackend(), FALLBACK: _FakeBackend()})
        router = _router()
        assert await router.primary_status(session) == "up"
        assert await router.fallback_status(session) == "up"

    async def test_fallback_status_unconfigured_without_fallback_url(self):
        session = _FakeSession({PRIMARY: _FakeBackend()})
        assert await _router(fallback_url=None).fallback_status(session) == "unconfigured"

    async def test_primary_status_down_when_probe_fails(self):
        session = _FakeSession({PRIMARY: _FakeBackend(down=True), FALLBACK: _FakeBackend()})
        assert await _router().primary_status(session) == "down"