test(router): observe scheduling through wait_for_turn, not add_request

The priority tests spied on Scheduler.add_request to check which priority a
request was scheduled at and that it was scheduled exactly once. A request a
healthy deployment can take no longer joins the queue, so add_request is never
called on that path and the spies saw nothing, while the zero-count checks
would have kept passing even if a request without a priority was scheduled.

wait_for_turn is the router's single entry into the scheduler and receives the
same FlowItem, so the assertions stay as they were.

tests/unit/proxy/proxy_server/test_routes_queue.py

@@ -154,11 +154,11 @@ def test_queue_chat_completions_reports_the_priority_the_router_queued_at(
 ) -> None:
     router: Final = _queue_router(monkeypatch, default_priority)
     payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], **priority_fields}
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         with auth_as():
             response: Final = client.post("/queue/chat/completions", json=payload)
     assert response.status_code == 200, response.text
-    assert [call.kwargs["request"].priority for call in add_request.await_args_list] == [expected_priority]
+    assert [call.kwargs["request"].priority for call in wait_for_turn.await_args_list] == [expected_priority]
     assert response.headers.get("x-litellm-priority") == str(expected_priority)
 
 
@@ -181,14 +181,14 @@ def test_queue_chat_completions_rejects_a_request_it_cannot_queue_as_an_invalid_
 ) -> None:
     router: Final = _queue_router(monkeypatch, default_priority)
     payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], **priority_fields}
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         with auth_as():
             response: Final = client.post("/queue/chat/completions", json=payload)
     assert response.status_code == 400, response.text
     error: Final = response.json()["error"]
     assert (error["type"], error["param"]) == ("invalid_request_error", "priority")
     assert message in error["message"]
-    assert add_request.await_count == 0
+    assert wait_for_turn.await_count == 0
 
 
 def test_streamed_queue_responses_report_the_priority_the_router_queued_at(

tests/unit/test_router/test_router.py

@@ -26325,20 +26325,20 @@ async def test_acompletion_queues_the_resolved_priority_exactly_once(
     default_priority: int | None, request_kwargs: dict[str, object], queued_priority: int
 ) -> None:
     router: Final = _priority_router(default_priority)
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.acompletion(
             model="qa-chat", messages=_PRIORITY_MESSAGES, mock_response="pong", **request_kwargs
         )
-    assert add_request.await_count == 1
-    assert add_request.await_args.kwargs["request"].priority == queued_priority
+    assert wait_for_turn.await_count == 1
+    assert wait_for_turn.await_args.kwargs["request"].priority == queued_priority
     assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True
 
 
 async def test_acompletion_without_any_priority_never_enters_the_scheduler() -> None:
     router: Final = _priority_router(None)
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.acompletion(model="qa-chat", messages=_PRIORITY_MESSAGES, mock_response="pong")
-    assert add_request.await_count == 0
+    assert wait_for_turn.await_count == 0
     assert "x-litellm-request-prioritization-used" not in response._hidden_params.get("additional_headers", {})
 
 
@@ -26402,12 +26402,12 @@ async def test_acompletion_with_an_integer_priority_sends_no_priority_to_the_pro
 
 async def test_schedule_acompletion_under_a_default_priority_queues_the_explicit_priority_once() -> None:
     router: Final = _priority_router(3)
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.schedule_acompletion(
             model="qa-chat", messages=_PRIORITY_MESSAGES, priority=1, mock_response="pong"
         )
-    assert add_request.await_count == 1
-    assert add_request.await_args.kwargs["request"].priority == 1
+    assert wait_for_turn.await_count == 1
+    assert wait_for_turn.await_args.kwargs["request"].priority == 1
     assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True
 
 
@@ -26434,11 +26434,11 @@ async def test_acompletion_drops_a_non_integer_priority_under_the_router_or_depl
         default_priority=3,
         default_litellm_params=router_defaults,
     )
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.acompletion(
             model="qa-chat", messages=_PRIORITY_MESSAGES, mock_response="pong", priority="high"
         )
