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 (
    EmbeddingDimensionError,
    _vector_literal,
    check_ollama_health,
    embed_batch,
    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 _FakeBatchResponse(_FakeEmbedResponse):
    pass


class _FakeBatchOllamaSession:
    """Serves a fixed set of embeddings for every /api/embed POST."""

    def __init__(self, dim=1024, n_override=None, bad_index=None):
        self._dim = dim
        self._n_override = n_override
        self._bad_index = bad_index
        self.requests: list[dict] = []

    def post(self, url, json):
        self.requests.append({"url": url, "json": json})
        texts = json["input"]
        n = self._n_override if self._n_override is not None else len(texts)
        embeddings = [[0.01] * self._dim for _ in range(n)]
        if self._bad_index is not None and self._bad_index < len(embeddings):
            embeddings[self._bad_index] = [0.01] * (self._dim - 1)
        return _FakeBatchResponse({"embeddings": embeddings})


class TestEmbedBatch:
    async def test_returns_embeddings_and_elapsed(self):
        session = _FakeBatchOllamaSession(dim=1024)
        texts = ["hello", "world", "third"]
        embeddings, elapsed = await embed_batch(session, "http://fake-ollama", "bge-m3", texts)
        assert len(embeddings) == 3
        assert all(len(e) == 1024 for e in embeddings)
        assert elapsed >= 0
        assert session.requests == [
            {"url": "http://fake-ollama/api/embed", "json": {"model": "bge-m3", "input": texts}}
        ]

    async def test_length_mismatch_raises_dimension_error(self):
        session = _FakeBatchOllamaSession(dim=1024, n_override=2)
        with pytest.raises(EmbeddingDimensionError):
            await embed_batch(session, "http://fake-ollama", "bge-m3", ["a", "b", "c"])

    async def test_wrong_dimension_on_one_item_raises(self):
        session = _FakeBatchOllamaSession(dim=1024, bad_index=1)
        with pytest.raises(EmbeddingDimensionError):
            await embed_batch(session, "http://fake-ollama", "bge-m3", ["a", "b", "c"])

    async def test_missing_embeddings_key_raises_value_error(self):
        class _Session:
            def post(self, url, json):
                return _FakeBatchResponse({})

        with pytest.raises(ValueError):
            await embed_batch(_Session(), "http://fake-ollama", "bge-m3", ["a"])


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