diff --git a/tests/test_litellm_rust/README.md b/tests/test_litellm_rust/README.md index 389543c178f..4c117fb846b 100644 --- a/tests/test_litellm_rust/README.md +++ b/tests/test_litellm_rust/README.md @@ -4,10 +4,10 @@ This suite covers OCR requests through LiteLLM's compiled Rust extension. OCR be A test name identifies the OCR entrypoint or callback under test and its expected observable result. Parameter IDs state the execution mode or credential case. Keep multiple assertions together only when they prove one request, mutation, failure, or callback lifecycle behavior. Record callback observations and assert them after the callback returns because production logging can swallow callback exceptions -`ocr/test_requests.py` covers provider payloads, file preparation, endpoint and credential resolution, normalized responses, errors, timeouts, and Azure token-provider behavior. `ocr/test_callbacks.py` covers OCR callback inputs, mutations, ordering, context, failure handling, concurrency, and cleanup. `ocr/test_guardrails.py` covers OCR post-call blocking and response replacement. These contract modules call the Rust bridge directly. `test_ocr.py` has the single public API dispatch test, covering enabled native dispatch and disabled Python dispatch +`ocr/test_requests.py` covers provider payloads, file preparation, endpoint and credential resolution, normalized responses, errors, timeouts, and Azure token-provider behavior. `ocr/test_callbacks.py` covers OCR callback inputs, mutations, ordering, context, failure handling, concurrency, and cleanup. `ocr/test_guardrails.py` covers OCR post-call blocking and response replacement. These contract modules call the Rust bridge directly. `ocr/test_dispatch.py` has the single public API dispatch test, covering enabled native dispatch and disabled Python dispatch. `test_ocr.py` is a strict smoke test of the compiled Rust OCR transport Run `make test-rust-extension` as the acceptance command. It builds a fresh wheel, installs that wheel into a temporary environment, requires `LITELLM_RUST=1`, and runs this suite with isolated Python imports -Collection fails when `LITELLM_RUST=1` is set but the compiled `_native` module cannot be imported. The autouse fixture isolates callback and configuration state but does not select a backend. Native contract tests call `litellm.rust_bridge.ocr` directly, while the strict dispatch test explicitly enables and disables Rust and verifies the selected transport at the wire boundary +Collection fails when `LITELLM_RUST=1` is set but the compiled `_native` module cannot be imported. The autouse fixture isolates callback and configuration state but does not select a backend. Native contract tests call `litellm.rust_bridge.ocr` directly, while the strict dispatch test explicitly enables and disables Rust and records which OCR entrypoint runs The OCR contract modules are non-strict expected failures until the retained callback implementation from #40070 lands. The public dispatch test remains strict. Passing contract cases appear as XPASS so staging coverage stays visible diff --git a/tests/test_litellm_rust/ocr/test_dispatch.py b/tests/test_litellm_rust/ocr/test_dispatch.py new file mode 100644 index 00000000000..a6c76bc5d0e --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_dispatch.py @@ -0,0 +1,44 @@ +from typing import Final +from unittest.mock import Mock + +import pytest + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.ocr import main as ocr_main +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.requests import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.fixture +def ocr_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) + return recording_server + + +@pytest.mark.parametrize("rust_enabled", [True, False], ids=["enabled", "disabled"]) +def test_public_ocr_dispatches_according_to_rust_setting( + ocr_server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + rust_enabled: bool, +) -> None: + rust_call: Final = Mock(wraps=ocr_main.rust_ocr_bridge.ocr) + python_call: Final = Mock(wraps=ocr_main.base_llm_http_handler.ocr) + monkeypatch.setattr(ocr_main.rust_ocr_bridge, "ocr", rust_call) + monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", python_call) + litellm.rust(rust_enabled) + + response: Final = litellm.ocr( + model=OCR_MODEL, + document=OCR_DOCUMENT, + api_key="test-key", + api_base=ocr_server.base_url, + ) + + assert isinstance(response, OCRResponse) + assert response.pages[0].markdown == "native OCR response" + assert rust_call.call_count == int(rust_enabled) + assert python_call.call_count == int(not rust_enabled) + assert len(ocr_server.requests) == 1 diff --git a/tests/test_litellm_rust/support/response_marker.py b/tests/test_litellm_rust/support/response_marker.py deleted file mode 100644 index bffda20e580..00000000000 --- a/tests/test_litellm_rust/support/response_marker.py +++ /dev/null @@ -1,4 +0,0 @@ -def has_rust_response_marker(response: object) -> bool: - from litellm.rust_bridge.provenance import has_rust_response_marker as implementation - - return implementation(response) diff --git a/tests/test_litellm_rust/test_ocr.py b/tests/test_litellm_rust/test_ocr.py index b5445c3a721..1293de9ee0e 100644 --- a/tests/test_litellm_rust/test_ocr.py +++ b/tests/test_litellm_rust/test_ocr.py @@ -7,14 +7,13 @@ from typing import Final import pytest -import litellm +from litellm.rust_bridge import ocr as rust_ocr_bridge pytestmark = pytest.mark.requires_rust_extension @dataclass(frozen=True, slots=True) class RecordedOCRRequest: - headers: dict[str, str] body: object @@ -26,7 +25,6 @@ def ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[RecordedOCRRequest def do_POST(self) -> None: requests.append( RecordedOCRRequest( - headers={name.lower(): value for name, value in self.headers.items()}, body=json.loads(self.rfile.read(int(self.headers["Content-Length"]))), ) ) @@ -57,26 +55,28 @@ def ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[RecordedOCRRequest thread.join() -@pytest.mark.parametrize("rust_enabled", [True, False], ids=["enabled", "disabled"]) -def test_public_ocr_dispatches_according_to_rust_setting( - ocr_server: tuple[ThreadingHTTPServer, list[RecordedOCRRequest]], rust_enabled: bool +def test_native_ocr_with_compiled_rust_extension( + ocr_server: tuple[ThreadingHTTPServer, list[RecordedOCRRequest]], ) -> None: - litellm.rust(rust_enabled) server, requests = ocr_server address: Final = server.server_address host: Final = str(address[0]) port: Final = int(address[1]) - response: Final = litellm.ocr( - model="mistral/mistral-ocr-latest", + response: Final = rust_ocr_bridge.ocr( + model="mistral-ocr-latest", document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, api_key="test-key", api_base=f"http://{host}:{port}", + custom_llm_provider="mistral", + extra_headers=None, + optional_params={}, + timeout=None, ) - assert response.pages[0].markdown == "native OCR response" + assert response is not None + assert response["pages"][0]["markdown"] == "native OCR response" assert len(requests) == 1 - assert ("user-agent" in requests[0].headers) == (not rust_enabled) assert requests[0].body == { "model": "mistral-ocr-latest", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},