From e9825f1d269d77185f5f238d6704d227a18e4edc Mon Sep 17 00:00:00 2001
From: yucheng <yucheng@berri.ai>
Date: Thu, 17 Sep 2026 22:39:15 +0000
Subject: [PATCH] test(proxy): drive the heuristics responsiveness check
 without mutable state

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../hooks/test_prompt_injection_detection.py  | 31 +++++++++----------
 1 file changed, 15 insertions(+), 16 deletions(-)

diff --git a/tests/test_litellm/proxy/hooks/test_prompt_injection_detection.py b/tests/test_litellm/proxy/hooks/test_prompt_injection_detection.py
index a04ed9345ee..bbd35404136 100644
--- a/tests/test_litellm/proxy/hooks/test_prompt_injection_detection.py
+++ b/tests/test_litellm/proxy/hooks/test_prompt_injection_detection.py
@@ -1,6 +1,7 @@
 import asyncio
 import importlib
 import time
+from collections.abc import AsyncIterator
 from concurrent.futures import ThreadPoolExecutor
 
 import pytest
@@ -73,29 +74,27 @@ async def test_heuristics_check_keeps_event_loop_responsive():
         prompt_injection_params=LiteLLMPromptInjectionParams(heuristics_check=True)
     )
     data = {"model": "test-model", "messages": [{"role": "user", "content": LONG_SAFE_PROMPT}]}
-    ticks_during_scan: list[float] = []
-    scan_done = asyncio.Event()
 
-    async def ticker() -> None:
-        while not scan_done.is_set():
+    async def ticks_until_done(task: asyncio.Task[dict]) -> AsyncIterator[float]:
+        while not task.done():
             await asyncio.sleep(0.01)
-            ticks_during_scan.append(time.perf_counter())
+            yield time.perf_counter()
 
-    ticker_task = asyncio.create_task(ticker())
-    started = time.perf_counter()
-    result = await detector.async_pre_call_hook(
-        user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
-        cache=DualCache(),
-        data=data,
-        call_type="acompletion",
+    scan = asyncio.create_task(
+        detector.async_pre_call_hook(
+            user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
+            cache=DualCache(),
+            data=data,
+            call_type="acompletion",
+        )
     )
+    started = time.perf_counter()
+    ticks_during_scan = tuple([tick async for tick in ticks_until_done(scan)])
     finished = time.perf_counter()
-    scan_done.set()
-    await ticker_task
+    result = await scan
 
     assert result == data
-    ticks_before_finish = [tick for tick in ticks_during_scan if tick < finished]
-    assert len(ticks_before_finish) >= int((finished - started) / 0.05)
+    assert len(ticks_during_scan) >= int((finished - started) / 0.05)
 
 
 @pytest.mark.asyncio