From 7e8be61fa04c253fb9500b2d146a6b169472d232 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:01:36 +0000 Subject: [PATCH] test: fix flaky team delete, refusal sink, loop-blocking, Moyai tamper and model-cost leak tests (#47135) * test(integration): stabilize auth and team delete checks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(unit): make loop-blocking and tampered-signature checks deterministic Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(unit): isolate model-cost state and synchronize loop progress Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(unit): fail loop-progress handshakes when the blocking call never starts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_team_delete_large_membership.py | 2 +- ...est_anthropic_refusal_stop_details_wire.py | 6 +- .../vertex_ai/count_tokens/test_context.py | 58 +++++++++++------- tests/unit/proxy/test_moyai_endpoints.py | 6 +- tests/unit/proxy/test_proxy_token_counter.py | 61 +++++++++++-------- tests/unit/test_utils.py | 3 +- 6 files changed, 82 insertions(+), 54 deletions(-) diff --git a/tests/integration/management/test_team_delete_large_membership.py b/tests/integration/management/test_team_delete_large_membership.py index 3e4bb568394..e752c835e4f 100644 --- a/tests/integration/management/test_team_delete_large_membership.py +++ b/tests/integration/management/test_team_delete_large_membership.py @@ -727,7 +727,7 @@ def _models_status(proxy: Gateway, key: str) -> int: def _statuses(proxy: Gateway, keys: Sequence[str]) -> dict[int, int]: - with ThreadPoolExecutor(max_workers=16) as pool: + with ThreadPoolExecutor(max_workers=2) as pool: statuses: Final = tuple(pool.map(lambda key: _models_status(proxy, key), keys)) return {status: statuses.count(status) for status in frozenset(statuses)} diff --git a/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py b/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py index b47dd242429..cd42e816ee3 100644 --- a/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py +++ b/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py @@ -896,7 +896,7 @@ def test_failing_sinks_get_one_attempt_each_and_the_proxy_stays_healthy(rig: Rig assert (readiness.status_code, readiness.json()["status"]) == (200, "healthy"), readiness.text -def test_unauthenticated_refusal_request_never_reaches_the_upstream_or_the_sink(rig: Rig) -> None: +def test_unauthenticated_refusal_request_never_reaches_upstream_and_logs_failure(rig: Rig) -> None: with rig.gateway.scenario() as scenario: model: Final = _deployment(scenario, "anthropic", rig) stranger: Final = "sk-not-a-virtual-key-" + _marker() @@ -908,7 +908,9 @@ def test_unauthenticated_refusal_request_never_reaches_the_upstream_or_the_sink( ) assert response.status_code == 401, response.text assert _targets(rig.wire) == () - assert rig.sink.for_key(stranger) == () + failures: Final = rig.sink.await_key(stranger) + assert len(failures) == 1, failures + assert failures[0]["status"] == "failure", failures[0] @pytest.mark.parametrize("fixture", ["null", "missing"]) diff --git a/tests/unit/llms/vertex_ai/count_tokens/test_context.py b/tests/unit/llms/vertex_ai/count_tokens/test_context.py index 948f66ab84d..a1643f0f699 100644 --- a/tests/unit/llms/vertex_ai/count_tokens/test_context.py +++ b/tests/unit/llms/vertex_ai/count_tokens/test_context.py @@ -1,12 +1,12 @@ import asyncio import base64 import logging -import time import uuid from collections.abc import Generator from contextlib import contextmanager +from dataclasses import dataclass, field from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from threading import Thread +from threading import Event, Thread from typing import Final, Literal from unittest.mock import patch @@ -297,25 +297,37 @@ async def test_openai_messages_keep_their_chat_conversion(custom_llm_provider: L _PNG_BYTES: Final = base64.b64decode( "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" ) -_DOWNLOAD_SECONDS: Final = 0.5 +_LOOP_PROOF_TIMEOUT_SECONDS: Final = 5 -class _SlowPngHandler(BaseHTTPRequestHandler): - def do_GET(self) -> None: - time.sleep(_DOWNLOAD_SECONDS) - self.send_response(200) - self.send_header("Content-Type", "image/png") - self.send_header("Content-Length", str(len(_PNG_BYTES))) - self.end_headers() - self.wfile.write(_PNG_BYTES) +@dataclass(frozen=True, slots=True) +class _LoopProbe: + request_arrived: Event = field(default_factory=Event) + loop_ran: Event = field(default_factory=Event) + loop_ran_while_download_pending: Event = field(default_factory=Event) - def log_message(self, format: str, *args: object) -> None: - return None + +def _png_handler_waiting_for_loop(probe: _LoopProbe) -> type[BaseHTTPRequestHandler]: + class _PngHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + probe.request_arrived.set() + if probe.loop_ran.wait(_LOOP_PROOF_TIMEOUT_SECONDS): + probe.loop_ran_while_download_pending.set() + self.send_response(200) + self.send_header("Content-Type", "image/png") + self.send_header("Content-Length", str(len(_PNG_BYTES))) + self.end_headers() + self.wfile.write(_PNG_BYTES) + + def log_message(self, format: str, *args: object) -> None: + return None + + return _PngHandler @contextmanager -def _slow_png_server() -> Generator[str, None, None]: - with ThreadingHTTPServer(("127.0.0.1", 0), _SlowPngHandler) as server: +def _png_server_waiting_for_loop(probe: _LoopProbe) -> Generator[str, None, None]: + with ThreadingHTTPServer(("127.0.0.1", 0), _png_handler_waiting_for_loop(probe)) as server: thread: Final = Thread(target=server.serve_forever, daemon=True) thread.start() try: @@ -324,10 +336,9 @@ def _slow_png_server() -> Generator[str, None, None]: server.shutdown() -async def _event_loop_wake_delay(nap_seconds: float) -> float: - started: Final = time.monotonic() - await asyncio.sleep(nap_seconds) - return time.monotonic() - started +async def _run_loop_once_download_starts(probe: _LoopProbe) -> None: + assert await asyncio.to_thread(probe.request_arrived.wait, _LOOP_PROOF_TIMEOUT_SECONDS) + probe.loop_ran.set() @pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) @@ -335,10 +346,11 @@ async def test_remote_image_downloads_without_blocking_the_event_loop( custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "user_url_validation", False) - with _slow_png_server() as base_url: + probe: Final = _LoopProbe() + with _png_server_waiting_for_loop(probe) as base_url: image_url: Final = f"{base_url}/{uuid.uuid4().hex}.png" - wake_delay, contents = await asyncio.gather( - _event_loop_wake_delay(0.05), + _, contents = await asyncio.gather( + _run_loop_once_download_starts(probe), gemini_count_contents( ( { @@ -365,7 +377,7 @@ async def test_remote_image_downloads_without_blocking_the_event_loop( ], } ] - assert wake_delay < _DOWNLOAD_SECONDS / 2 + assert probe.loop_ran_while_download_pending.is_set() class _MissingImageHandler(BaseHTTPRequestHandler): diff --git a/tests/unit/proxy/test_moyai_endpoints.py b/tests/unit/proxy/test_moyai_endpoints.py index 92a1261d952..aecda75c661 100644 --- a/tests/unit/proxy/test_moyai_endpoints.py +++ b/tests/unit/proxy/test_moyai_endpoints.py @@ -186,11 +186,13 @@ async def test_exchange_without_database_fails_before_nonce(monkeypatch: pytest. @pytest.mark.asyncio async def test_exchange_rejects_tampered_signature(monkeypatch: pytest.MonkeyPatch) -> None: from fastapi import HTTPException - from litellm.proxy.moyai_endpoints import _sign_connect_code + from litellm.proxy.moyai_endpoints import _b64url, _b64url_decode, _sign_connect_code await _exchange_env(monkeypatch) code = _sign_connect_code("sk-master", "https://moyai.example.com", "admin-user") - tampered = code[:-2] + ("aa" if not code.endswith("aa") else "bb") + payload_b64, signature_b64 = code.split(".", 1) + signature = _b64url_decode(signature_b64) + tampered = f"{payload_b64}.{_b64url(bytes((signature[0] ^ 1,)) + signature[1:])}" with pytest.raises(HTTPException) as exc: await _exchange_call(tampered, "https://moyai.example.com") diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 46e88614bb6..fa0c9eda02d 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -9,7 +9,6 @@ import itertools import json import logging import threading -import time import traceback from collections.abc import Callable, Coroutine, Mapping from functools import partial @@ -2730,7 +2729,7 @@ async def test_slow_cost_map_lookup_does_not_stall_other_requests_on_the_worker( request=TokenCountRequest(model="slow-lookup-model", messages=[{"role": "user", "content": "hi"}]) ) ) - await asyncio.to_thread(lookup_started.wait, 5) + assert await asyncio.to_thread(lookup_started.wait, 5) other_request_ran.set() await counting @@ -2854,17 +2853,24 @@ async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geog assert (counter.calls[0][0], counter.regions[0]) == ("bedrock_mantle/anthropic.claude-opus-5-5", counted_in) -class _SlowVault: +_LOOP_PROGRESS_TIMEOUT_SECONDS: Final = 5 + + +class _VaultWaitingForLoopProgress: + def __init__( + self, + lookup_started: threading.Event, + loop_progressed: threading.Event, + loop_progressed_during_lookup: threading.Event, + ) -> None: + self._lookup_started = lookup_started + self._loop_progressed = loop_progressed + self._loop_progressed_during_lookup = loop_progressed_during_lookup + def sync_read_secret(self, secret_name: str) -> None: - time.sleep(0.5) - - -async def _longest_stall(done: asyncio.Event) -> float: - if done.is_set(): - return 0.0 - started: Final = time.monotonic() - await asyncio.sleep(0.01) - return max(time.monotonic() - started, await _longest_stall(done)) + self._lookup_started.set() + if self._loop_progressed.wait(_LOOP_PROGRESS_TIMEOUT_SECONDS): + self._loop_progressed_during_lookup.set() @pytest.mark.asyncio @@ -2884,25 +2890,30 @@ async def test_region_lookup_for_a_shipped_count_model_leaves_the_event_loop_fre ), ) monkeypatch.setattr(proxy_server, "_get_provider_token_counter", bedrock_counter) - monkeypatch.setattr(litellm, "secret_manager_client", _SlowVault()) + lookup_started: Final = threading.Event() + loop_progressed: Final = threading.Event() + loop_progressed_during_lookup: Final = threading.Event() + monkeypatch.setattr( + litellm, + "secret_manager_client", + _VaultWaitingForLoopProgress(lookup_started, loop_progressed, loop_progressed_during_lookup), + ) monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.HASHICORP_VAULT) monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) - done: Final = asyncio.Event() async def count() -> TokenCountResponse: - try: - return await token_counter( - request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), - call_endpoint=True, - ) - finally: - done.set() + return await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), + call_endpoint=True, + ) - started: Final = time.monotonic() - stall, response = await asyncio.gather(_longest_stall(done), count()) + async def prove_loop_progress() -> None: + assert await asyncio.to_thread(lookup_started.wait, _LOOP_PROGRESS_TIMEOUT_SECONDS) + loop_progressed.set() - assert time.monotonic() - started >= 1.0 - assert stall < 0.4 + response, _ = await asyncio.gather(count(), prove_loop_progress()) + + assert loop_progressed_during_lookup.is_set() assert response.total_tokens == 447 diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index a1bdc4d5723..24d3ed678c9 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -4841,7 +4841,8 @@ def test_reapply_runtime_registrations_replays_register_model_overrides( # installed would make this depend on when that happens. monkeypatch.setattr(litellm_utils._LiveDeploymentReplay, "callback", None) - saved_model_cost = litellm.model_cost + saved_model_cost: Final = litellm.model_cost + monkeypatch.setattr(litellm, "model_cost", dict(saved_model_cost)) try: litellm.register_model( model_cost={