diff --git a/tests/unit/proxy/test_proxy_utils.py b/tests/unit/proxy/test_proxy_utils.py index 92bb163d6ec..b5b8930671f 100644 --- a/tests/unit/proxy/test_proxy_utils.py +++ b/tests/unit/proxy/test_proxy_utils.py @@ -2589,7 +2589,7 @@ async def test_handle_logging_proxy_only_error_syncs_normalized_call_type( async def test_during_call_hook_parallel_execution(): """ Test that multiple guardrails in during_call_hook are executed in parallel. - Verifies parallel execution by checking timing and execution order. + Each guardrail blocks until all of them have started, so sequential execution times out. """ from litellm.caching.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail @@ -2599,6 +2599,8 @@ async def test_during_call_hook_parallel_execution(): cache = DualCache() proxy_logging = ProxyLogging(user_api_key_cache=cache) execution_order = [] + guardrail_count: Final = 3 + all_started: Final = asyncio.Event() class TestGuardrail(CustomGuardrail): def __init__(self, name): @@ -2611,24 +2613,23 @@ async def test_during_call_hook_parallel_execution(): async def async_moderation_hook(self, data, user_api_key_dict, call_type): execution_order.append(f"{self.name}_start") - await asyncio.sleep(0.1) + if sum(1 for item in execution_order if item.endswith("_start")) == guardrail_count: + all_started.set() + await asyncio.wait_for(all_started.wait(), timeout=5) execution_order.append(f"{self.name}_end") return data original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] try: - litellm.callbacks = [TestGuardrail(f"g{i}") for i in range(3)] + litellm.callbacks = [TestGuardrail(f"g{i}") for i in range(guardrail_count)] - start_time = asyncio.get_event_loop().time() result = await proxy_logging.during_call_hook( data={"model": "gpt-4", "messages": [{"role": "user", "content": "test"}]}, user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), call_type="completion", ) - execution_time = asyncio.get_event_loop().time() - start_time - # Verify parallel execution: all start before any end first_end_idx = next( i for i, item in enumerate(execution_order) if "end" in item ) @@ -2636,13 +2637,8 @@ async def test_during_call_hook_parallel_execution(): 1 for item in execution_order[:first_end_idx] if "start" in item ) assert ( - starts_before_end == 3 - ), f"Expected 3 starts before first end, got {starts_before_end}" - - # Verify timing: parallel ~0.1s vs sequential ~0.3s - assert ( - execution_time < 0.2 - ), f"Parallel execution took {execution_time}s, expected < 0.2s" + starts_before_end == guardrail_count + ), f"Expected {guardrail_count} starts before first end, got {starts_before_end}" assert result["model"] == "gpt-4" finally: litellm.callbacks = original_callbacks diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.characterization.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.characterization.test.tsx index 9628c284bd2..076b2ada77f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.characterization.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.characterization.test.tsx @@ -163,9 +163,9 @@ describe("AddGuardrailForm create payload characterization", () => { await user.click((await screen.findAllByText("logging_only")).at(-1) as HTMLElement); await user.click(await screen.findByLabelText("Logging only scope")); + expect(await screen.findByRole("option", { name: "Both (request and response)" })).toBeInTheDocument(); expect(screen.queryByRole("option", { name: "Input only (request)" })).not.toBeInTheDocument(); expect(screen.queryByRole("option", { name: "Output only (response)" })).not.toBeInTheDocument(); - expect(screen.getByRole("option", { name: "Both (request and response)" })).toBeInTheDocument(); }); it("hides logging-only scope and omits it from a pre-call payload", async () => {