mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
test: add tag rpm limit and tag budget reset integration coverage (#42859)
Adds two integration tests to tests/integration/spend/test_tag_budget_enforcement.py test_key_tag_rpm_limit_rejects_the_second_request_carrying_that_tag proves a key metadata tag_rpm_limit of 1 rejects the second request carrying that tag with 429 while a request carrying a different tag still passes test_tag_budget_duration_resets_spend_and_unblocks_the_tag boots an owned proxy with a 2 to 3 second budget rescheduler, creates a tag with max_budget 0.0001 and budget_duration 5s, observes the spend block, then observes the tag serving again once ResetBudgetJob zeroes the tag spend Mutation evidence get_key_tag_rpm_limit forced to return None: the second tagged request returned 200 instead of 429, test red _queue_budget_linked_resets for uow.tags disabled in _commit_budget_cascade_once: the tag stayed blocked at 422 for the full 70 second recovery window, test red Co-authored-by: kerry <kerry@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
a4f69e058c
commit
a3d791f348
1 changed files with 98 additions and 0 deletions
|
|
@ -1,7 +1,9 @@
|
|||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.process import owned_proxy
|
||||
|
||||
|
||||
def test_spend_over_a_tag_max_budget_rejects_the_next_request(gateway: Gateway) -> None:
|
||||
|
|
@ -58,3 +60,99 @@ def test_spend_over_a_tag_max_budget_rejects_the_next_request(gateway: Gateway)
|
|||
},
|
||||
)
|
||||
assert control.status_code == 200, control.text
|
||||
|
||||
|
||||
def test_key_tag_rpm_limit_rejects_the_second_request_carrying_that_tag(gateway: Gateway) -> None:
|
||||
tag: Final = f"tag-rpm-{uuid.uuid4().hex}"
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(input_cost_per_token=0.01, output_cost_per_token=0.01)
|
||||
key: Final = scenario.key(metadata={"tag_rpm_limit": {tag: 1}})
|
||||
first: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"tag rpm {tag}"}],
|
||||
"metadata": {"tags": [tag]},
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert first.status_code == 200, first.text
|
||||
second: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"tag rpm {tag}"}],
|
||||
"metadata": {"tags": [tag]},
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert second.status_code == 429, second.text
|
||||
assert "rpm" in second.text.lower() or "rate" in second.text.lower(), second.text
|
||||
control: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"other tag rpm {tag}"}],
|
||||
"metadata": {"tags": [f"other-{tag}"]},
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert control.status_code == 200, control.text
|
||||
|
||||
|
||||
def test_tag_budget_duration_resets_spend_and_unblocks_the_tag(gateway: Gateway, tmp_path: Path) -> None:
|
||||
tag: Final = f"tag-reset-{uuid.uuid4().hex}"
|
||||
with (
|
||||
owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"PROXY_BUDGET_RESCHEDULER_MIN_TIME": "2", "PROXY_BUDGET_RESCHEDULER_MAX_TIME": "3"},
|
||||
) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
|
||||
def delete_tag() -> None:
|
||||
candidate.post("/tag/delete", {"name": tag})
|
||||
|
||||
model: Final = scenario.model(input_cost_per_token=0.01, output_cost_per_token=0.01)
|
||||
candidate.post("/tag/new", {"name": tag, "max_budget": 0.0001, "budget_duration": "5s"})
|
||||
scenario.cleanups.callback(delete_tag)
|
||||
first: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"tag spend {tag}"}],
|
||||
"metadata": {"tags": [tag]},
|
||||
},
|
||||
)
|
||||
assert first.status_code == 200, first.text
|
||||
|
||||
def rejection() -> int:
|
||||
return candidate.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"tag reset probe {tag}"}],
|
||||
"metadata": {"tags": [tag]},
|
||||
},
|
||||
).status_code
|
||||
|
||||
status: Final = eventually(rejection, lambda code: code != 200, seconds=70)
|
||||
assert status in (400, 422, 429), status
|
||||
blocked: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"tag reset probe {tag}"}],
|
||||
"metadata": {"tags": [tag]},
|
||||
},
|
||||
)
|
||||
assert "budget" in blocked.text.lower(), blocked.text
|
||||
recovered: Final = eventually(rejection, lambda code: code == 200, seconds=70)
|
||||
assert recovered == 200, recovered
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue