test(proxy): wait on the broadcast instead of a sleep and type the peer eviction tests

tests/unit/proxy/auth/test_auth_checks.py

@@ -2382,13 +2382,25 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_
     assert redis.entries == {}
 
 
-class _WedgedDeleteRedisCache(_AuthRedisCache):
+class _SignallingPublishRedisClient(_SlowPublishRedisClient):
     def __init__(self) -> None:
+        super().__init__()
+        self.sent = asyncio.Event()
+
+    async def _send(self, messages: Sequence[str]) -> None:
+        await super()._send(messages)
+        self.sent.set()
+
+
+class _DeleteAfterBroadcastRedisCache(_AuthRedisCache):
+    def __init__(self, broadcast_sent: asyncio.Event) -> None:
         super().__init__({})
-        self.release = asyncio.Event()
+        self._broadcast_sent = broadcast_sent
+        self.finished = False
 
     async def delete_cache_keys(self, keys: Sequence[str]) -> None:
-        await self.release.wait()
+        await self._broadcast_sent.wait()
+        self.finished = True
 
 
 @pytest.mark.asyncio
@@ -2397,9 +2409,8 @@ async def test_bulk_eviction_broadcasts_while_the_usage_cache_redis_is_wedged(
 ) -> None:
     monkeypatch.setattr(pubsub_module, "_in_flight_publishes", asyncio.Semaphore(16))
     monkeypatch.setattr(pubsub_module, "_pending_publishes", set())
-    monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5)
-    client: Final = _SlowPublishRedisClient()
-    usage_redis: Final = _WedgedDeleteRedisCache()
+    client: Final = _SignallingPublishRedisClient()
+    usage_redis: Final = _DeleteAfterBroadcastRedisCache(broadcast_sent=client.sent)
     proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=None)
     proxy_logging_obj.internal_usage_cache = InternalUsageCache(
         dual_cache=DualCache(redis_cache=usage_redis)  # pyright: ignore[reportArgumentType]  # a dict-backed fake of the Redis layer
@@ -2410,15 +2421,10 @@ async def test_bulk_eviction_broadcasts_while_the_usage_cache_redis_is_wedged(
         "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
         return_value=_PubSubRedisCache(client),
     ):
-        eviction: Final = asyncio.create_task(
-            delete_cache_key_objects(
-                hashed_tokens=hashed_tokens, user_api_key_cache=UserApiKeyCache(), proxy_logging_obj=proxy_logging_obj
-            )
+        await delete_cache_key_objects(
+            hashed_tokens=hashed_tokens, user_api_key_cache=UserApiKeyCache(), proxy_logging_obj=proxy_logging_obj
         )
-        await asyncio.sleep(0.1)
-        broadcast_while_wedged: Final = {json.loads(message)["cache_key"] for message in client.published}
-        usage_redis.release.set()
-        await eviction
         await _drain_publishes()
 
-    assert broadcast_while_wedged == set(hashed_tokens)
+    assert usage_redis.finished
+    assert {json.loads(message)["cache_key"] for message in client.published} == set(hashed_tokens)

tests/unit/proxy/management_endpoints/test_key_management_endpoints.py

@@ -5860,7 +5860,7 @@ async def _until_evicted(cache: UserApiKeyCache, hashed_token: str) -> None:
 
 
 @pytest.mark.asyncio
-async def test_delete_verification_tokens_stops_the_key_on_a_peer_worker(monkeypatch):
+async def test_delete_verification_tokens_stops_the_key_on_a_peer_worker(monkeypatch: pytest.MonkeyPatch) -> None:
     """Auth serves key objects cache-first, so a peer worker that is not told about the delete keeps
     authenticating the deleted key from its own copy until user_api_key_cache_ttl expires"""
     raw_key: Final = "sk-" + uuid4().hex
@@ -5899,7 +5899,7 @@ async def test_delete_verification_tokens_stops_the_key_on_a_peer_worker(monkeyp
 
 
 @pytest.mark.asyncio
-async def test_delete_verification_tokens_leaves_unowned_cache_entries_alone_on_every_worker(monkeypatch):
+async def test_delete_verification_tokens_leaves_unowned_cache_entries_alone_on_every_worker(monkeypatch: pytest.MonkeyPatch) -> None:
     """/key/delete accepts any string, and a name that is not a key row passes through unhashed. A key
     owner padding the request with another team's spend counter name must not wipe that counter from
     every worker's memory: only the rows the caller was allowed to delete are evicted and broadcast"""