opencode-config/tests/test_embedder.py

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), pytest.raises(httpx.HTTPStatusError):
        embed_texts(["test"])


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