From 7beafeebd8177eb8d5f9310c6b4bfcf13b6e7cfc Mon Sep 17 00:00:00 2001
From: Yujong Lee <yujong@berri.ai>
Date: Tue, 8 Sep 2026 10:00:37 -0700
Subject: [PATCH] test(rust): stabilize CI harnesses

---
 Makefile                                     | 5 +++--
 tests/test_litellm_rust/recording_server.py  | 8 +++++++-
 tests/test_litellm_rust/test_integrations.py | 1 +
 3 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/Makefile b/Makefile
index 7ca4585556e..e6e54e100a9 100644
--- a/Makefile
+++ b/Makefile
@@ -319,11 +319,12 @@ test-rust-ocr:
 		[ "$$#" -eq 1 ] || exit 1; \
 		wheel="$$1"; \
 	fi && \
-	UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --no-default-groups --group dev --extra proxy && \
+	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 "$$wheel" && \
 	LITELLM_LOCAL_MODEL_COST_MAP=True "$$temporary/venv/bin/python" -I -c 'from litellm.rust_bridge import _native; assert callable(_native.ocr) and callable(_native.aocr)' && \
 	LITELLM_REQUIRE_NATIVE_OCR=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
-	"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib tests/test_litellm/ocr/test_rust_bridge.py -v
+	"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib \
+	tests/test_litellm_rust/test_ocr.py tests/test_litellm_rust/test_ocr_callbacks.py -v
 
 test-rust-python: install-rust-python-test-deps
 	@python=$$($(UV_RUN) python -c 'import sys; print(sys.executable)') && \
diff --git a/tests/test_litellm_rust/recording_server.py b/tests/test_litellm_rust/recording_server.py
index 7ea20406cac..de22e908a88 100644
--- a/tests/test_litellm_rust/recording_server.py
+++ b/tests/test_litellm_rust/recording_server.py
@@ -1,5 +1,6 @@
-import json
+import asyncio
 import copy
+import json
 import threading
 import time
 from collections.abc import Iterator
@@ -44,6 +45,11 @@ class RecordingServer:
     def enqueue(self, response: ResponseSpec) -> None:
         self.responses.append(response)
 
+    async def wait_for_requests(self, count: int) -> None:
+        async with asyncio.timeout(2):
+            while len(self.requests) < count:
+                await asyncio.sleep(0.01)
+
 
 @contextmanager
 def recording_service() -> Iterator[RecordingServer]:
diff --git a/tests/test_litellm_rust/test_integrations.py b/tests/test_litellm_rust/test_integrations.py
index bb0be4dbd49..f77b91331e8 100644
--- a/tests/test_litellm_rust/test_integrations.py
+++ b/tests/test_litellm_rust/test_integrations.py
@@ -129,6 +129,7 @@ async def test_generic_guardrail_logging_only_verdict_is_exported_over_http(
     response: Final = await route.invoke(provider, callbacks=[guardrail, logger, recorder])
     await recorder.wait_for_async("async_log_success_event")
     await drain_logging()
+    await provider.wait_for_requests(3)
 
     scan: Final = provider.requests[1]
     assert scan.path == "/beta/litellm_basic_guardrail_api"