embedding-runtime/tests/test_server.py

from __future__ import annotations

import importlib
import sys
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from types import ModuleType

import pytest
from fastapi.testclient import TestClient

APP_ROOT = Path(__file__).resolve().parents[1]


@dataclass(frozen=True, slots=True)
class FakeVector:
    values: tuple[float, ...]

    def tolist(self) -> list[float]:
        return list(self.values)

    def dot(self, other: FakeVector) -> float:
        return sum(left * right for left, right in zip(self.values, other.values))


@dataclass(frozen=True, slots=True)
class FakeScores:
    values: tuple[float, ...]

    def __neg__(self) -> FakeScores:
        return FakeScores(tuple(-value for value in self.values))

    def __getitem__(self, index: int) -> float:
        return self.values[index]


class FakeMatrix:
    def __init__(self, rows: tuple[tuple[float, ...], ...]) -> None:
        self._vectors = tuple(FakeVector(row) for row in rows)

    def __iter__(self) -> Iterator[FakeVector]:
        return iter(self._vectors)

    def __getitem__(self, index: int) -> FakeVector:
        return self._vectors[index]

    def __matmul__(self, other: FakeVector) -> FakeScores:
        return FakeScores(tuple(vector.dot(other) for vector in self._vectors))


@dataclass(frozen=True, slots=True)
class EncodeCall:
    texts: tuple[str, ...]
    task: str | None
    instruction: str | None


class FailingSentenceTransformer:
    def __init__(self, model_name: str, *, device: str, trust_remote_code: bool) -> None:
        raise AssertionError(f"real model constructor must not run for {model_name} on {device}")


class FakeCuda:
    @staticmethod
    def is_available() -> bool:
        return False


def build_fake_torch_module() -> ModuleType:
    module = ModuleType("torch")
    setattr(module, "cuda", FakeCuda())
    return module


def build_fake_sentence_transformers_module() -> ModuleType:
    module = ModuleType("sentence_transformers")
    setattr(module, "SentenceTransformer", FailingSentenceTransformer)
    return module


def build_fake_numpy_module() -> ModuleType:
    module = ModuleType("numpy")
    setattr(module, "ndarray", FakeMatrix)

    def argsort(scores: FakeScores) -> list[int]:
        return sorted(range(len(scores.values)), key=lambda index: scores.values[index])

    setattr(module, "argsort", argsort)
    return module


@pytest.fixture()
def server_module(monkeypatch: pytest.MonkeyPatch) -> ModuleType:
    monkeypatch.syspath_prepend(str(APP_ROOT))
    monkeypatch.setitem(sys.modules, "torch", build_fake_torch_module())
    monkeypatch.setitem(sys.modules, "sentence_transformers", build_fake_sentence_transformers_module())
    monkeypatch.setitem(sys.modules, "numpy", build_fake_numpy_module())
    sys.modules.pop("server", None)
    return importlib.import_module("server")


def test_model_loading_disables_remote_code_TASK_8c19d6a7(
    server_module: ModuleType, monkeypatch: pytest.MonkeyPatch
) -> None:
    # Given: a recording model constructor at the real loader boundary.
    constructor_calls: list[tuple[str, str, bool]] = []

    class RecordingSentenceTransformer:
        def __init__(self, model_name: str, *, device: str, trust_remote_code: bool) -> None:
            constructor_calls.append((model_name, device, trust_remote_code))

    monkeypatch.setattr(server_module, "SentenceTransformer", RecordingSentenceTransformer)
    setattr(server_module, "_model", None)

    # When: the runtime lazily loads its configured model.
    server_module._load_model()

    # Then: repository-supplied model code is never executed.
    assert constructor_calls == [(server_module.MODEL_NAME, server_module.DEVICE, False)]


def test_health_remains_available_without_loading_model_SPEC_KSERVE_002_SC_KSERVE_002(
    server_module: ModuleType,
) -> None:
    client = TestClient(server_module.app)

    response = client.get("/health")

    assert response.status_code == 200
    assert response.json() == {"status": "ok", "device": "cpu", "model": "/mnt/models"}


def test_embeddings_accept_string_and_return_openai_list_shape_SPEC_KSERVE_002_SC_KSERVE_002(
    monkeypatch: pytest.MonkeyPatch,
    server_module: ModuleType,
) -> None:
    calls: list[EncodeCall] = []

    def encode(texts: list[str], task: str | None, instruction: str | None) -> FakeMatrix:
        calls.append(EncodeCall(tuple(texts), task, instruction))
        return FakeMatrix(((0.25, 0.75),))

    monkeypatch.setattr(server_module, "_encode", encode)
    client = TestClient(server_module.app)

    response = client.post("/v1/embeddings", json={"input": "hello world", "model": "test-model"})

    assert response.status_code == 200
    payload = response.json()
    assert payload["object"] == "list"
    assert payload["model"] == "test-model"
    assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": [0.25, 0.75]}]
    assert payload["usage"] == {"prompt_tokens": 2, "total_tokens": 2}
    assert calls == [EncodeCall(("hello world",), None, None)]


def test_embeddings_accept_list_input_and_preserve_indices_SPEC_KSERVE_002_SC_KSERVE_002(
    monkeypatch: pytest.MonkeyPatch,
    server_module: ModuleType,
) -> None:
    calls: list[EncodeCall] = []

    def encode(texts: list[str], task: str | None, instruction: str | None) -> FakeMatrix:
        calls.append(EncodeCall(tuple(texts), task, instruction))
        return FakeMatrix(((1.0, 0.0), (0.0, 1.0)))

    monkeypatch.setattr(server_module, "_encode", encode)
    client = TestClient(server_module.app)

    response = client.post(
        "/v1/embeddings",
        json={"input": ["first text", "second"], "task": "classification"},
    )

    assert response.status_code == 200
    payload = response.json()
    assert payload["object"] == "list"
    assert payload["data"] == [
        {"object": "embedding", "index": 0, "embedding": [1.0, 0.0]},
        {"object": "embedding", "index": 1, "embedding": [0.0, 1.0]},
    ]
    assert payload["usage"] == {"prompt_tokens": 3, "total_tokens": 3}
    assert calls == [EncodeCall(("first text", "second"), "classification", None)]


def test_rerank_returns_documents_ordered_by_score_SPEC_KSERVE_002_SC_KSERVE_002(
    monkeypatch: pytest.MonkeyPatch,
    server_module: ModuleType,
) -> None:
    calls: list[EncodeCall] = []

    def encode(texts: list[str], task: str | None, instruction: str | None) -> FakeMatrix:
        calls.append(EncodeCall(tuple(texts), task, instruction))
        if texts == ["needle"]:
            return FakeMatrix(((1.0, 0.0),))
        return FakeMatrix(((0.20, 0.0), (0.95, 0.0), (0.50, 0.0)))

    monkeypatch.setattr(server_module, "_encode", encode)
    client = TestClient(server_module.app)

    response = client.post(
        "/v1/rerank",
        json={"query": "needle", "documents": ["low", "high", "middle"], "top_n": 2},
    )

    assert response.status_code == 200
    payload = response.json()
    assert payload["model"] == "/mnt/models"
    assert payload["results"] == [
        {"index": 1, "relevance_score": 0.95, "document": "high"},
        {"index": 2, "relevance_score": 0.5, "document": "middle"},
    ]
    assert calls == [
        EncodeCall(("needle",), "retrieval.query", None),
        EncodeCall(("low", "high", "middle"), "retrieval.passage", None),
    ]