From d701b818f9806e89e6feb285dc43fdec89abc6d2 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 9 Sep 2026 08:04:03 -0700 Subject: [PATCH] fix(tests): narrow retained callback xfails --- .../prometheus.py => _prometheus_helpers.py} | 5 ++ tests/test_litellm_rust/README.md | 2 +- tests/test_litellm_rust/conftest.py | 25 +++++-- .../integrations/test_exporters.py | 2 +- .../integrations/test_guardrails.py | 14 ++-- tests/test_litellm_rust/test_ocr.py | 72 ------------------- 6 files changed, 35 insertions(+), 85 deletions(-) rename tests/{test_litellm_rust/support/prometheus.py => _prometheus_helpers.py} (77%) delete mode 100644 tests/test_litellm_rust/test_ocr.py diff --git a/tests/test_litellm_rust/support/prometheus.py b/tests/_prometheus_helpers.py similarity index 77% rename from tests/test_litellm_rust/support/prometheus.py rename to tests/_prometheus_helpers.py index 000e6041a1e..f537253237f 100644 --- a/tests/test_litellm_rust/support/prometheus.py +++ b/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 diff --git a/tests/test_litellm_rust/README.md b/tests/test_litellm_rust/README.md index 772d6e82e00..9cb63366822 100644 --- a/tests/test_litellm_rust/README.md +++ b/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 diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index ccf321a1c8b..0a6ec0c82ba 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/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) diff --git a/tests/test_litellm_rust/integrations/test_exporters.py b/tests/test_litellm_rust/integrations/test_exporters.py index 08ce78e2a92..3df3aa7758a 100644 --- a/tests/test_litellm_rust/integrations/test_exporters.py +++ b/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, diff --git a/tests/test_litellm_rust/integrations/test_guardrails.py b/tests/test_litellm_rust/integrations/test_guardrails.py index 03d96efb768..bd3d83a14f9 100644 --- a/tests/test_litellm_rust/integrations/test_guardrails.py +++ b/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) diff --git a/tests/test_litellm_rust/test_ocr.py b/tests/test_litellm_rust/test_ocr.py deleted file mode 100644 index d5b1fce1139..00000000000 --- a/tests/test_litellm_rust/test_ocr.py +++ /dev/null @@ -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"}, - }