"""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