diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index 29a036a4142..361c031a6be 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/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) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 8032ef4659d..dd6e40c0d62 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/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"""