didi-lot1-ai/ai_platform/modules/llm-inference/tests/test_api.py

"""Tests for FastAPI routes."""

from unittest.mock import AsyncMock, patch

import pytest
from fastapi.testclient import TestClient

from llm_inference.api.app import create_app
from llm_inference.api.dependencies import init_concurrency_limiter
from llm_inference.client import LLMClient
from llm_inference.config import LLMSettings, SettingsCache
from llm_inference.schemas import CompletionResponse


@pytest.fixture
def app_client(test_settings: LLMSettings) -> TestClient:
    """Create a test client with mocked client."""
    # Inject test settings into cache BEFORE creating app
    SettingsCache.set(test_settings)

    # Initialize concurrency limiter for tests
    init_concurrency_limiter(test_settings.max_concurrent_completions)

    app = create_app()
    app.state.settings = test_settings
    app.state.client = LLMClient(settings=test_settings)
    return TestClient(app)


class TestHealthRoutes:
    """Tests for health check routes."""

    def test_health_endpoint(self, app_client: TestClient) -> None:
        """Test /health endpoint."""
        response = app_client.get("/health")

        assert response.status_code == 200
        data = response.json()
        assert "status" in data
        assert "backends" in data
        assert data["status"] in ["healthy", "degraded", "unhealthy"]

    def test_ready_endpoint(self, app_client: TestClient) -> None:
        """Test /ready endpoint."""
        response = app_client.get("/ready")

        assert response.status_code == 200
        data = response.json()
        assert data["ready"] is True


class TestModelsRoutes:
    """Tests for model management routes."""

    def test_list_models(self, app_client: TestClient) -> None:
        """Test GET /v1/models."""
        response = app_client.get("/v1/models")

        assert response.status_code == 200
        data = response.json()
        assert "data" in data
        assert "object" in data
        assert data["object"] == "list"

    def test_list_models_with_backend_filter(self, app_client: TestClient) -> None:
        """Test GET /v1/models with backend filter."""
        response = app_client.get("/v1/models?backend=litellm")

        assert response.status_code == 200
        data = response.json()
        # All models should be from litellm backend
        for model in data["data"]:
            assert model["backend"] == "litellm"

    def test_list_models_invalid_backend(self, app_client: TestClient) -> None:
        """Test GET /v1/models with invalid backend."""
        response = app_client.get("/v1/models?backend=invalid")

        assert response.status_code == 400

    def test_list_backends(self, app_client: TestClient) -> None:
        """Test GET /v1/backends."""
        response = app_client.get("/v1/backends")

        assert response.status_code == 200
        data = response.json()
        assert "backends" in data
        assert "litellm" in data["backends"]


class TestCompletionsRoutes:
    """Tests for completion routes."""

    def test_chat_completions(
        self,
        app_client: TestClient,
        mock_completion_response: CompletionResponse,
    ) -> None:
        """Test POST /v1/chat/completions."""
        # Mock the client's complete method
        with patch.object(
            app_client.app.state.client,
            "complete",
            new_callable=AsyncMock,
            return_value=mock_completion_response,
        ):
            response = app_client.post(
                "/v1/chat/completions",
                json={
                    "messages": [{"role": "user", "content": "Hello"}],
                    "model": "gpt-3.5-turbo",
                },
            )

            assert response.status_code == 200
            data = response.json()
            assert "id" in data
            assert "choices" in data
            assert "model" in data

    def test_chat_completions_with_parameters(
        self,
        app_client: TestClient,
        mock_completion_response: CompletionResponse,
    ) -> None:
        """Test completion with optional parameters."""
        with patch.object(
            app_client.app.state.client,
            "complete",
            new_callable=AsyncMock,
            return_value=mock_completion_response,
        ):
            response = app_client.post(
                "/v1/chat/completions",
                json={
                    "messages": [{"role": "user", "content": "Hello"}],
                    "model": "gpt-4",
                    "temperature": 0.5,
                    "max_tokens": 100,
                    "top_p": 0.9,
                },
            )

            assert response.status_code == 200

    def test_chat_completions_with_backend_override(
        self,
        app_client: TestClient,
        mock_completion_response: CompletionResponse,
    ) -> None:
        """Test completion with backend override."""
        with patch.object(
            app_client.app.state.client,
            "complete",
            new_callable=AsyncMock,
            return_value=mock_completion_response,
        ):
            response = app_client.post(
                "/v1/chat/completions",
                json={
                    "messages": [{"role": "user", "content": "Hello"}],
                    "model": "gpt-3.5-turbo",
                    "backend": "litellm",
                },
            )

            assert response.status_code == 200

    def test_chat_completions_missing_messages(self, app_client: TestClient) -> None:
        """Test completion without messages returns error."""
        response = app_client.post(
            "/v1/chat/completions",
            json={"model": "gpt-3.5-turbo"},
        )

        assert response.status_code == 422  # Validation error

    def test_chat_completions_missing_model(self, app_client: TestClient) -> None:
        """Test completion without model returns error."""
        response = app_client.post(
            "/v1/chat/completions",
            json={"messages": [{"role": "user", "content": "Hello"}]},
        )

        assert response.status_code == 422  # Validation error


class TestModelLoadRoutes:
    """Tests for model load/unload routes."""

    def test_load_model_unsupported_backend(self, app_client: TestClient) -> None:
        """Test loading model on backend that doesn't support it."""
        response = app_client.post(
            "/v1/models/load",
            json={"model": "test-model", "backend": "litellm"},
        )

        # LiteLLM doesn't support model loading
        assert response.status_code == 400

    def test_unload_model_unsupported_backend(self, app_client: TestClient) -> None:
        """Test unloading model on backend that doesn't support it."""
        response = app_client.post(
            "/v1/models/unload",
            json={"model": "test-model", "backend": "litellm"},
        )

        # LiteLLM doesn't support model unloading
        assert response.status_code == 400