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>

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)}
 

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"])

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):

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")

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
 
 

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={