test: add Rust extension pytest contract (#40181)

* test: add Rust extension pytest contract

* test: prove native OCR execution

* test: isolate Rust extension pytest collection

* ci: register Rust extension test coverage

* test: prove native OCR at wire boundary

.github/workflows/test-rust.yml

@@ -117,6 +117,9 @@ jobs:
 
       - run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
 
+      - name: Run pytest tests/test_litellm_rust with the compiled extension
+        run: make test-rust-extension
+
       - run: >-
           uv build --wheel --out-dir panic-dist
           --config-setting "maturin.build-args=--features panic-test,extension-module"

Makefile

@@ -4,6 +4,7 @@
 .PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc \
 	test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \
 	test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \
+	test-rust-extension \
 	info lint lint-inner lint-dev lint-checks format \
 	lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \
 	lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
@@ -54,6 +55,7 @@ help:
 	@echo "  make test-proxy-unit-b  - Run proxy_unit_tests (p-z, ~28 files)"
 	@echo "  make test-integration   - Run integration tests"
 	@echo "  make test-unit-helm     - Run helm unit tests"
+	@echo "  make test-rust-extension - Build the Rust extension and run its public Python tests"
 	@echo ""
 	@echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide"
 	@echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine."
@@ -289,6 +291,17 @@ pre-commit:
 	@$(MAKE) check
 
 # Testing targets
+test-rust-extension:
+	@temporary=$$(mktemp -d) && \
+	trap 'rm -rf "$$temporary"' EXIT HUP INT TERM && \
+	$(UV) build --python 3.12 --wheel --out-dir "$$temporary/wheels" && \
+	set -- "$$temporary"/wheels/*.whl && \
+	[ "$$#" -eq 1 ] && \
+	UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \
+	$(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \
+	LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
+	"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
+
 test: install-test-deps
 	$(UV_RUN) pytest tests/
 

pyproject.toml

@@ -340,6 +340,7 @@ markers = [
     "asyncio: mark test as an asyncio test",
     "limit_leaks: mark test with memory limit for leak detection (e.g., '40 MB')",
     "no_parallel: mark test to run sequentially (not in parallel) - typically for memory measurement tests",
+    "requires_rust_extension: public Python contract requiring an enabled, compiled Rust extension",
 ]
 filterwarnings = [
     # Suppress Pydantic serializer warnings from mock server responses (non-critical for memory tests)

tests/test_litellm_rust/conftest.py

@@ -0,0 +1,24 @@
+import os
+
+import pytest
+
+
+def pytest_collection_modifyitems(items):
+    rust_enabled = os.environ.get("LITELLM_RUST", "").strip().lower() in {
+        "1",
+        "true",
+        "yes",
+        "on",
+    }
+    if not rust_enabled:
+        skip = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension")
+        for item in items:
+            item.add_marker(skip)
+        return
+
+    try:
+        from litellm.rust_bridge import _native  # noqa: F401  # validates the installed extension
+    except ImportError as error:
+        raise pytest.UsageError(
+            "LITELLM_RUST=1 requires a compiled litellm.rust_bridge._native extension"
+        ) from error

tests/test_litellm_rust/test_ocr.py

@@ -0,0 +1,72 @@
+import json
+import threading
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+import pytest
+
+import litellm
+
+pytestmark = pytest.mark.requires_rust_extension
+
+
+@pytest.fixture
+def ocr_server():
+    requests = []
+
+    class Handler(BaseHTTPRequestHandler):
+        def do_POST(self):
+            requests.append(
+                {
+                    "headers": {name.lower(): value for name, value in self.headers.items()},
+                    "body": json.loads(self.rfile.read(int(self.headers["Content-Length"]))),
+                }
+            )
+            if self.headers.get("User-Agent", "").startswith("python-httpx"):
+                self.send_response(418)
+                self.end_headers()
+                return
+            response = json.dumps(
+                {
+                    "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}],
+                    "model": "mistral-ocr-latest",
+                    "usage_info": {"pages_processed": 1, "doc_size_bytes": 3},
+                }
+            ).encode()
+            self.send_response(200)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(response)))
+            self.end_headers()
+            self.wfile.write(response)
+
+        def log_message(self, format, *args):
+            pass
+
+    server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+    thread = threading.Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True)
+    thread.start()
+    try:
+        yield server, requests
+    finally:
+        server.shutdown()
+        server.server_close()
+        thread.join()
+
+
+def test_ocr_with_rust_extension(ocr_server):
+    server, requests = ocr_server
+    host, port = server.server_address
+
+    response = litellm.ocr(
+        model="mistral/mistral-ocr-latest",
+        document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
+        api_key="test-key",
+        api_base=f"http://{host}:{port}",
+    )
+
+    assert response.pages[0].markdown == "native OCR response"
+    assert len(requests) == 1
+    assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx")
+    assert requests[0]["body"] == {
+            "model": "mistral-ocr-latest",
+            "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
+    }