81 lines
2.3 KiB
Python
81 lines
2.3 KiB
Python
"""Pytest fixtures for rerank tests."""
|
|
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
from rerank.api.app import create_app
|
|
from rerank.api.dependencies import init_concurrency_limiter
|
|
from rerank.client import RerankClient
|
|
from rerank.config import RerankSettings, SettingsCache
|
|
from rerank.schemas import RerankResponse
|
|
from rerank.types import BackendType, RerankResult, RerankUsage
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clear_settings_cache() -> None:
|
|
"""Clear settings cache before each test."""
|
|
SettingsCache.clear()
|
|
|
|
|
|
@pytest.fixture
|
|
def test_settings() -> RerankSettings:
|
|
"""Create test settings with mocked values."""
|
|
return RerankSettings(
|
|
default_backend="vllm",
|
|
enable_vllm=True,
|
|
enable_llamacpp=False,
|
|
external_url="http://localhost:54200",
|
|
host="127.0.0.1",
|
|
port=54200,
|
|
api_tokens=None,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_rerank_response() -> RerankResponse:
|
|
"""Create a mock rerank response."""
|
|
return RerankResponse(
|
|
model="test-model",
|
|
results=[
|
|
RerankResult(index=2, relevance_score=0.95),
|
|
RerankResult(index=0, relevance_score=0.82),
|
|
RerankResult(index=1, relevance_score=0.10),
|
|
],
|
|
usage=RerankUsage(total_tokens=150),
|
|
backend="vllm",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def client_with_mock_backend(
|
|
test_settings: RerankSettings, mock_rerank_response: RerankResponse
|
|
) -> RerankClient:
|
|
"""Create a RerankClient with mocked backend."""
|
|
SettingsCache.set(test_settings)
|
|
|
|
client = RerankClient(settings=test_settings)
|
|
|
|
# Mock the backend's rerank method
|
|
backend = client.registry.get(BackendType.VLLM)
|
|
results = [(r.index, r.relevance_score) for r in mock_rerank_response.results]
|
|
usage = mock_rerank_response.usage
|
|
backend.rerank = AsyncMock(return_value=(results, usage))
|
|
|
|
return client
|
|
|
|
|
|
@pytest.fixture
|
|
def app_client(test_settings: RerankSettings) -> TestClient:
|
|
"""Create a FastAPI TestClient with test settings."""
|
|
SettingsCache.set(test_settings)
|
|
|
|
init_concurrency_limiter(test_settings.max_concurrent_reranks)
|
|
|
|
app = create_app()
|
|
|
|
app.state.settings = test_settings
|
|
app.state.client = RerankClient(settings=test_settings)
|
|
|
|
return TestClient(app)
|