litellm/tests/test_litellm/proxy/test_openapi_schema_validation.py

"""
Test that the OpenAPI schema generated by FastAPI is valid for specific endpoints.

Validates fixes for:
- /spend/calculate response schema (must use proper OpenAPI 3.x content wrapper)
- /credentials/by_model/{model_id} path parameter (must not leak credential_name)

Related issue: https://github.com/BerriAI/litellm/issues/21305
"""

import pytest


class TestSpendCalculateOpenAPISchema:
    """Test /spend/calculate response schema is valid OpenAPI 3.x."""

    def test_response_schema_has_description(self):
        """The 200 response must have a 'description' field per OpenAPI 3.x spec."""
        from litellm.proxy.spend_tracking.spend_management_endpoints import router

        for route in router.routes:
            if hasattr(route, "path") and route.path == "/spend/calculate":
                responses = route.responses or {}
                response_200 = responses.get(200, {})
                assert (
                    "description" in response_200
                ), "/spend/calculate 200 response must have a 'description' field"
                break
        else:
            pytest.fail("/spend/calculate route not found in router")

    def test_response_schema_has_content_wrapper(self):
        """The 200 response must use 'content' wrapper, not bare properties."""
        from litellm.proxy.spend_tracking.spend_management_endpoints import router

        for route in router.routes:
            if hasattr(route, "path") and route.path == "/spend/calculate":
                responses = route.responses or {}
                response_200 = responses.get(200, {})
                # Must NOT have 'cost' as a top-level key (invalid OpenAPI)
                assert "cost" not in response_200, (
                    "/spend/calculate 200 response must not have 'cost' as a "
                    "top-level property - use 'content' wrapper instead"
                )
                # Must have 'content' wrapper
                assert (
                    "content" in response_200
                ), "/spend/calculate 200 response must have a 'content' field"
                content = response_200["content"]
                assert "application/json" in content
                assert "schema" in content["application/json"]
                break
        else:
            pytest.fail("/spend/calculate route not found in router")


class TestCredentialEndpointsOpenAPISchema:
    """Test /credentials endpoints have correct path parameters."""

    def test_by_name_and_by_model_are_separate_handlers(self):
        """
        /credentials/by_name/{credential_name} and /credentials/by_model/{model_id}
        must be separate handler functions so each only declares its own path params.
        """
        from litellm.proxy.credential_endpoints.endpoints import router

        by_name_routes = []
        by_model_routes = []
        for route in router.routes:
            if not hasattr(route, "path"):
                continue
            if "by_name" in route.path:
                by_name_routes.append(route)
            elif "by_model" in route.path:
                by_model_routes.append(route)

        assert len(by_name_routes) == 1, "Expected exactly one by_name route"
        assert len(by_model_routes) == 1, "Expected exactly one by_model route"

        # They must be different endpoint functions
        by_name_endpoint = by_name_routes[0].endpoint
        by_model_endpoint = by_model_routes[0].endpoint
        assert by_name_endpoint is not by_model_endpoint, (
            "by_name and by_model must be separate handler functions "
            "to avoid path parameter conflicts in OpenAPI spec"
        )

    def test_by_model_route_does_not_require_credential_name(self):
        """
        The /credentials/by_model/{model_id} route must NOT have
        credential_name as a parameter.
        """
        import inspect
        from litellm.proxy.credential_endpoints.endpoints import (
            get_credential_by_model,
        )

        sig = inspect.signature(get_credential_by_model)
        param_names = list(sig.parameters.keys())
        assert (
            "credential_name" not in param_names
        ), "get_credential_by_model must not have a credential_name parameter"

    def test_by_name_route_does_not_require_model_id(self):
        """
        The /credentials/by_name/{credential_name} route must NOT have
        model_id as a parameter.
        """
        import inspect
        from litellm.proxy.credential_endpoints.endpoints import (
            get_credential_by_name,
        )

        sig = inspect.signature(get_credential_by_name)
        param_names = list(sig.parameters.keys())
        assert (
            "model_id" not in param_names
        ), "get_credential_by_name must not have a model_id parameter"

    def test_by_model_has_model_id_path_param(self):
        """The by_model handler must accept model_id as a path parameter."""
        import inspect
        from litellm.proxy.credential_endpoints.endpoints import (
            get_credential_by_model,
        )

        sig = inspect.signature(get_credential_by_model)
        assert (
            "model_id" in sig.parameters
        ), "get_credential_by_model must have a model_id parameter"

    def test_by_name_has_credential_name_path_param(self):
        """The by_name handler must accept credential_name as a path parameter."""
        import inspect
        from litellm.proxy.credential_endpoints.endpoints import (
            get_credential_by_name,
        )

        sig = inspect.signature(get_credential_by_name)
        assert (
            "credential_name" in sig.parameters
        ), "get_credential_by_name must have a credential_name parameter"