-    assert add_request.await_args.kwargs["request"].priority == 3
+    assert wait_for_turn.await_args.kwargs["request"].priority == 3
     assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True
 
 
@@ -26482,7 +26482,7 @@ async def test_deployment_drop_params_drop_a_non_integer_priority_on_every_deplo
 ) -> None:
     monkeypatch.setattr(litellm, "drop_params", False)
     router: Final = Router(model_list=[deployment], default_priority=3)
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.acompletion(
             model=requested_model,
             messages=_PRIORITY_MESSAGES,
@@ -26490,7 +26490,7 @@ async def test_deployment_drop_params_drop_a_non_integer_priority_on_every_deplo
             priority="high",
             **routing_fields,
         )
-    assert add_request.await_args.kwargs["request"].priority == 3
+    assert wait_for_turn.await_args.kwargs["request"].priority == 3
     assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True
 
 
@@ -26518,12 +26518,12 @@ async def test_schedule_acompletion_queues_a_prompt_management_model(monkeypatch
         litellm_call_id="prompt-management-priority",
         function_id="f",
     )
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.schedule_acompletion(
             model="cached-chat", messages=_PRIORITY_MESSAGES, priority=1, litellm_logging_obj=logging_obj
         )
-    assert add_request.await_count == 1
-    assert add_request.await_args.kwargs["request"].priority == 1
+    assert wait_for_turn.await_count == 1
+    assert wait_for_turn.await_args.kwargs["request"].priority == 1
     assert provider.await_count == 1
     assert "priority" not in provider.await_args.kwargs
     assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True
@@ -26550,9 +26550,9 @@ def test_an_invalid_default_priority_is_ignored_with_a_warning(
 
 async def test_an_invalid_default_priority_leaves_every_request_unscheduled() -> None:
     router: Final = _priority_router("high")
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.acompletion(model="qa-chat", messages=_PRIORITY_MESSAGES, mock_response="pong")
-    assert add_request.await_count == 0
+    assert wait_for_turn.await_count == 0
     assert "x-litellm-request-prioritization-used" not in response._hidden_params.get("additional_headers", {})
 
 
@@ -26569,12 +26569,12 @@ async def test_schedule_acompletion_without_a_priority_or_a_default_is_rejected_
 
 async def test_schedule_acompletion_without_a_priority_queues_the_default_once() -> None:
     router: Final = _priority_router(3)
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.schedule_acompletion(
             model="qa-chat", messages=_PRIORITY_MESSAGES, mock_response="pong"
         )
-    assert add_request.await_count == 1
-    assert add_request.await_args.kwargs["request"].priority == 3
+    assert wait_for_turn.await_count == 1
+    assert wait_for_turn.await_args.kwargs["request"].priority == 3
     assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True
 
 
@@ -26598,10 +26598,10 @@ async def test_atext_completion_queues_the_resolved_priority_exactly_once(
 ) -> None:
     provider: Final = _priority_provider(respx_mock, monkeypatch)
     router: Final = _priority_router(default_priority)
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.atext_completion(model="qa-chat", prompt="pong?", **request_kwargs)
-    assert add_request.await_count == 1
-    assert add_request.await_args.kwargs["request"].priority == queued_priority
+    assert wait_for_turn.await_count == 1
+    assert wait_for_turn.await_args.kwargs["request"].priority == queued_priority
     assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True
     assert "priority" not in json.loads(provider.calls.last.request.content)
 
@@ -26611,9 +26611,9 @@ async def test_atext_completion_without_any_priority_never_enters_the_scheduler(
 ) -> None:
     provider: Final = _priority_provider(respx_mock, monkeypatch)
     router: Final = _priority_router(None)
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.atext_completion(model="qa-chat", prompt="pong?")
-    assert add_request.await_count == 0
+    assert wait_for_turn.await_count == 0
     assert "x-litellm-request-prioritization-used" not in response._hidden_params.get("additional_headers", {})
     assert provider.called