From 7afd9259ac022ddea0df93811f83a6718ce92d8c Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Sat, 18 Jul 2026 04:00:43 +0000
Subject: [PATCH] test(ocr): cover required Rust bridge behavior

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
---
 tests/test_litellm/ocr/test_rust_bridge.py | 29 ++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py
index 4172ab575eb..a05e1c065b5 100644
--- a/tests/test_litellm/ocr/test_rust_bridge.py
+++ b/tests/test_litellm/ocr/test_rust_bridge.py
@@ -301,6 +301,25 @@ def test_load_rust_aocr_returns_injected_impl():
     assert rust_bridge.load_rust_aocr() is bridge
 
 
+def test_use_litellm_rust_controls_injected_bridges(monkeypatch):
+    bridge = RecordingBridge()
+    async_bridge = RecordingAsyncBridge()
+
+    rust_bridge.use_litellm_rust(ocr=bridge, aocr=async_bridge)
+    assert rust_bridge.rust_ocr_enabled() is True
+    assert rust_bridge.load_rust_ocr() is bridge
+    assert rust_bridge.load_rust_aocr() is async_bridge
+
+    monkeypatch.setattr(
+        importlib.import_module("litellm.rust_bridge"),
+        "get_native_bridge",
+        lambda: None,
+    )
+    rust_bridge.use_litellm_rust(False)
+    assert rust_bridge.load_rust_ocr() is None
+    assert rust_bridge.load_rust_aocr() is None
+
+
 def test_bridge_injection_preserves_unspecified_impl():
     bridge = RecordingBridge()
     async_bridge = RecordingAsyncBridge()
@@ -627,6 +646,16 @@ def test_ocr_requires_rust_bridge_when_unavailable(monkeypatch):
     assert "Rust OCR bridge is required" in str(exc_info.value)
 
 
+@pytest.mark.asyncio
+async def test_aocr_requires_rust_bridge_when_unavailable(monkeypatch):
+    monkeypatch.setattr(ocr_main, "load_rust_aocr", lambda: None)
+
+    with pytest.raises(Exception) as exc_info:
+        await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
+
+    assert "Rust OCR bridge is required" in str(exc_info.value)
+
+
 RUST_OCR_ERROR_CASES = [
     pytest.param(400, litellm.BadRequestError, 400, id="400_bad_request"),
     pytest.param(401, litellm.AuthenticationError, 401, id="401_authentication"),