test(rust_bridge): cover binding validation, async upstream errors, and OCR preparation failures

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/test_litellm/ocr/test_main.py

@@ -257,3 +257,70 @@ def test_direct_ocr_call_bills_request_level_per_page_pricing() -> None:
     )
 
     assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.05 * 3)
+
+
+def _prepare(model: str, document: object, **kwargs: object) -> object:
+    return _prepare_ocr_request(
+        model=model,
+        document=document,  # pyright: ignore[reportArgumentType]  # exercises the runtime guard for untyped callers
+        api_key="test-key",
+        api_base=None,
+        timeout=None,
+        custom_llm_provider=None,
+        extra_headers=None,
+        kwargs={"litellm_logging_obj": Mock(), **kwargs},
+    )
+
+
+@pytest.mark.parametrize(
+    ("document", "match"),
+    (
+        ("https://example.com/file.pdf", "document must be a dict"),
+        ({"type": "video_url", "video_url": "https://example.com/clip.mp4"}, "Invalid document type: video_url"),
+    ),
+)
+def test_prepare_ocr_request_rejects_malformed_documents(document: object, match: str) -> None:
+    with pytest.raises(ValueError, match=match):
+        _prepare("mistral/mistral-ocr-latest", document)
+
+
+def test_prepare_ocr_request_rejects_provider_without_ocr_support() -> None:
+    with pytest.raises(ValueError, match="OCR is not supported for provider: openai"):
+        _prepare("openai/gpt-4o", dict(PRICING_DOCUMENT))
+
+
+@pytest.mark.parametrize(
+    ("request_format", "match"),
+    (("markdown", "Invalid `req_format`"), ("native", "`req_format='native'` is not supported")),
+)
+def test_prepare_ocr_request_rejects_unsupported_request_format(request_format: str, match: str) -> None:
+    with pytest.raises(litellm.UnsupportedParamsError, match=match):
+        _prepare("mistral/mistral-ocr-latest", dict(PRICING_DOCUMENT), req_format=request_format)
+
+
+@pytest.mark.asyncio
+async def test_python_none_provider_response_raises_public_error(
+    provider: Mock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    from litellm.ocr import main
+
+    monkeypatch.setattr(main.base_llm_http_handler, "ocr", Mock(return_value=None))
+
+    with pytest.raises(litellm.APIConnectionError, match="unexpected None response") as error:
+        await litellm.aocr(model="mistral/mistral-ocr-latest", document=dict(PRICING_DOCUMENT), api_key="test-key")
+    assert error.value.llm_provider == "mistral"
+    assert provider.call_count == 0
+
+
+@pytest.mark.parametrize(
+    ("model", "expected_provider"),
+    (("mistral-ocr-latest", "mistral"), ("azure_ai/doc-intelligence/prebuilt-layout", "azure_ai")),
+)
+def test_preparation_errors_map_to_public_exception_for_inferred_provider(
+    provider: Mock, model: str, expected_provider: str
+) -> None:
+    with pytest.raises(litellm.APIConnectionError) as error:
+        litellm.ocr(model=model, document="not-a-document")  # pyright: ignore[reportArgumentType]  # exercises the runtime guard
+    assert error.value.llm_provider == expected_provider
+    assert "document must be a dict" in str(error.value)
+    assert provider.call_count == 0

tests/test_litellm/rust_bridge/test_bindings.py

@@ -4,6 +4,11 @@ from typing import Final
 import pytest
 
 from litellm.rust_bridge import bindings
+from litellm.rust_bridge.chat_completions import entrypoints as chat_completions
+from litellm.rust_bridge.messages import entrypoints as messages
+from litellm.rust_bridge.ocr import entrypoints as ocr
+from litellm.rust_bridge.responses import entrypoints as responses
+from litellm.rust_bridge.transcription import native as transcription
 
 
 def test_binding_distinguishes_disable_from_reset(monkeypatch) -> None:
@@ -33,3 +38,36 @@ def test_binding_validates_native_attribute(
     binding: Final = bindings.NativeBinding("route", validate=lambda item: item if isinstance(item, int) else None)
 
     assert binding.load() == expected
+
+
+ROUTE_BINDINGS: Final = (
+    ("completion", chat_completions.NATIVE_COMPLETION),
+    ("acompletion", chat_completions.NATIVE_ACOMPLETION),
+    ("anthropic_messages_handler", messages.NATIVE_MESSAGES),
+    ("anthropic_messages", messages.NATIVE_AMESSAGES),
+    ("responses", responses.NATIVE_RESPONSES),
+    ("aresponses", responses.NATIVE_ARESPONSES),
+    ("ocr", ocr.NATIVE_OCR),
+    ("aocr", ocr.NATIVE_AOCR),
+    ("transcription", transcription.NATIVE_TRANSCRIPTION),
+    ("atranscription", transcription.NATIVE_ATRANSCRIPTION),
+)
+
+
+@pytest.mark.parametrize(
+    ("attribute", "route_binding"), ROUTE_BINDINGS, ids=[attribute for attribute, _ in ROUTE_BINDINGS]
+)
+def test_route_bindings_only_accept_callable_native_attributes(
+    monkeypatch: pytest.MonkeyPatch, attribute: str, route_binding: bindings.NativeBinding[object]
+) -> None:
+    def native_route() -> None:
+        pass
+
+    monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: "not callable"}))
+    route_binding.reset()
+    assert route_binding.load() is None
+
+    monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: native_route}))
+    route_binding.reset()
+    assert route_binding.load() is native_route
+    route_binding.reset()

tests/test_litellm/rust_bridge/test_dispatch.py

@@ -28,7 +28,9 @@ def test_route_without_rules_forwards_before_request_projection() -> None:
     def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request:
         pytest.fail("Python-only routes must not project the request")
 
