import importlib import os os.environ.setdefault("OPENAI_BASE_URL", "http://test/v1") from typing import ClassVar, NoReturn from unittest.mock import patch import httpx import pytest import src.memory.embedder as embedder_mod from src.memory.embedder import embed_texts def test_embed_texts_success() -> None: fake_embedding = [0.1, 0.2, 0.3] class FakeResponse: status_code = 200 def json(self): return { "data": [{"embedding": fake_embedding, "index": 0}], "model": "gemini-embedding-2-preview", } def raise_for_status(self) -> None: pass def mock_post(self, url, **kwargs): assert "embeddings" in url return FakeResponse() with patch.object(httpx.Client, "post", mock_post): result = embed_texts(["test text"]) assert len(result) == 1 assert result[0] == fake_embedding def test_embed_texts_empty() -> None: assert embed_texts([]) == [] def test_embed_texts_api_error() -> None: class FakeErrorResponse: status_code = 401 def raise_for_status(self) -> NoReturn: msg = "Unauthorized" raise httpx.HTTPStatusError( msg, request=None, response=self, ) def json(self): return {} def mock_post(self, url, **kwargs): return FakeErrorResponse() with patch.object(httpx.Client, "post", mock_post): result = embed_texts(["test"]) assert result is None def test_embed_texts_no_env_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("OPENAI_BASE_URL", raising=False) assert embed_texts(["test"]) is None def test_embed_texts_empty_input() -> None: assert embed_texts([]) == [] def test_embed_texts_default_model(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("OPENAI_EMBEDDING_MODEL", raising=False) importlib.reload(embedder_mod) captured: dict[str, str | list[str]] = {} class FakeResponse: status_code = 200 def json(self): return {"data": [{"embedding": [0.1], "index": 0}]} def raise_for_status(self) -> None: pass def mock_post(self, url, **kwargs): captured.update(kwargs["json"]) return FakeResponse() with patch.object(httpx.Client, "post", mock_post): embedder_mod.embed_texts(["text"]) assert captured["model"] == "gemini-embedding-2-preview" importlib.reload(embedder_mod) def test_embed_texts_custom_model(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small") importlib.reload(embedder_mod) captured: dict[str, str | list[str]] = {} class FakeResponse: status_code = 200 def json(self): return {"data": [{"embedding": [0.1], "index": 0}]} def raise_for_status(self) -> None: pass def mock_post(self, url, **kwargs): captured.update(kwargs["json"]) return FakeResponse() with patch.object(httpx.Client, "post", mock_post): embedder_mod.embed_texts(["text"]) assert captured["model"] == "text-embedding-3-small" monkeypatch.delenv("OPENAI_EMBEDDING_MODEL", raising=False) importlib.reload(embedder_mod) def test_embed_texts_trailing_slash(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("OPENAI_BASE_URL", "https://api.test.com/v1/") importlib.reload(embedder_mod) captured_url: dict[str, str] = {} class FakeResponse: status_code = 200 def json(self): return {"data": [{"embedding": [0.1], "index": 0}]} def raise_for_status(self) -> None: pass def mock_post(self, url, **kwargs): captured_url["url"] = url return FakeResponse() with patch.object(httpx.Client, "post", mock_post): embedder_mod.embed_texts(["text"]) assert captured_url["url"] == "https://api.test.com/v1/embeddings" monkeypatch.setenv("OPENAI_BASE_URL", "http://test/v1") importlib.reload(embedder_mod) def test_embed_texts_batches(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(embedder_mod, "BATCH_SIZE", 2048) calls: list[int] = [] class FakeResponse: def __init__(self, count: int) -> None: self.count = count status_code = 200 def json(self): return {"data": [{"embedding": [0.1], "index": i} for i in range(self.count)]} def raise_for_status(self) -> None: pass def mock_post(self, url, **kwargs): count = len(kwargs["json"]["input"]) calls.append(count) return FakeResponse(count) with patch.object(httpx.Client, "post", mock_post): result = embedder_mod.embed_texts(["text"] * 3000) assert len(calls) == 2 assert calls == [2048, 952] assert len(result) == 3000 def test_embed_texts_batch_progress( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] ) -> None: monkeypatch.setattr(embedder_mod, "BATCH_SIZE", 2048) class FakeResponse: status_code = 200 def json(self): return {"data": [{"embedding": [0.1], "index": 0}]} def raise_for_status(self) -> None: pass def mock_post(self, url, **kwargs): return FakeResponse() with patch.object(httpx.Client, "post", mock_post): embedder_mod.embed_texts(["text"] * 3000) captured = capsys.readouterr() assert "batch 1/2" in captured.out assert "batch 2/2" in captured.out def test_retry_after_header() -> None: class RetryResponse: status_code = 429 headers: ClassVar[dict[str, str]] = {"Retry-After": "2"} def raise_for_status(self) -> NoReturn: raise httpx.HTTPStatusError( "Too Many Requests", request=None, response=self, ) def json(self): return {} class OkResponse: status_code = 200 def json(self): return {"data": [{"embedding": [0.1], "index": 0}]} def raise_for_status(self) -> None: pass responses = [RetryResponse(), OkResponse()] def mock_post(self, url, **kwargs): return responses.pop(0) with ( patch.object(httpx.Client, "post", mock_post), patch("src.memory.embedder.time.sleep") as mock_sleep, ): result = embed_texts(["text"]) mock_sleep.assert_any_call(2) assert result == [[0.1]]