Livrare LOT 1 - Didi

ai_platform/modules/rerank/tests/conftest.py

@@ -0,0 +1,81 @@
+"""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)