scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

refactor(tests): consolidate rust bridge route coverage

Browse source
This commit is contained in:
Yujong Lee 2026-09-04 13:47:26 -07:00
parent 1e3d2e7f6b
commit e7b8c3140f
4 changed files with 163 additions and 115 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

243
tests/test_litellm/rust_bridge/native_route_wheel_test.py
View file

@ -10,6 +10,8 @@ import sys
import tempfile
import threading
import zipfile
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from http.client import HTTPMessage
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
@ -27,6 +29,123 @@ ANTHROPIC_RESPONSE: Final = (
)
def assert_ocr_request(path: str, headers: HTTPMessage, body: dict[str, object]) -> None:
assert path == "/v1/ocr"
assert headers.get("authorization") == "Bearer sk-native"
assert body == {
"model": "mistral-ocr-latest",
"document": {"type": "document_url", "document_url": "https://example.com/document.pdf"},
"include_image_base64": True,
}
def assert_transcription_request(path: str, headers: HTTPMessage, body: dict[str, object]) -> None:
assert path == "/model/mistral.voxtral-mini-3b-2507/converse"
assert headers.get("authorization", "").startswith("AWS4-HMAC-SHA256 ")
assert headers.get("x-amz-date")
messages: Final = body["messages"]
assert isinstance(messages, list)
assert messages[0]["content"][0]["audio"]["source"]["bytes"] == "AQI="
assert "The audio language is en" in messages[0]["content"][1]["text"]
def assert_messages_request(path: str, headers: HTTPMessage, body: dict[str, object]) -> None:
assert path == "/v1/messages"
assert headers.get("x-api-key") == "sk-native"
assert body == {
"model": "claude-sonnet-4-5",
"max_tokens": 16,
"messages": [{"role": "user", "content": "hello-from-messages"}],
}
def assert_chat_completions_request(path: str, headers: HTTPMessage, body: dict[str, object]) -> None:
assert path == "/v1/messages"
assert headers.get("x-api-key") == "sk-native"
assert body == {
"model": "claude-sonnet-4-5",
"max_tokens": 17,
"messages": [{"role": "user", "content": [{"type": "text", "text": "hello-from-chat"}]}],
}
@dataclass(frozen=True, slots=True)
class RouteCase:
name: str
request_kwargs: Mapping[str, object]
assert_request: Callable[[str, HTTPMessage, dict[str, object]], None]
success_value: Callable[[dict[str, object]], object]
expected_success: str
native_upstream_error: bool
ROUTE_CASES: Final = (
RouteCase(
name="ocr",
request_kwargs={
"model": "mistral-ocr-latest",
"document": {"type": "document_url", "document_url": "https://example.com/document.pdf"},
"api_key": "sk-native",
"custom_llm_provider": "mistral",
"optional_params": {"include_image_base64": True},
},
assert_request=assert_ocr_request,
success_value=lambda response: response["pages"][0]["markdown"],
expected_success="native-ocr",
native_upstream_error=True,
),
RouteCase(
name="transcription",
request_kwargs={
"model": "mistral.voxtral-mini-3b-2507",
"audio": {"data": "AQI=", "format": "wav", "filename": "audio.wav"},
"custom_llm_provider": "bedrock",
"optional_params": {
"aws_access_key_id": "native-access-key",
"aws_secret_access_key": "native-secret-key",
"aws_region_name": "us-east-1",
"language": "en",
},
},
assert_request=assert_transcription_request,
success_value=lambda response: response["text"],
expected_success="native-transcription",
native_upstream_error=False,
),
RouteCase(
name="messages",
request_kwargs={
"model": "claude-sonnet-4-5",
"body": {
"model": "claude-sonnet-4-5",
"max_tokens": 16,
"messages": [{"role": "user", "content": "hello-from-messages"}],
},
"api_key": "sk-native",
"custom_llm_provider": "anthropic",
},
assert_request=assert_messages_request,
success_value=lambda response: response["content"][0]["text"],
expected_success="native-message",
native_upstream_error=False,
),
RouteCase(
name="chat_completions",
request_kwargs={
"model": "anthropic/claude-sonnet-4-5",
"messages": [{"role": "user", "content": "hello-from-chat"}],
"optional_params": {"max_tokens": 17},
"api_key": "sk-native",
},
assert_request=assert_chat_completions_request,
success_value=lambda response: response["choices"][0]["message"]["content"],
expected_success="native-message",
native_upstream_error=True,
),
)
ROUTE_CASE_BY_NAME: Final = {case.name: case for case in ROUTE_CASES}
class NativeRouteHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
@ -73,35 +192,14 @@ def assert_native_request(
headers: HTTPMessage,
body: object,
) -> None:
if route not in {"ocr", "transcription", "messages", "chat_completions"}:
case: Final = ROUTE_CASE_BY_NAME.get(route or "")
if case is None:
raise AssertionError(f"unexpected route marker: {route!r}")
if outcome not in {"success", "429", "hang"}:
raise AssertionError(f"unexpected outcome marker: {outcome!r}")
if not isinstance(body, dict):
raise TypeError(f"{route} sent {type(body).__name__}, expected a JSON object")
if route == "ocr":
assert path == "/v1/ocr"
assert headers.get("authorization") == "Bearer sk-native"
assert body["model"] == "mistral-ocr-latest"
assert body["document"]["document_url"] == "https://example.com/document.pdf"
assert body["include_image_base64"] is True
return
if route == "transcription":
assert path == "/model/mistral.voxtral-mini-3b-2507/converse"
assert headers.get("authorization", "").startswith("AWS4-HMAC-SHA256 ")
assert headers.get("x-amz-date")
assert body["messages"][0]["content"][0]["audio"]["source"]["bytes"] == "AQI="
assert "The audio language is en" in body["messages"][0]["content"][1]["text"]
return
assert path == "/v1/messages"
assert headers.get("x-api-key") == "sk-native"
assert body["model"] == "claude-sonnet-4-5"
if route == "messages":
assert body["max_tokens"] == 16
assert body["messages"][0]["content"] == "hello-from-messages"
return
assert body["max_tokens"] == 17
assert body["messages"][0]["content"] == [{"type": "text", "text": "hello-from-chat"}]
case.assert_request(path, headers, body)
def native_response(status: int, route: str | None) -> bytes:
@ -124,75 +222,29 @@ def load_native(native_path: Path) -> object:
def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]:
case: Final = ROUTE_CASE_BY_NAME.get(route)
if case is None:
raise AssertionError(f"unknown route: {route}")
common: Final = {
"api_base": api_base,
"extra_headers": {"x-test-outcome": outcome, "x-test-route": route},
"timeout_seconds": 3.0,
}
if route == "ocr":
return common | {
"model": "mistral-ocr-latest",
"document": {"type": "document_url", "document_url": "https://example.com/document.pdf"},
"api_key": "sk-native",
"custom_llm_provider": "mistral",
"optional_params": {"include_image_base64": True},
}
if route == "transcription":
return common | {
"model": "mistral.voxtral-mini-3b-2507",
"audio": {"data": "AQI=", "format": "wav", "filename": "audio.wav"},
"custom_llm_provider": "bedrock",
"optional_params": {
"aws_access_key_id": "native-access-key",
"aws_secret_access_key": "native-secret-key",
"aws_region_name": "us-east-1",
"language": "en",
},
}
if route == "messages":
return common | {
"model": "claude-sonnet-4-5",
"body": {
"model": "claude-sonnet-4-5",
"max_tokens": 16,
"messages": [{"role": "user", "content": "hello-from-messages"}],
},
"api_key": "sk-native",
"custom_llm_provider": "anthropic",
}
if route == "chat_completions":
return common | {
"model": "anthropic/claude-sonnet-4-5",
"messages": [{"role": "user", "content": "hello-from-chat"}],
"optional_params": {"max_tokens": 17},
"api_key": "sk-native",
}
raise AssertionError(f"unknown route: {route}")
return common | dict(case.request_kwargs)
def assert_success(route: str, response: object) -> None:
case: Final = ROUTE_CASE_BY_NAME[route]
if not isinstance(response, dict):
raise TypeError(f"{route} returned {type(response).__name__}, expected dict")
actual: Final = success_value(route, response)
expected: Final = (
"native-ocr" if route == "ocr" else "native-transcription" if route == "transcription" else "native-message"
)
if actual != expected:
raise AssertionError(f"{route} returned {actual!r}, expected {expected!r}")
def success_value(route: str, response: dict[object, object]) -> object:
if route == "ocr":
return response["pages"][0]["markdown"]
if route == "transcription":
return response["text"]
if route == "messages":
return response["content"][0]["text"]
return response["choices"][0]["message"]["content"]
actual: Final = case.success_value(response)
if actual != case.expected_success:
raise AssertionError(f"{route} returned {actual!r}, expected {case.expected_success!r}")
def assert_rate_limit(native: object, route: str, error: BaseException) -> None:
if route in {"ocr", "chat_completions"}:
case: Final = ROUTE_CASE_BY_NAME[route]
if case.native_upstream_error:
upstream_error: Final = native.RustUpstreamError
if not isinstance(error, upstream_error) or error.args[0] != 429:
raise AssertionError(f"{route} returned the wrong 429 error: {error!r}")
@ -202,37 +254,32 @@ def assert_rate_limit(native: object, route: str, error: BaseException) -> None:
def exercise_sync(native: object, api_base: str) -> None:
for route in ("ocr", "transcription", "messages", "chat_completions"):
function: Final = getattr(native, route)
assert_success(route, function(**route_kwargs(route, api_base, "success")))
for case in ROUTE_CASES:
function: Final = getattr(native, case.name)
assert_success(case.name, function(**route_kwargs(case.name, api_base, "success")))
try:
function(**route_kwargs(route, api_base, "429"))
function(**route_kwargs(case.name, api_base, "429"))
except (RuntimeError, native.RustUpstreamError) as error:
assert_rate_limit(native, route, error)
assert_rate_limit(native, case.name, error)
else:
raise AssertionError(f"{route} accepted a 429 response")
raise AssertionError(f"{case.name} accepted a 429 response")
async def exercise_async(native: object, api_base: str) -> None:
for route in ("ocr", "transcription", "messages", "chat_completions"):
function: Final = getattr(native, f"a{route}")
assert_success(route, await function(**route_kwargs(route, api_base, "success")))
for case in ROUTE_CASES:
function: Final = getattr(native, f"a{case.name}")
assert_success(case.name, await function(**route_kwargs(case.name, api_base, "success")))
try:
await function(**route_kwargs(route, api_base, "429"))
await function(**route_kwargs(case.name, api_base, "429"))
except (RuntimeError, native.RustUpstreamError) as error:
assert_rate_limit(native, route, error)
assert_rate_limit(native, case.name, error)
else:
raise AssertionError(f"a{route} accepted a 429 response")
raise AssertionError(f"a{case.name} accepted a 429 response")
async def exercise_async_concurrency(native: object, api_base: str) -> None:
responses: Final = await asyncio.wait_for(
asyncio.gather(
*(
native.amessages(**route_kwargs("messages", api_base, "success"))
for _ in range(32)
)
),
asyncio.gather(*(native.amessages(**route_kwargs("messages", api_base, "success")) for _ in range(32))),
timeout=15,
)
for response in responses:

19
tests/test_litellm/rust_bridge/test_chat_completions.py
View file

@ -163,11 +163,7 @@ def test_gate_forwards_enablement_to_preflight(
if enablement == "environment":
monkeypatch.setenv("LITELLM_RUST", "true")
litellm_params: Final = (
{"rust": True}
if enablement == "request-true"
else {"rust": False}
if enablement == "request-false"
else {}
{"rust": True} if enablement == "request-true" else {"rust": False} if enablement == "request-false" else {}
)
assert accepts(litellm_params=litellm_params) is expected
@ -218,11 +214,14 @@ def test_gate_short_circuits_python_only_requests(
bridge.set_rust_chat_completions(decline=gate)
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", bedrock_metadata_fields)
assert accepts(
custom_llm_provider=provider,
litellm_params=litellm_params,
stream=stream,
) is False
assert (
accepts(
custom_llm_provider=provider,
litellm_params=litellm_params,
stream=stream,
)
is False
)
assert gate.calls == []

13
tests/test_litellm/rust_bridge/test_runtime.py
View file

@ -288,11 +288,14 @@ def test_accepts_only_eligible_supported_requests(
bridge: Final = runtime.RustBridge(route="preflight", load=load, enabled=enabled)
assert bridge.accepts(
check=check,
request_override=False if state == "disabled" else None,
eligible=state != "ineligible",
) is expected
assert (
bridge.accepts(
check=check,
request_override=False if state == "disabled" else None,
eligible=state != "ineligible",
)
is expected
)
assert tuple(events) == expected_events

3
tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py
View file

@ -176,8 +176,7 @@ class _InvalidWheelCase:
_InvalidWheelCase(
metadata_tags=("cp312-cp312-linux_x86_64",),
expected_error=(
"WHEEL tags do not match filename: expected cp310-abi3-linux_x86_64, "
"found cp312-cp312-linux_x86_64"
"WHEEL tags do not match filename: expected cp310-abi3-linux_x86_64, found cp312-cp312-linux_x86_64"
),
),
id="metadata-tag-mismatched",