-    dispatch: Final = PublicDispatch(route=Route.CHAT_COMPLETIONS, request=reject_request, context=lambda _: Context(Route.CHAT_COMPLETIONS))
+    dispatch: Final = PublicDispatch(
+        route=Route.CHAT_COMPLETIONS, request=reject_request, context=lambda _: Context(Route.CHAT_COMPLETIONS)
+    )
     result: Final = dispatch.run(
         ("model",),
         {"stream": True},
@@ -132,7 +134,9 @@ async def test_async_route_without_rules_preserves_async_iterator_result() -> No
     async def python(*args: object, **kwargs: object) -> AsyncGenerator[int, None]:  # kwargs-ok: pass-through shape
         return stream
 
-    dispatch: Final = PublicDispatch(route=Route.RESPONSES, request=reject_request, context=lambda _: Context(Route.RESPONSES))
+    dispatch: Final = PublicDispatch(
+        route=Route.RESPONSES, request=reject_request, context=lambda _: Context(Route.RESPONSES)
+    )
     result: Final = await dispatch.arun(
         ("model",),
         {"stream": True},
@@ -148,9 +152,7 @@ async def test_async_route_without_rules_preserves_async_iterator_result() -> No
 @pytest.mark.asyncio
 async def test_async_dispatch_accepts_websocket_style_none_result() -> None:
     request: Final = Request(model="realtime-model")
-    rules: Final[Rules] = (
-        Rule(Route.RESPONSES, Rollout.RUST_REQUIRED, deliveries=frozenset({Delivery.WEBSOCKET})),
-    )
+    rules: Final[Rules] = (Rule(Route.RESPONSES, Rollout.RUST_REQUIRED, deliveries=frozenset({Delivery.WEBSOCKET})),)
     dispatch: Final = PublicDispatch(
         route=Route.RESPONSES,
         request=lambda args, kwargs: request,
@@ -163,9 +165,9 @@ async def test_async_dispatch_accepts_websocket_style_none_result() -> None:
     async def native(request: Request, args: tuple[object, ...], kwargs: Mapping[str, object]) -> None:
         return None
 
-    native_binding: Final[NativeBinding[Callable[[Request, tuple[object, ...], Mapping[str, object]], Awaitable[None]]]] = NativeBinding(
-        "websocket", validate=lambda _: None
-    )
+    native_binding: Final[
+        NativeBinding[Callable[[Request, tuple[object, ...], Mapping[str, object]], Awaitable[None]]]
+    ] = NativeBinding("websocket", validate=lambda _: None)
     native_binding.override(native)
 
     result: Final = await dispatch.arun(
@@ -177,3 +179,51 @@ async def test_async_dispatch_accepts_websocket_style_none_result() -> None:
         rules=rules,
     )
     assert result is None
+
+
+def test_rules_for_other_routes_and_constrained_python_rules_skip_projection() -> None:
+    rules: Final[Rules] = (
+        Rule(Route.MESSAGES, Rollout.RUST_REQUIRED),
+        Rule(Route.OCR, Rollout.PYTHON_ONLY, providers=frozenset({"mistral"})),
+    )
+
+    def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request:
+        pytest.fail("Rules that cannot select Rust must not project the request")
+
+    dispatch: Final = PublicDispatch(route=Route.OCR, request=reject_request, context=lambda _: Context(Route.OCR))
+    expected: Final = object()
+    result: Final = dispatch.run(
+        ("model",),
+        {},
+        python=lambda *args, **kwargs: expected,
+        binding=binding(),
+        native=lambda hook, request, args, kwargs: pytest.fail("Rules that cannot select Rust must not call native"),
+        rules=rules,
+    )
+    assert result is expected
+
+
+@pytest.mark.asyncio
+async def test_async_bypass_forwards_to_python_without_native() -> None:
+    request: Final = Request(model="bypassed-model")
+    rules: Final[Rules] = (Rule(Route.RESPONSES, Rollout.RUST_REQUIRED),)
+    dispatch: Final = PublicDispatch(
+        route=Route.RESPONSES,
+        request=lambda args, kwargs: request,
+        context=lambda value: Context(Route.RESPONSES, model=value.model),
+        bypass=lambda value: value.model == "bypassed-model",
+    )
+    expected: Final = object()
+
+    async def python(*args: object, **kwargs: object) -> object:  # kwargs-ok: public pass-through shape
+        return expected
+
+    result: Final = await dispatch.arun(
+        ("bypassed-model",),
+        {},
+        python=python,
+        binding=binding(),
+        native=lambda hook, value, args, kwargs: pytest.fail("Bypassed requests must not call native"),
+        rules=rules,
+    )
+    assert result is expected

tests/test_litellm/rust_bridge/test_runtime.py

@@ -283,3 +283,26 @@ async def test_arun_required_route_rejects_unavailable_bridge() -> None:
             python=python,
             rules=rules(Rollout.RUST_REQUIRED),
         )
+
+
+@pytest.mark.asyncio
+async def test_arun_upstream_error_maps_to_api_error_without_fallback() -> None:
+    calls: Final = recorder(RustUpstreamError(503, "upstream unavailable"))
+
+    async def native(fn: NativeFn) -> str:
+        return fn()
+
+    async def python() -> str:
+        return calls.python()
+
+    with pytest.raises(APIError, match="upstream unavailable") as caught:
+        await runtime.arun(
+            CONTEXT,
+            binding=binding(calls.rust),
+            native=native,
+            python=python,
+            rules=rules(Rollout.RUST_OPT_OUT),
+        )
+
+    assert caught.value.status_code == 503
+    assert calls.calls == (RUST,)