homelab-codex-ws/packages/kb-retrieval/tests/test_embed.py

"""Unit tests for the Ollama embedding client -- no real HTTP, no real Ollama."""
from __future__ import annotations

import pytest

from kb_retrieval.embed import _vector_literal, check_ollama_health, embed_chunk


class TestVectorLiteral:
    def test_formats_as_bracketed_csv(self):
        assert _vector_literal([0.1, 0.2, -0.3]) == "[0.1,0.2,-0.3]"


class _FakeEmbedResponse:
    def __init__(self, payload, status=200):
        self._payload = payload
        self._status = status

    async def __aenter__(self):
        return self

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

    async def json(self):
        return self._payload

    def raise_for_status(self):
        if self._status >= 400:
            raise RuntimeError(f"HTTP {self._status}")


class _FakeOllamaSession:
    """Serves a fixed embedding vector for every /api/embeddings POST, or errors by prompt."""

    def __init__(self, dim=1024, fail_for=None):
        self._dim = dim
        self._fail_for = fail_for or set()
        self.requests: list[dict] = []

    def post(self, url, json):
        self.requests.append({"url": url, "json": json})
        if json["prompt"] in self._fail_for:
            return _FakeEmbedResponse({}, status=500)
        return _FakeEmbedResponse({"embedding": [0.01] * self._dim})


class TestEmbedChunk:
    async def test_returns_embedding_and_elapsed(self):
        session = _FakeOllamaSession(dim=1024)
        embedding, elapsed = await embed_chunk(session, "http://fake-ollama", "bge-m3", "hello world")
        assert len(embedding) == 1024
        assert elapsed >= 0
        assert session.requests == [
            {"url": "http://fake-ollama/api/embeddings", "json": {"model": "bge-m3", "prompt": "hello world"}}
        ]

    async def test_missing_embedding_key_raises(self):
        class _EmptyResponse(_FakeEmbedResponse):
            pass

        class _Session:
            def post(self, url, json):
                return _EmptyResponse({})

        with pytest.raises(ValueError):
            await embed_chunk(_Session(), "http://fake-ollama", "bge-m3", "text")


class _FakeTagsResponse:
    def __init__(self, status):
        self.status = status

    async def __aenter__(self):
        return self

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


class _FakeHealthSession:
    def __init__(self, status=200, raises=None):
        self._status = status
        self._raises = raises

    def get(self, url, timeout=None):
        if self._raises is not None:
            raise self._raises
        return _FakeTagsResponse(self._status)


class TestCheckOllamaHealth:
    async def test_up_on_2xx(self):
        session = _FakeHealthSession(status=200)
        assert await check_ollama_health(session, "http://fake-ollama", 0.5) is True

    async def test_down_on_error_status(self):
        session = _FakeHealthSession(status=500)
        assert await check_ollama_health(session, "http://fake-ollama", 0.5) is False

    async def test_down_on_connection_error(self):
        import aiohttp

        session = _FakeHealthSession(raises=aiohttp.ClientConnectionError("refused"))
        assert await check_ollama_health(session, "http://fake-ollama", 0.5) is False

    async def test_down_on_timeout(self):
        session = _FakeHealthSession(raises=TimeoutError())
        assert await check_ollama_health(session, "http://fake-ollama", 0.5) is False