scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-29 01:42:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

test(integration): cover tag object rpm and tpm limits

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-25 00:18:41 +00:00
parent c27170b384
commit cc312c30af
1 changed files with 106 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

106
tests/integration/spend/test_tag_budget_enforcement.py
View file

@ -1,9 +1,11 @@
import os
import uuid
from pathlib import Path
from typing import Final
from integration._support.client import Gateway, eventually
from integration._support.process import owned_proxy
from redis import Redis
def test_spend_over_a_tag_max_budget_rejects_the_next_request(gateway: Gateway) -> None:
@ -62,6 +64,110 @@ def test_spend_over_a_tag_max_budget_rejects_the_next_request(gateway: Gateway)
assert control.status_code == 200, control.text
def test_tag_object_rpm_limit_rejects_the_second_request_across_keys(gateway: Gateway) -> None:
tag: Final = f"tag-rpm-{uuid.uuid4().hex}"
def delete_tag() -> None:
gateway.post("/tag/delete", {"name": tag})
with gateway.scenario() as scenario:
model: Final = scenario.model(input_cost_per_token=0.01, output_cost_per_token=0.01)
key_a: Final = scenario.key()
key_b: Final = scenario.key()
gateway.post("/tag/new", {"name": tag, "rpm_limit": 1})
scenario.cleanups.callback(delete_tag)
first: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": f"tag rpm {tag}"}],
"metadata": {"tags": [tag]},
},
key=key_a,
)
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_b,
)
assert second.status_code == 429, second.text
assert "tag" 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_b,
)
assert control.status_code == 200, control.text
def test_tag_object_tpm_limit_rejects_the_next_request_once_tokens_are_charged(gateway: Gateway) -> None:
tag: Final = f"tag-tpm-{uuid.uuid4().hex}"
def delete_tag() -> None:
gateway.post("/tag/delete", {"name": tag})
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()
gateway.post("/tag/new", {"name": tag, "tpm_limit": 39})
scenario.cleanups.callback(delete_tag)
first: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": f"tag tpm {tag}"}],
"metadata": {"tags": [tag]},
},
key=key,
)
assert first.status_code == 200, first.text
with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache:
charged: Final = eventually(
lambda: int(cache.get(f"{{tag:{tag}}}:tokens") or 0),
lambda tokens: tokens >= 40,
seconds=70,
)
assert charged >= 40, charged
second: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": f"tag tpm {tag}"}],
"metadata": {"tags": [tag]},
},
key=key,
)
assert second.status_code == 429, second.text
assert "tag" in second.text.lower(), second.text
assert "token" in second.text.lower(), second.text
control: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": f"other tag tpm {tag}"}],
"metadata": {"tags": [f"other-{tag}"]},
},
key=key,
)
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: