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:
devin-ai-integration[bot] 2026-09-24 02:49:01 +00:00 • committed by GitHub
parent a4f69e058c
commit a3d791f348
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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