fix(tests): narrow retained callback xfails

tests/_prometheus_helpers.py

@@ -5,6 +5,11 @@ from typing import Final
 from prometheus_client import REGISTRY, CollectorRegistry
 
 
+def clear_prometheus_registry() -> None:
+    for collector in tuple(REGISTRY._collector_to_names):  # pyright: ignore[reportPrivateUsage]  # prometheus_client has no public collector enumeration
+        REGISTRY.unregister(collector)
+
+
 @contextmanager
 def isolated_prometheus_registry(registry: CollectorRegistry = REGISTRY) -> Iterator[None]:
     original: Final = tuple(registry._collector_to_names)  # pyright: ignore[reportPrivateUsage]  # prometheus_client has no public collector enumeration

tests/test_litellm_rust/README.md

@@ -6,6 +6,6 @@ Use `backend` with `"python"` and `"rust"` for parity tests. Keep callback mutat
 
 Run the suite with `LITELLM_RUST=1 uv run pytest tests/test_litellm_rust`.
 
-The newly ported suite is marked as a non-strict expected failure until the retained callback implementation from #40070 lands. Passing cases appear as XPASS so staging coverage remains visible
+The retained callback contract modules are marked as non-strict expected failures until the implementation from #40070 lands. Harness tests remain strict, and passing contract cases appear as XPASS so staging coverage stays visible
 
 `isolated_backend` restores the previous backend override and callback state on exit, including after exceptions or cancellation. Scopes can nest in the same task. Run backend comparisons sequentially: overlapping scopes in different tasks raise before changing process-global state. Use separate worker processes for parallel backend comparisons

tests/test_litellm_rust/conftest.py

@@ -21,8 +21,8 @@ from litellm.rust_bridge.configuration import (  # pyright: ignore[reportPrivate
     _CONFIGURATION,
     _parse_env_bool,
 )
+from tests._prometheus_helpers import isolated_prometheus_registry
 from tests.test_litellm_rust.support.callback_recorder import drain_logging
-from tests.test_litellm_rust.support.prometheus import isolated_prometheus_registry
 from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service
 
 CALLBACK_ATTRIBUTES: Final = (
@@ -35,6 +35,20 @@ CALLBACK_ATTRIBUTES: Final = (
     "_async_failure_callback",
 )
 Backend = Literal["python", "rust"]
+EXPECTED_FAILURE_FILES: Final = frozenset(
+    {
+        "chat/test_callback_mutation.py",
+        "integrations/test_composition.py",
+        "integrations/test_exporters.py",
+        "integrations/test_guardrails.py",
+        "messages/test_callback_mutation.py",
+        "messages/test_streaming.py",
+        "ocr/test_callbacks.py",
+        "ocr/test_dispatch.py",
+        "ocr/test_requests.py",
+        "test_provenance.py",
+    }
+)
 
 
 class _BackendScopes:
@@ -143,16 +157,19 @@ def recording_server() -> Generator[RecordingServer]:
 
 
 def pytest_collection_modifyitems(items: list[pytest.Item]) -> None:
-    expected_failure = pytest.mark.xfail(
+    expected_failure: Final = pytest.mark.xfail(
         reason="requires the retained callback implementation from #40070",
         strict=False,
     )
     for item in items:
-        if "test_litellm_rust" in item.path.parts and item.path.name != "test_ocr.py":
+        if "test_litellm_rust" not in item.path.parts:
+            continue
+        relative_path: Final = "/".join(item.path.parts[item.path.parts.index("test_litellm_rust") + 1 :])
+        if relative_path in EXPECTED_FAILURE_FILES:
             item.add_marker(expected_failure)
 
     if not _parse_env_bool(os.environ.get("LITELLM_RUST")):
-        skip = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension")
+        skip: Final = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension")
         for item in items:
             if "test_litellm_rust" in item.path.parts:
                 item.add_marker(skip)

tests/test_litellm_rust/integrations/test_exporters.py

@@ -8,8 +8,8 @@ from prometheus_client import CollectorRegistry, Counter
 import litellm
 from litellm.integrations.prometheus import PrometheusLogger
 from litellm.litellm_core_utils import litellm_logging
+from tests._prometheus_helpers import isolated_prometheus_registry
 from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
-from tests.test_litellm_rust.support.prometheus import isolated_prometheus_registry
 from tests.test_litellm_rust.support.requests import MESSAGES_EVENTS
 from tests.test_litellm_rust.integrations import (
     ALL_ROUTES,

tests/test_litellm_rust/integrations/test_guardrails.py

@@ -174,6 +174,12 @@ def azure_text_moderation(server: RecordingServer) -> AzureContentSafetyTextMode
     )
 
 
+async def _observed_stream(source: AsyncIterator[object], first_chunk: asyncio.Event) -> AsyncIterator[object]:
+    async for chunk in source:
+        first_chunk.set()
+        yield chunk
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("backend", ("python", "rust"))
 async def test_stream_masking_reaches_client_loggers_and_exporters(
@@ -358,13 +364,7 @@ async def test_guarded_stream_cancel_drains_and_releases_roots(backend: Backend)
                 metadata={"retained": root},
             )
             first_chunk: Final = asyncio.Event()
-
-            async def observed_source() -> AsyncIterator[bytes]:
-                async for chunk in source:
-                    first_chunk.set()
-                    yield chunk
-
-            guarded = arm_guarded_stream(observed_source(), request_data, logger)
+            guarded = arm_guarded_stream(_observed_stream(source, first_chunk), request_data, logger)
             del root
             task: Final = asyncio.create_task(anext(guarded))
             await provider.wait_for_requests(1)

tests/test_litellm_rust/test_ocr.py

@@ -1,72 +0,0 @@
-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"},
-    }