From e70e3d60349bb37cd4faed0bd8f21acc314c534d Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 23 Jun 2026 12:27:59 -0700 Subject: [PATCH] test(ocr): cover compiled-extension load and Python fallback paths Adds two tests so the Rust bridge module hits 100% and the ocr() fallback-to-Python branch is exercised: - load_rust_ocr() returning the compiled extension's ocr callable - ocr() degrading to the HTTP handler when no bridge is available --- tests/test_litellm/ocr/test_rust_bridge.py | 33 ++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 31443dec4bd..7e028064e4c 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -1,6 +1,8 @@ """Tests for the optional Rust-backed OCR path (``litellm/ocr/rust_bridge.py``).""" import importlib +import sys +import types import httpx import pytest @@ -135,6 +137,18 @@ def test_load_rust_ocr_none_when_extension_absent(): assert rust_bridge.load_rust_ocr() is None +def test_load_rust_ocr_uses_compiled_extension(monkeypatch): + """With no injected impl but a compiled ``litellm_python_bridge`` importable, + the loader returns the extension's ``ocr`` callable. The native wheel isn't + built in CI, so stand in a fake module via ``sys.modules``.""" + fake_module = types.ModuleType("litellm_python_bridge") + fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "litellm_python_bridge", fake_module) + + litellm.use_litellm_rust(True) # enabled, no impl injected -> import the extension + assert rust_bridge.load_rust_ocr() is fake_module.ocr + + def test_timeout_to_seconds_handles_float_timeout_and_none(): assert ocr_main._timeout_to_seconds(12.5) == 12.5 assert ocr_main._timeout_to_seconds(None) is None @@ -298,3 +312,22 @@ def test_ocr_does_not_route_to_rust_when_disabled(): # The impl stays available for injection, but the disabled flag gates usage, # so ocr() never reaches the Rust path (asserted via the enabled-path test). assert bridge.calls == [] + + +def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): + """Rust enabled but no bridge available (no injected impl, no compiled wheel): + ocr() must degrade to the Python HTTP handler instead of raising.""" + litellm.use_litellm_rust(True) # enabled, but load_rust_ocr() returns None in CI + + captured = {} + + def fake_handler_ocr(**kwargs): + captured["called"] = True + return OCRResponse(pages=[], model="mistral-ocr-latest", object="ocr") + + monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", fake_handler_ocr) + + response = litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + + assert captured.get("called") is True # Python path was used + assert isinstance(response, OCRResponse)