mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
424687ce09
* test(proxy): keep answering requests while a cold /openapi.json builds * fix(proxy): build the OpenAPI schema off the event loop so a cold /openapi.json no longer stalls the worker * perf(proxy): reuse the rendered /openapi.json body until the schema changes * fix(proxy): scope the openapi build's pydantic patch to the build The schema build swapped a process-wide pydantic hook, which live requests saw once the build moved to a thread, and overlapping builds could leave it installed. The hook now checks a context variable set only inside the build. The proxy test counts liveliness answers instead of timing them. * test(proxy): hold the cold openapi build on an event instead of pacing on a clock * fix(proxy): serve a replaced app.openapi like FastAPI's own docs route * test(proxy): build the cold-docs test on the proxy module's current app * test(proxy): freeze the model that holds the cold docs build
488 lines
17 KiB
Python
488 lines
17 KiB
Python
"""Behavior pins for proxy_server OpenAPI customization + CORS helpers.
|
|
|
|
Pins covered:
|
|
- ``_generate_stable_operation_id``
|
|
- ``_strip_operation_id_method_suffix``
|
|
- ``ensure_unique_openapi_operation_ids``
|
|
- ``_inject_websocket_stubs_into_openapi_schema``
|
|
- ``get_openapi_schema``
|
|
- ``custom_openapi``
|
|
- ``mount_swagger_ui``
|
|
- ``_get_cors_config``
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import threading
|
|
from types import SimpleNamespace
|
|
from typing import Final
|
|
|
|
import httpx
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from pydantic import BaseModel, ConfigDict, GetJsonSchemaHandler
|
|
from pydantic.json_schema import JsonSchemaValue
|
|
from pydantic_core import CoreSchema
|
|
|
|
from litellm.proxy.proxy_server import (
|
|
_generate_stable_operation_id,
|
|
_get_cors_config,
|
|
_inject_websocket_stubs_into_openapi_schema,
|
|
_strip_operation_id_method_suffix,
|
|
app,
|
|
custom_openapi,
|
|
ensure_unique_openapi_operation_ids,
|
|
get_openapi_schema,
|
|
mount_swagger_ui,
|
|
)
|
|
|
|
from .conftest import normalize
|
|
|
|
UNBLOCK_TIMEOUT_SECONDS: Final = 5
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _generate_stable_operation_id
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_generate_stable_operation_id_single_method_appends_suffix():
|
|
route = SimpleNamespace(
|
|
name="list_models",
|
|
path_format="/v1/models",
|
|
methods={"GET"},
|
|
)
|
|
observed = {
|
|
"operation_id": _generate_stable_operation_id(route),
|
|
"name": route.name,
|
|
"path": route.path_format,
|
|
}
|
|
assert normalize(observed) == {
|
|
"operation_id": "list_models_v1_models_get",
|
|
"name": "list_models",
|
|
"path": "/v1/models",
|
|
}
|
|
|
|
|
|
def test_generate_stable_operation_id_multi_method_no_suffix():
|
|
route = SimpleNamespace(
|
|
name="multi_op",
|
|
path_format="/v1/things/{id}",
|
|
methods={"GET", "POST"},
|
|
)
|
|
observed = {
|
|
"operation_id": _generate_stable_operation_id(route),
|
|
"method_count": len(route.methods),
|
|
"has_method_suffix": _generate_stable_operation_id(route).endswith(
|
|
("_get", "_post")
|
|
),
|
|
}
|
|
assert normalize(observed) == {
|
|
"operation_id": "multi_op_v1_things__id_",
|
|
"method_count": 2,
|
|
"has_method_suffix": False,
|
|
}
|
|
|
|
|
|
def test_generate_stable_operation_id_missing_attrs_raises_error():
|
|
bad_route = SimpleNamespace() # missing name/path_format/methods
|
|
with pytest.raises(AttributeError):
|
|
_generate_stable_operation_id(bad_route)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _strip_operation_id_method_suffix
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_strip_operation_id_method_suffix_removes_known_method():
|
|
observed = {
|
|
"with_get": _strip_operation_id_method_suffix("list_models_v1_models_get"),
|
|
"with_post": _strip_operation_id_method_suffix("create_thing_post"),
|
|
"with_delete": _strip_operation_id_method_suffix("drop_thing_delete"),
|
|
}
|
|
assert observed == {
|
|
"with_get": "list_models_v1_models",
|
|
"with_post": "create_thing",
|
|
"with_delete": "drop_thing",
|
|
}
|
|
|
|
|
|
def test_strip_operation_id_method_suffix_invalid_suffix_unchanged():
|
|
# "foo" is not a known HTTP method; "nounderscore" has no separator at all.
|
|
observed = {
|
|
"unknown_suffix": _strip_operation_id_method_suffix("operation_foo"),
|
|
"no_underscore": _strip_operation_id_method_suffix("nounderscore"),
|
|
"empty": _strip_operation_id_method_suffix(""),
|
|
}
|
|
assert observed == {
|
|
"unknown_suffix": "operation_foo",
|
|
"no_underscore": "nounderscore",
|
|
"empty": "",
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ensure_unique_openapi_operation_ids
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_ensure_unique_openapi_operation_ids_rewrites_duplicates():
|
|
schema = {
|
|
"paths": {
|
|
"/a": {"get": {"operationId": "dup_get"}},
|
|
"/b": {"get": {"operationId": "dup_get"}},
|
|
"/c": {"post": {"operationId": "unique_post"}},
|
|
}
|
|
}
|
|
result = ensure_unique_openapi_operation_ids(schema)
|
|
observed = {
|
|
"a_get": result["paths"]["/a"]["get"]["operationId"],
|
|
"b_get": result["paths"]["/b"]["get"]["operationId"],
|
|
"c_post": result["paths"]["/c"]["post"]["operationId"],
|
|
"ids_are_distinct": len(
|
|
{
|
|
result["paths"]["/a"]["get"]["operationId"],
|
|
result["paths"]["/b"]["get"]["operationId"],
|
|
result["paths"]["/c"]["post"]["operationId"],
|
|
}
|
|
)
|
|
== 3,
|
|
}
|
|
assert normalize(observed) == {
|
|
"a_get": "dup_get",
|
|
"b_get": "dup_get_2",
|
|
"c_post": "unique_post",
|
|
"ids_are_distinct": True,
|
|
}
|
|
|
|
|
|
def test_ensure_unique_openapi_operation_ids_respects_reserved():
|
|
# operationId already ends with "_get" (an HTTP method), so the suffix is
|
|
# stripped before re-appending the current method, yielding "reserved_get".
|
|
schema = {
|
|
"paths": {
|
|
"/a": {"get": {"operationId": "reserved_get"}},
|
|
}
|
|
}
|
|
reserved = {"reserved_get"}
|
|
result = ensure_unique_openapi_operation_ids(
|
|
schema, reserved_operation_ids=reserved
|
|
)
|
|
observed = {
|
|
"rewritten": result["paths"]["/a"]["get"]["operationId"],
|
|
"still_includes_original": "reserved_get" in reserved,
|
|
"reserved_grew": len(reserved) > 1,
|
|
}
|
|
assert normalize(observed) == {
|
|
"rewritten": "reserved_get_2",
|
|
"still_includes_original": True,
|
|
"reserved_grew": True,
|
|
}
|
|
|
|
|
|
def test_ensure_unique_openapi_operation_ids_missing_paths_invalid_returns_empty():
|
|
"""No ``paths`` key — function must not crash and must return the schema as-is."""
|
|
schema = {"info": {"title": "x"}}
|
|
result = ensure_unique_openapi_operation_ids(schema)
|
|
assert result is schema
|
|
assert "paths" not in result
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _inject_websocket_stubs_into_openapi_schema
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_inject_websocket_stubs_into_openapi_schema_adds_stub():
|
|
schema = {"paths": {}}
|
|
route = SimpleNamespace(path="/ws/chat", name="ws_chat", dependant=None)
|
|
result = _inject_websocket_stubs_into_openapi_schema(schema, [route])
|
|
stub = result["paths"]["/ws/chat"]["get"]
|
|
assert normalize(stub) == {
|
|
"summary": "WebSocket: ws_chat",
|
|
"description": "WebSocket connection endpoint",
|
|
"operationId": "websocket_ws_chat",
|
|
"parameters": [],
|
|
"responses": {"101": {"description": "WebSocket Protocol Switched"}},
|
|
"tags": ["WebSocket"],
|
|
}
|
|
|
|
|
|
def test_inject_websocket_stubs_into_openapi_schema_does_not_overwrite_existing_get():
|
|
# Existing GET on the same path must not be replaced by the stub.
|
|
existing_get = {"summary": "real http get", "operationId": "real_get"}
|
|
schema = {"paths": {"/ws/chat": {"get": existing_get}}}
|
|
route = SimpleNamespace(path="/ws/chat", name="ws_chat", dependant=None)
|
|
result = _inject_websocket_stubs_into_openapi_schema(schema, [route])
|
|
assert result["paths"]["/ws/chat"]["get"] is existing_get
|
|
|
|
|
|
def test_inject_websocket_stubs_into_openapi_schema_missing_paths_key_raises_error():
|
|
schema = {} # no "paths" key — setdefault on missing schema["paths"] will KeyError
|
|
route = SimpleNamespace(path="/ws/x", name="ws_x", dependant=None)
|
|
with pytest.raises(KeyError):
|
|
_inject_websocket_stubs_into_openapi_schema(schema, [route])
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# get_openapi_schema
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_get_openapi_schema_returns_well_formed_schema(monkeypatch):
|
|
"""Patch ps.app to a fresh FastAPI so we get a deterministic minimal schema
|
|
without depending on whatever the session app currently has cached."""
|
|
import litellm.proxy.proxy_server as ps
|
|
|
|
fresh = FastAPI(title="pinned-title", version="0.0.1")
|
|
|
|
@fresh.get("/ping")
|
|
def _ping():
|
|
return {"ok": True}
|
|
|
|
monkeypatch.setattr(ps, "app", fresh, raising=True)
|
|
schema = get_openapi_schema()
|
|
observed = {
|
|
"openapi_present": "openapi" in schema,
|
|
"has_paths": isinstance(schema.get("paths"), dict),
|
|
"has_info": isinstance(schema.get("info"), dict),
|
|
"title": schema["info"]["title"],
|
|
"ping_path_in_schema": "/ping" in schema["paths"],
|
|
}
|
|
assert normalize(observed) == {
|
|
"openapi_present": True,
|
|
"has_paths": True,
|
|
"has_info": True,
|
|
"title": "pinned-title",
|
|
"ping_path_in_schema": True,
|
|
}
|
|
|
|
|
|
def test_get_openapi_schema_returns_cached_when_present(monkeypatch):
|
|
"""When the patched app already has openapi_schema set, the function
|
|
returns it untouched (no regeneration)."""
|
|
import litellm.proxy.proxy_server as ps
|
|
|
|
fresh = FastAPI()
|
|
sentinel = {"openapi": "3.0.0", "paths": {}, "info": {"title": "cached"}}
|
|
fresh.openapi_schema = sentinel
|
|
monkeypatch.setattr(ps, "app", fresh, raising=True)
|
|
result = get_openapi_schema()
|
|
observed = {
|
|
"is_sentinel": result is sentinel,
|
|
"title": result["info"]["title"],
|
|
"paths_empty": result["paths"] == {},
|
|
}
|
|
assert normalize(observed) == {
|
|
"is_sentinel": True,
|
|
"title": "cached",
|
|
"paths_empty": True,
|
|
}
|
|
|
|
|
|
def test_get_openapi_schema_missing_app_attribute_raises_error(monkeypatch):
|
|
"""If the module-level ``app`` is replaced by something without
|
|
``openapi_schema`` and without ``routes``, the function fails fast."""
|
|
import litellm.proxy.proxy_server as ps
|
|
|
|
monkeypatch.setattr(ps, "app", SimpleNamespace(), raising=True)
|
|
with pytest.raises(AttributeError):
|
|
get_openapi_schema()
|
|
|
|
|
|
async def test_a_request_is_answered_while_a_cold_openapi_json_builds(monkeypatch: pytest.MonkeyPatch):
|
|
build_started: Final = threading.Event()
|
|
liveliness_answered: Final = threading.Event()
|
|
build_saw_the_answer: Final = threading.Event()
|
|
|
|
class HoldsTheDocsBuild(BaseModel):
|
|
model_config = ConfigDict(frozen=True)
|
|
|
|
@classmethod
|
|
def __get_pydantic_json_schema__(cls, schema: CoreSchema, handler: GetJsonSchemaHandler) -> JsonSchemaValue:
|
|
build_started.set()
|
|
if liveliness_answered.wait(UNBLOCK_TIMEOUT_SECONDS):
|
|
build_saw_the_answer.set()
|
|
return handler(schema)
|
|
|
|
import litellm.proxy.proxy_server as ps
|
|
|
|
proxy_app: Final = ps.app
|
|
monkeypatch.setattr(proxy_app, "openapi_schema", None)
|
|
monkeypatch.setattr(proxy_app.router, "routes", [*proxy_app.router.routes])
|
|
proxy_app.add_api_route("/holds-the-docs-build", lambda: None, response_model=HoldsTheDocsBuild)
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=proxy_app), base_url="http://proxy") as client:
|
|
docs_request: Final = asyncio.ensure_future(client.get("/openapi.json"))
|
|
await asyncio.to_thread(build_started.wait, UNBLOCK_TIMEOUT_SECONDS)
|
|
liveliness: Final = await client.get("/health/liveliness")
|
|
liveliness_answered.set()
|
|
docs: Final = await docs_request
|
|
|
|
assert (liveliness.status_code, docs.status_code, build_saw_the_answer.is_set()) == (200, 200, True)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# custom_openapi
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_custom_openapi_filters_to_openai_routes(monkeypatch):
|
|
"""custom_openapi() filters paths down to the OpenAI-compatible set and
|
|
caches the result on the patched app."""
|
|
import litellm.proxy.proxy_server as ps
|
|
|
|
fresh = FastAPI(title="pinned-custom", version="0.0.1")
|
|
|
|
@fresh.get("/ping")
|
|
def _ping():
|
|
return {"ok": True}
|
|
|
|
monkeypatch.setattr(ps, "app", fresh, raising=True)
|
|
schema = custom_openapi()
|
|
observed = {
|
|
"openapi_present": "openapi" in schema,
|
|
"paths_is_dict": isinstance(schema.get("paths"), dict),
|
|
"info_title": schema["info"]["title"],
|
|
"cached_now": fresh.openapi_schema is schema,
|
|
"non_openai_path_filtered": "/ping" not in schema["paths"],
|
|
}
|
|
assert normalize(observed) == {
|
|
"openapi_present": True,
|
|
"paths_is_dict": True,
|
|
"info_title": "pinned-custom",
|
|
"cached_now": True,
|
|
"non_openai_path_filtered": True,
|
|
}
|
|
|
|
|
|
def test_custom_openapi_returns_cached_when_present(monkeypatch):
|
|
import litellm.proxy.proxy_server as ps
|
|
|
|
fresh = FastAPI()
|
|
sentinel = {"openapi": "3.0.0", "paths": {}, "info": {"title": "cached"}}
|
|
fresh.openapi_schema = sentinel
|
|
monkeypatch.setattr(ps, "app", fresh, raising=True)
|
|
result = custom_openapi()
|
|
observed = {
|
|
"is_sentinel": result is sentinel,
|
|
"title": result["info"]["title"],
|
|
"paths_empty": result["paths"] == {},
|
|
}
|
|
assert normalize(observed) == {
|
|
"is_sentinel": True,
|
|
"title": "cached",
|
|
"paths_empty": True,
|
|
}
|
|
|
|
|
|
def test_custom_openapi_missing_app_attribute_raises_error(monkeypatch):
|
|
import litellm.proxy.proxy_server as ps
|
|
|
|
monkeypatch.setattr(ps, "app", SimpleNamespace(), raising=True)
|
|
with pytest.raises(AttributeError):
|
|
custom_openapi()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# mount_swagger_ui
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_mount_swagger_ui_mounts_static_route(monkeypatch):
|
|
"""mount_swagger_ui mutates the global app — patch the module's `app` to a
|
|
fresh FastAPI() so we don't pollute the session app's mount table."""
|
|
import litellm.proxy.proxy_server as ps
|
|
from fastapi import applications as fa_applications
|
|
|
|
fresh_app = FastAPI()
|
|
monkeypatch.setattr(ps, "app", fresh_app, raising=True)
|
|
original_get_swagger = fa_applications.get_swagger_ui_html
|
|
|
|
try:
|
|
mount_swagger_ui()
|
|
finally:
|
|
# Restore the swagger monkey-patch so other tests are unaffected.
|
|
fa_applications.get_swagger_ui_html = original_get_swagger
|
|
|
|
mount_names = [getattr(r, "name", None) for r in fresh_app.routes]
|
|
observed = {
|
|
"swagger_mounted": "swagger" in mount_names,
|
|
"patched_get_swagger": (
|
|
fa_applications.get_swagger_ui_html is original_get_swagger
|
|
),
|
|
"route_count_positive": len(fresh_app.routes) > 0,
|
|
}
|
|
assert normalize(observed) == {
|
|
"swagger_mounted": True,
|
|
"patched_get_swagger": True,
|
|
"route_count_positive": True,
|
|
}
|
|
|
|
|
|
def test_mount_swagger_ui_missing_directory_raises_error(monkeypatch, tmp_path):
|
|
"""If the swagger directory is missing, StaticFiles raises RuntimeError."""
|
|
import litellm.proxy.proxy_server as ps
|
|
from fastapi import applications as fa_applications
|
|
|
|
fresh_app = FastAPI()
|
|
monkeypatch.setattr(ps, "app", fresh_app, raising=True)
|
|
monkeypatch.setattr(
|
|
ps, "current_dir", str(tmp_path / "does_not_exist"), raising=True
|
|
)
|
|
original_get_swagger = fa_applications.get_swagger_ui_html
|
|
|
|
try:
|
|
with pytest.raises(RuntimeError):
|
|
mount_swagger_ui()
|
|
finally:
|
|
fa_applications.get_swagger_ui_html = original_get_swagger
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _get_cors_config
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_get_cors_config_explicit_origins_and_credentials():
|
|
origins, allow_creds = _get_cors_config(
|
|
cors_origins_env="https://a.example,https://b.example",
|
|
cors_credentials_env="true",
|
|
)
|
|
observed = {
|
|
"origins": origins,
|
|
"allow_credentials": allow_creds,
|
|
"origin_count": len(origins),
|
|
}
|
|
assert normalize(observed) == {
|
|
"origins": ["https://a.example", "https://b.example"],
|
|
"allow_credentials": True,
|
|
"origin_count": 2,
|
|
}
|
|
|
|
|
|
def test_get_cors_config_wildcard_defaults_credentials_false(monkeypatch):
|
|
# Clear env to ensure we test the default branch deterministically.
|
|
monkeypatch.delenv("LITELLM_CORS_ORIGINS", raising=False)
|
|
monkeypatch.delenv("LITELLM_CORS_ALLOW_CREDENTIALS", raising=False)
|
|
origins, allow_creds = _get_cors_config()
|
|
observed = {
|
|
"origins": origins,
|
|
"allow_credentials": allow_creds,
|
|
"wildcard_in_origins": "*" in origins,
|
|
}
|
|
assert normalize(observed) == {
|
|
"origins": ["*"],
|
|
"allow_credentials": False,
|
|
"wildcard_in_origins": True,
|
|
}
|
|
|
|
|
|
def test_get_cors_config_invalid_credentials_value_treated_as_false():
|
|
"""Anything other than the literal "true" (case-insensitive) is false —
|
|
misconfigured strings should not silently enable credentialed CORS."""
|
|
_, allow_creds = _get_cors_config(
|
|
cors_origins_env="https://a.example",
|
|
cors_credentials_env="yes-please",
|
|
)
|
|
assert allow_creds is False
|