scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-27 06:13:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 72
litellm/tests/unit/proxy/proxy_server/test_openapi_customization.py
Mateo Wang 464932b1d1
refactor(typing): clear 1181 basedpyright Any errors across 431 files (#47789) 
* refactor(typing): narrow Any sources across providers, proxy and guardrails (wave 1)

* refactor(typing): narrow Any sources across hubs, providers and proxy (wave 2)

* refactor(typing): validate RunwayML task JSON and let dotprompt loaders accept exported prompts

* fix: restore pass-through behavior on paths the typing pass tightened

MCP client IP extraction trusts X-Forwarded-For only when use_x_forwarded_for is True, dotprompt skips a malformed JSON prompt entry instead of raising, RunwayML task payloads pass through unvalidated, and a malformed provider body raises the same AttributeError as before at the CompactifAI, Gemini, and mock transport parse sites

* fix(runwayml): keep a malformed 2xx task body on the 500 path

* fix(edenai,cli): keep a non-object body in the error and fail on a malformed migration report

EdenAI's embedding and image generation parse dropped a non-object body from the error
they raised, so the message no longer named what the provider sent. A helper now raises
with the body in the message, and both paths have a regression test.

`lite encryption migrate` exited 0 with no message when the proxy's report was not an
object or its residual_legacy was not a number, where base crashed. It now fails with a
message naming the value.

The CompactifAI and Gemini image-edit non-object body tests accept TypeError or
AttributeError, which both base and this branch raise, so they guard against a
ValidationError instead of pinning one class. RunwayML's async parse gets the same
non-object body test as the sync one.

* fix: keep base's error classes, duck-typed checks and fail-closed guardrails the typing pass changed

* test: annotate the new regression tests' locals with Final

* test: annotate the remaining new test locals with Final
2026-10-24 11:31:09 -07:00

500 lines
18 KiB
Python
Raw Blame History

"""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 tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder
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_documents_query_params_that_forward_their_fields():
query_param: Final = AttributeForwarder(SimpleNamespace(name="session_id", required=True))
route: Final = SimpleNamespace(path="/ws/chat", name="ws_chat", dependant=SimpleNamespace(query_params=[query_param]))
result: Final = _inject_websocket_stubs_into_openapi_schema({"paths": {}}, [route])
assert result["paths"]["/ws/chat"]["get"]["parameters"] == [
{"name": "session_id", "in": "query", "required": True, "schema": {"type": "string"}}
]
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
Reference in a new issue View git blame Copy permalink