class TestWebSocketStubInjection:
    """
    Regression test for the v1.82.3 bug where adding a WebSocket route on a path
    that already had an HTTP route silently dropped the HTTP operation from the
    OpenAPI schema.

    Related case: 2026-05-05-madhu-swagger-responses-missing
    """

    def _make_fake_ws_route(self, path: str, name: str = "fake_ws"):
        """Minimal stand-in for fastapi.routing.APIWebSocketRoute for the helper's purposes."""
        from types import SimpleNamespace

        return SimpleNamespace(path=path, name=name, dependant=None)

    def test_websocket_stub_does_not_clobber_existing_post(self):
        """
        When a WebSocket route shares its path with an existing POST operation,
        the POST must survive — the WebSocket stub is added alongside, not on top.
        """
        from litellm.proxy.proxy_server import (
            _inject_websocket_stubs_into_openapi_schema,
        )

        schema = {
            "paths": {
                "/v1/responses": {
                    "post": {"summary": "responses_api", "operationId": "responses_api"}
                }
            }
        }
        ws_routes = [self._make_fake_ws_route("/v1/responses", name="responses_ws")]

        result = _inject_websocket_stubs_into_openapi_schema(schema, ws_routes)

        assert (
            "post" in result["paths"]["/v1/responses"]
        ), "POST operation must be preserved when a WebSocket route shares the path"
        assert (
            result["paths"]["/v1/responses"]["post"]["operationId"] == "responses_api"
        )
        assert (
            "get" in result["paths"]["/v1/responses"]
        ), "WebSocket stub should also be added under 'get'"
        assert result["paths"]["/v1/responses"]["get"]["tags"] == ["WebSocket"]

    def test_websocket_stub_added_when_path_is_new(self):
        """
        When a WebSocket route's path is not already in the schema, the stub
        creates a fresh entry — preserving the original behavior for WebSocket-only
        paths.
        """
        from litellm.proxy.proxy_server import (
            _inject_websocket_stubs_into_openapi_schema,
        )

        schema = {"paths": {}}
        ws_routes = [self._make_fake_ws_route("/ws_only", name="ws_only")]

        result = _inject_websocket_stubs_into_openapi_schema(schema, ws_routes)

        assert "/ws_only" in result["paths"]
        assert "get" in result["paths"]["/ws_only"]
        assert result["paths"]["/ws_only"]["get"]["tags"] == ["WebSocket"]

    def test_websocket_stub_skipped_when_existing_get(self):
        """
        If a real GET is already documented on the path, the WebSocket stub is
        skipped — a real operation always wins over the synthetic stub. This
        closes the same trap for future GET-vs-WebSocket collisions.
        """
        from litellm.proxy.proxy_server import (
            _inject_websocket_stubs_into_openapi_schema,
        )

        schema = {
            "paths": {
                "/health": {
                    "get": {"summary": "health_check", "operationId": "real_get"}
                }
            }
        }
        ws_routes = [self._make_fake_ws_route("/health", name="health_ws")]

        result = _inject_websocket_stubs_into_openapi_schema(schema, ws_routes)

        assert (
            result["paths"]["/health"]["get"]["operationId"] == "real_get"
        ), "Real GET must take precedence over WebSocket stub"

    def test_responses_post_routes_registered_on_router(self):
        """
        Sanity check: the three POST routes for the responses API are still wired
        on the responses router. Guards against accidental removal at the source.
        """
        from litellm.proxy.response_api_endpoints.endpoints import router

        post_paths = {
            route.path
            for route in router.routes
            if hasattr(route, "methods")
            and "POST" in (route.methods or set())
            and route.path in {"/v1/responses", "/responses", "/openai/v1/responses"}
        }
        assert post_paths == {"/v1/responses", "/responses", "/openai/v1/responses"}