test(integration): measure dynamic limiter shares after saturation (#51714)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/integration/routing/test_dynamic_limiter_prometheus_metrics.py

@@ -18,9 +18,12 @@ from integration._support.prometheus_series import Sample, scrape
 from integration._support.wire import Reply, Request, wire_server
 
 OPENAI_MODEL: Final = "gpt-4.1-mini"
-MODEL_TPM: Final = 600
+MODEL_TPM: Final = 1000
 TEAM_TPM: Final = 100
-REPLY_TOKENS: Final = 30
+REPLY_TOKENS: Final = 60
+SATURATION_THRESHOLD: Final = 0.4
+IDLE_CALLS: Final = TEAM_TPM // REPLY_TOKENS + 1
+WARMUP_CALLS: Final = int(MODEL_TPM * SATURATION_THRESHOLD) // REPLY_TOKENS + 1
 UPSTREAM_REPLY: Final = json.dumps(
     {
         "id": "chatcmpl_dynamic_limiter_metrics",
@@ -28,7 +31,7 @@ UPSTREAM_REPLY: Final = json.dumps(
         "created": 1700000000,
         "model": OPENAI_MODEL,
         "choices": [{"index": 0, "message": {"role": "assistant", "content": "metered"}, "finish_reason": "stop"}],
-        "usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": REPLY_TOKENS},
+        "usage": {"prompt_tokens": 20, "completion_tokens": REPLY_TOKENS - 20, "total_tokens": REPLY_TOKENS},
     }
 ).encode()
 
@@ -48,7 +51,7 @@ def saturated_model(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[
         "priority_reservation": {"interactive": 0.6, "batch": 0.4},
         "priority_reservation_settings": {
             "default_priority": 1.0,
-            "saturation_threshold": 0.6,
+            "saturation_threshold": SATURATION_THRESHOLD,
             "saturation_check_cache_ttl": 0,
             "team_model_limit_borrows_idle_capacity_from_other_teams": True,
         },
@@ -72,6 +75,26 @@ def _team_series(samples: Sequence[Sample], name: str, team: str) -> Mapping[str
     }
 
 
+def _call(owned: Gateway, model: str, key: str) -> httpx.Response:
+    return owned.request(
+        "POST",
+        "/v1/chat/completions",
+        {"model": model, "messages": [{"role": "user", "content": f"hi {uuid.uuid4().hex[:12]}"}]},
+        key=key,
+    )
+
+
+def _settled_calls(owned: Gateway, model: str, team: str, key: str, count: int) -> None:
+    for index in range(count):
+        response: Final = _call(owned, model, key)
+        assert response.status_code == 200, response.text
+        eventually(
+            lambda: sum(_team_series(scrape(owned), "litellm_total_tokens_metric_total", team).values()),
+            lambda tokens, expected=(index + 1) * REPLY_TOKENS: tokens == expected,
+            seconds=30,
+        )
+
+
 def test_a_team_over_its_share_shows_up_as_borrowed_tokens_and_a_matching_throttle_reason(
     saturated_model: tuple[Gateway, str],
 ) -> None:
@@ -82,17 +105,13 @@ def test_a_team_over_its_share_shows_up_as_borrowed_tokens_and_a_matching_thrott
         )
         team: Final = scenario.team(metadata={"model_tpm_limit": {model: TEAM_TPM}})
         key: Final = scenario.key(team_id=team)
-        responses: Final = tuple(
-            owned.request(
-                "POST",
-                "/v1/chat/completions",
-                {"model": model, "messages": [{"role": "user", "content": f"metered {uuid.uuid4().hex}"}]},
-                key=key,
-            )
-            for _ in range(16)
-        )
+        _settled_calls(owned, model, team, key, IDLE_CALLS)
+        assert IDLE_CALLS * REPLY_TOKENS > TEAM_TPM
+        warmup: Final = scenario.team(metadata={"priority": "batch"})
+        _settled_calls(owned, model, warmup, scenario.key(team_id=warmup), WARMUP_CALLS)
+        responses: Final = tuple(_call(owned, model, key) for _ in range(IDLE_CALLS + 1))
         blocked: Final = tuple(response for response in responses if response.status_code == 429)
-        expected_borrowed: Final = max(0, (len(responses) - len(blocked)) * REPLY_TOKENS - TEAM_TPM)
+        expected_borrowed: Final = IDLE_CALLS * REPLY_TOKENS - TEAM_TPM
         samples: Final = eventually(
             lambda: scrape(owned),
             lambda scraped: (
@@ -106,7 +125,9 @@ def test_a_team_over_its_share_shows_up_as_borrowed_tokens_and_a_matching_thrott
         )
 
     assert responses[0].status_code == 200, responses[0].text
-    assert blocked, [r.text for r in responses]
+    assert tuple(response.status_code for response in responses) == (200,) * IDLE_CALLS + (429,), [
+        r.text for r in responses
+    ]
     reason: Final = blocked[0].headers.get("rate_limited_by")
     assert reason == "model_per_team_saturated", dict(blocked[0].headers)
 
@@ -133,6 +154,8 @@ def test_a_team_under_its_share_borrows_nothing(saturated_model: tuple[Gateway,
             model=f"openai/{OPENAI_MODEL}", api_base=f"{upstream}/v1", api_key="synthetic-openai-key", tpm=MODEL_TPM
         )
         team: Final = scenario.team(metadata={"model_tpm_limit": {model: TEAM_TPM}})
+        warmup: Final = scenario.team(metadata={"priority": "batch"})
+        _settled_calls(owned, model, warmup, scenario.key(team_id=warmup), WARMUP_CALLS)
         response: Final[httpx.Response] = owned.request(
             "POST",
             "/v1/chat/completions",