test(integration): cover budgets_spend endpoint contracts (#44736)

* test(integration): cover budgets_spend endpoint contracts

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

* test(integration): cover budgets_spend endpoint contracts

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

* test(integration): cover budgets_spend endpoint contracts

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

* test(integration): cover budgets_spend endpoint contracts

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

* test(integration): harden budgets_spend contracts per review

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

* test(integration): make spend report window midnight safe and drop redundant comments

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

* test(integration): pin organization spend report scope

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

* test(integration): close budgets and spend review findings

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

* test(integration): drop streamed cost header absence assertion

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

* test(spend): cover streamed spend on linked budget tiers

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

* test(spend): wait for async spend logs

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

---------

Co-authored-by: mateo <mateo@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-10-14 06:58:22 +00:00 • committed by GitHub
parent 60ee966885
commit 8822e086ec
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 1703 additions and 3 deletions

View file

@ -1,10 +1,13 @@
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Final
import pytest
from pydantic import JsonValue, TypeAdapter
from tests.integration._support.client import Gateway, string_value
from tests.integration._support.database import read_rows
from tests.integration._support.process import owned_proxy
def _persisted_reset_at(budget_id: str) -> datetime:
@ -27,3 +30,184 @@ def test_shortening_budget_duration_moves_reset_at_onto_the_new_schedule(gateway
updated: Final = _persisted_reset_at(budget_id)
assert updated < ten_day_reset_at, f"{updated} not before {ten_day_reset_at}"
assert before < updated <= before + timedelta(days=1, minutes=5), f"{updated} not within 1d of {before}"
BUDGET_COLUMNS: Final = (
"budget_id, max_budget, soft_budget, max_parallel_requests, tpm_limit, rpm_limit, tpd_limit, model_max_budget, "
"budget_duration, allowed_models, temp_budget_increase, temp_budget_expiry, created_by, updated_by"
)
VOLATILE_INFO_FIELDS: Final = frozenset({"created_at", "updated_at", "budget_reset_at"})
UNLOADED_RELATIONS: Final[dict[str, JsonValue]] = dict.fromkeys(
(
"organization",
"projects",
"keys",
"end_users",
"tags",
"model_access_groups",
"team_membership",
"organization_membership",
)
)
def _budget_row(budget_id: str) -> dict[str, JsonValue]:
rows: Final = read_rows(
f'SELECT {BUDGET_COLUMNS}, budget_reset_at::text AS reset_at FROM "LiteLLM_BudgetTable" WHERE budget_id = %s',
(budget_id,),
)
assert len(rows) == 1, rows
return rows[0]
def _budget_info(gateway: Gateway, budget_id: str) -> dict[str, JsonValue]:
response: Final = gateway.request("POST", "/budget/info", {"budgets": [budget_id]})
assert response.status_code == 200, response.text
entries: Final = TypeAdapter(list[dict[str, JsonValue]]).validate_json(response.content)
assert len(entries) == 1, entries
return entries[0]
def _reset_at(value: JsonValue) -> datetime | None:
if value is None:
return None
parsed: Final = datetime.fromisoformat(string_value(value).replace("Z", "+00:00"))
return parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=timezone.utc)
def _assert_budget(
gateway: Gateway,
budget_id: str,
expected: dict[str, JsonValue],
reset_after: datetime | None,
reset_within: timedelta,
) -> None:
row: Final = _budget_row(budget_id)
reset_at: Final = _reset_at(row.pop("reset_at"))
assert row == expected, row
info: Final = _budget_info(gateway, budget_id)
assert {name: value for name, value in info.items() if name not in VOLATILE_INFO_FIELDS} == {
**expected,
**UNLOADED_RELATIONS,
}, info
assert VOLATILE_INFO_FIELDS <= info.keys(), info
assert _reset_at(info["budget_reset_at"]) == reset_at, info
if reset_after is None:
assert reset_at is None, row
else:
assert reset_at is not None and reset_after < reset_at <= reset_after + reset_within, (reset_at, reset_after)
def _update(gateway: Gateway, body: dict[str, JsonValue]) -> None:
response: Final = gateway.request("POST", "/budget/update", body)
assert response.status_code == 200, response.text
def test_partial_update_keeps_sibling_budget_fields(gateway: Gateway) -> None:
model_budget: Final[dict[str, JsonValue]] = {"openai/gpt-4o-mini": {"max_budget": 1, "budget_duration": "1d"}}
with gateway.scenario() as scenario:
created_at: Final = datetime.now(timezone.utc)
budget_id: Final = scenario.budget(
max_budget=10.0,
soft_budget=5.0,
tpm_limit=1000,
rpm_limit=60,
tpd_limit=100000,
model_max_budget=model_budget,
budget_duration="30d",
)
creator: Final = string_value(_budget_row(budget_id)["created_by"])
expected: Final[dict[str, JsonValue]] = {
"budget_id": budget_id,
"max_budget": 10.0,
"soft_budget": 5.0,
"max_parallel_requests": None,
"tpm_limit": 1000,
"rpm_limit": 60,
"tpd_limit": 100000,
"model_max_budget": {"openai/gpt-4o-mini": {"max_budget": 1.0, "budget_duration": "1d"}},
"budget_duration": "30d",
"allowed_models": [],
"temp_budget_increase": None,
"temp_budget_expiry": None,
"created_by": creator,
"updated_by": creator,
}
_assert_budget(gateway, budget_id, expected, created_at, timedelta(days=30, minutes=5))
before_partial: Final = datetime.now(timezone.utc)
_update(gateway, {"budget_id": budget_id, "max_budget": 20.0})
expected["max_budget"] = 20.0
_assert_budget(gateway, budget_id, expected, created_at, timedelta(days=30, minutes=5))
_update(
gateway,
{
"budget_id": budget_id,
"model_max_budget": {"openai/gpt-4o-mini": {"max_budget": 2, "budget_duration": "1d"}},
},
)
expected["model_max_budget"] = {"openai/gpt-4o-mini": {"max_budget": 2.0, "budget_duration": "1d"}}
_assert_budget(gateway, budget_id, expected, created_at, timedelta(days=30, minutes=5))
_update(gateway, {"budget_id": budget_id, "soft_budget": None})
expected["soft_budget"] = None
_assert_budget(gateway, budget_id, expected, created_at, timedelta(days=30, minutes=5))
before_hourly: Final = datetime.now(timezone.utc)
_update(gateway, {"budget_id": budget_id, "budget_duration": "1h"})
expected["budget_duration"] = "1h"
_assert_budget(gateway, budget_id, expected, before_hourly, timedelta(hours=1, minutes=5))
assert before_partial <= before_hourly
_update(gateway, {"budget_id": budget_id, "budget_duration": None})
expected["budget_duration"] = None
_assert_budget(gateway, budget_id, expected, None, timedelta())
pinned: Final = datetime(2030, 1, 1, tzinfo=timezone.utc)
_update(gateway, {"budget_id": budget_id, "budget_duration": None, "budget_reset_at": pinned.isoformat()})
_assert_budget(gateway, budget_id, expected, pinned - timedelta(microseconds=1), timedelta(microseconds=1))
rejected: Final = gateway.request(
"POST",
"/budget/update",
{"budget_id": budget_id, "model_max_budget": {"openai/gpt-4o-mini": {"max_budget": "lots"}}},
)
assert rejected.status_code == 422, rejected.text
assert rejected.json() == {
"detail": [
{
"type": "float_parsing",
"loc": ["body", "model_max_budget", "openai/gpt-4o-mini", "max_budget"],
"msg": "Input should be a valid number, unable to parse string as a number",
}
]
}, rejected.text
_assert_budget(gateway, budget_id, expected, pinned - timedelta(microseconds=1), timedelta(microseconds=1))
def test_model_max_budget_update_is_refused_by_the_handler_without_a_license(gateway: Gateway, tmp_path: Path) -> None:
model_budget: Final[dict[str, JsonValue]] = {"openai/gpt-4o-mini": {"max_budget": 1, "budget_duration": "1d"}}
with gateway.scenario() as scenario:
budget_id: Final = scenario.budget(max_budget=10.0, model_max_budget=model_budget)
before: Final = _budget_row(budget_id)
with owned_proxy(gateway, tmp_path, {}, remove_environment=("LITELLM_LICENSE",)) as unlicensed:
rejected: Final = unlicensed.request(
"POST",
"/budget/update",
{"budget_id": budget_id, "model_max_budget": {"openai/gpt-4o-mini": {"max_budget": 2}}},
)
assert rejected.status_code == 400, rejected.text
assert rejected.json() == {
"detail": {
"error": "Invalid model_max_budget: You must have an enterprise license to set model_max_budget. "
"You must be a LiteLLM Enterprise user to use this feature. If you have a license please set "
"`LITELLM_LICENSE` in your env. Get a 7 day trial key here: https://www.litellm.ai/enterprise#trial. "
"\nPricing: https://www.litellm.ai/#pricing. Example of valid model_max_budget: "
"https://docs.litellm.ai/docs/proxy/users"
}
}, rejected.text
assert _budget_row(budget_id) == before
served: Final = unlicensed.request("POST", "/budget/update", {"budget_id": budget_id, "max_budget": 20.0})
assert served.status_code == 200, served.text
assert _budget_row(budget_id) == {**before, "max_budget": 20.0}

View file

@ -7,6 +7,7 @@ from hashlib import sha256
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from pydantic import JsonValue
@ -358,3 +359,101 @@ def test_loaded_router_preserves_cached_defaults_during_real_requests(gateway: G
assert info["output_cost_per_token"] == (0.0 if alias == aliases[2] else 0.0000006)
finally:
router.reset()
def _bounded_chat(gateway: Gateway, model: str, key: str, content: str) -> httpx.Response:
return gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "max_tokens": 20, "messages": [{"role": "user", "content": content}]},
key=key,
)
def _observed(upstream: httpx.Client) -> list[JsonValue]:
response: Final = upstream.get("/__observations")
response.raise_for_status()
requests: Final = object_value(response.json())["requests"]
assert isinstance(requests, list), response.text
return requests
def _upstream_chat(upstream_model: str, content: str) -> dict[str, JsonValue]:
return {
"path": "/v1/chat/completions",
"authorization": "Bearer integration-provider-key",
"body": {"model": upstream_model, "max_tokens": 20, "messages": [{"role": "user", "content": content}]},
"method": "POST",
"api_key": "",
}
def _block_unpriced(gateway: Gateway, enabled: bool) -> httpx.Response:
return gateway.request("PATCH", "/config/block_requests_for_models_without_pricing", {"enabled": enabled})
def test_block_requests_for_models_without_pricing_refuses_unpriced_deployments(gateway: Gateway) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
original: Final = gateway.get("/config/block_requests_for_models_without_pricing")["enabled"]
scenario.cleanups.callback(_block_unpriced, gateway, bool(original))
enabled: Final = _block_unpriced(gateway, True)
assert enabled.status_code == 200, enabled.text
assert enabled.json() == {"enabled": True}, enabled.text
assert gateway.get("/config/block_requests_for_models_without_pricing") == {"enabled": True}
settings_rows: Final = read_rows(
'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', ("litellm_settings",)
)
settings: Final = object_value(settings_rows[0]["param_value"])
assert settings["block_requests_for_models_without_pricing"] is True, settings
unpriced_upstream_model: Final = f"integration-unpriced-{uuid.uuid4().hex}"
unpriced: Final = scenario.model(model=f"openai/{unpriced_upstream_model}")
priced: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
key: Final = scenario.key(models=[unpriced, priced])
entries: Final = gateway.get("/model/info")["data"]
assert isinstance(entries, list)
unpriced_info: Final = next(
object_value(entry) for entry in entries if object_value(entry)["model_name"] == unpriced
)
unpriced_model_info: Final = object_value(unpriced_info["model_info"])
assert unpriced_model_info["input_cost_per_token"] is None, unpriced_model_info
assert unpriced_model_info["output_cost_per_token"] is None, unpriced_model_info
_observed(upstream)
blocked: Final = _bounded_chat(gateway, unpriced, key, f"unpriced {uuid.uuid4().hex}")
assert blocked.status_code == 403, blocked.text
assert blocked.json() == {
"error": {
"message": f"Model '{unpriced}' has no pricing in the cost map, so litellm cannot price the request. "
"Requests for unpriced models are blocked because 'block_requests_for_models_without_pricing' is "
"enabled. Add pricing (input_cost_per_token/output_cost_per_token) to allow the request.",
"type": "model_cost_map_missing",
"param": "model",
"code": "403",
}
}, blocked.text
assert _observed(upstream) == [], blocked.text
priced_content: Final = f"priced {uuid.uuid4().hex}"
served: Final = _bounded_chat(gateway, priced, key, priced_content)
assert served.status_code == 200, served.text
assert float(served.headers["x-litellm-response-cost"]) == pytest.approx(0.06), served.text
assert _observed(upstream) == [_upstream_chat("gpt-4o-mini", priced_content)], served.text
disabled: Final = _block_unpriced(gateway, False)
assert disabled.status_code == 200, disabled.text
assert disabled.json() == {"enabled": False}, disabled.text
assert gateway.get("/config/block_requests_for_models_without_pricing") == {"enabled": False}
settings_rows: Final = read_rows(
'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', ("litellm_settings",)
)
disabled_settings: Final = object_value(settings_rows[0]["param_value"])
assert disabled_settings["block_requests_for_models_without_pricing"] is False, disabled_settings
_observed(upstream)
unpriced_content: Final = f"unpriced {uuid.uuid4().hex}"
now_served: Final = _bounded_chat(gateway, unpriced, key, unpriced_content)
assert now_served.status_code == 200, now_served.text
assert _observed(upstream) == [_upstream_chat(unpriced_upstream_model, unpriced_content)], now_served.text

View file

@ -0,0 +1,320 @@
import json
import uuid
from collections.abc import Mapping
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, eventually, object_value, string_value
from integration._support.database import read_rows
from pydantic import BaseModel, JsonValue
class CostEstimate(BaseModel):
model: str
cost_per_request: float
input_cost_per_request: float
output_cost_per_request: float
margin_cost_per_request: float
def _settings_row() -> dict[str, JsonValue]:
rows: Final = read_rows('SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', ("litellm_settings",))
assert len(rows) == 1, rows
value: Final = rows[0]["param_value"]
return object_value(json.loads(value) if isinstance(value, str) else value)
def _patch(gateway: Gateway, path: str, body: Mapping[str, JsonValue], message: str) -> dict[str, JsonValue]:
response: Final = gateway.request("PATCH", path, body)
assert response.status_code == 200, response.text
parsed: Final = object_value(response.json())
assert parsed == {"message": message, "status": "success", "values": dict(body)}, response.text
return parsed
class ChatChunk(BaseModel):
id: str
object: str
def _observed(upstream: httpx.Client) -> list[JsonValue]:
response: Final = upstream.get("/__observations")
response.raise_for_status()
requests: Final = object_value(response.json())["requests"]
assert isinstance(requests, list), response.text
return requests
def _upstream_chat(content: str, streamed: bool) -> dict[str, JsonValue]:
return {
"path": "/v1/chat/completions",
"authorization": "Bearer integration-provider-key",
"body": {
"model": "gpt-4o-mini",
"max_tokens": 20,
"messages": [{"role": "user", "content": content}],
**({"stream": True, "stream_options": {"include_usage": True}} if streamed else {}),
},
"method": "POST",
"api_key": "",
}
def _chat_cost(gateway: Gateway, upstream: httpx.Client, model: str, key: str, expected: float) -> str:
content: Final = f"margin {uuid.uuid4().hex}"
_observed(upstream)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "max_tokens": 20, "messages": [{"role": "user", "content": content}]},
key=key,
)
assert response.status_code == 200, response.text
request_id: Final = string_value(object_value(response.json())["id"])
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected), response.text
assert _observed(upstream) == [_upstream_chat(content, streamed=False)], response.text
return request_id
def _streamed_chat(gateway: Gateway, upstream: httpx.Client, model: str, key: str) -> str:
content: Final = f"margin stream {uuid.uuid4().hex}"
_observed(upstream)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "max_tokens": 20, "stream": True, "messages": [{"role": "user", "content": content}]},
key=key,
)
assert response.status_code == 200, response.text
events: Final = [line.removeprefix("data: ") for line in response.text.splitlines() if line.startswith("data: ")]
assert events[-1] == "[DONE]", response.text
chunks: Final = [ChatChunk.model_validate_json(event) for event in events[:-1]]
assert {chunk.object for chunk in chunks} == {"chat.completion.chunk"}, response.text
assert len({chunk.id for chunk in chunks}) == 1, response.text
assert _observed(upstream) == [_upstream_chat(content, streamed=True)], response.text
return chunks[0].id
def _assert_logged_spends(expectations: tuple[tuple[str, float], ...]) -> None:
placeholders: Final = ", ".join("%s" for _ in expectations)
request_ids: Final = tuple(request_id for request_id, _ in expectations)
rows: Final = eventually(
lambda: read_rows(
f'SELECT request_id, spend FROM "LiteLLM_SpendLogs" WHERE request_id IN ({placeholders})', request_ids
),
lambda values: len(values) == len(expectations),
seconds=40,
)
actual: Final = {string_value(row["request_id"]): float(str(row["spend"])) for row in rows}
assert actual == pytest.approx(dict(expectations)), rows
def _estimated(gateway: Gateway, model: str) -> CostEstimate:
response: Final = gateway.request(
"POST", "/cost/estimate", {"model": model, "input_tokens": 20, "output_tokens": 20}
)
assert response.status_code == 200, response.text
return CostEstimate.model_validate(response.json())
def test_discount_and_margin_reprice_a_request_and_survive_rejection(gateway: Gateway) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
original_discount: Final = gateway.get("/config/cost_discount_config")["values"]
original_margin: Final = gateway.get("/config/cost_margin_config")["values"]
def restore() -> None:
_patch(
gateway,
"/config/cost_discount_config",
original_discount,
"Cost discount configuration updated successfully",
)
_patch(
gateway,
"/config/cost_margin_config",
original_margin,
"Cost margin configuration updated successfully",
)
scenario.cleanups.callback(restore)
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
key: Final = scenario.key(models=[model])
_patch(gateway, "/config/cost_discount_config", {}, "Cost discount configuration updated successfully")
_patch(
gateway,
"/config/cost_margin_config",
{},
"Cost margin configuration updated successfully",
)
base_request_id: Final = _chat_cost(gateway, upstream, model, key, 0.06)
_patch(
gateway,
"/config/cost_discount_config",
{"openai": 0.05},
"Cost discount configuration updated successfully",
)
discounted_request_id: Final = _chat_cost(gateway, upstream, model, key, 0.057)
_patch(
gateway,
"/config/cost_margin_config",
{"openai": 0.05},
"Cost margin configuration updated successfully",
)
bare_margin_request_id: Final = _chat_cost(gateway, upstream, model, key, 0.06 * 0.95 * 1.05)
_patch(
gateway,
"/config/cost_margin_config",
{"openai": {"fixed_amount": 0.001}},
"Cost margin configuration updated successfully",
)
fixed_margin_request_id: Final = _chat_cost(gateway, upstream, model, key, 0.057 + 0.001)
_patch(
gateway,
"/config/cost_margin_config",
{"openai": {"percentage": 0.1, "fixed_amount": 0.001}, "global": 0.05},
"Cost margin configuration updated successfully",
)
provider_margin_request_id: Final = _chat_cost(gateway, upstream, model, key, 0.057 * 1.1 + 0.001)
streamed_margin_request_id: Final = _streamed_chat(gateway, upstream, model, key)
assert gateway.get("/config/cost_discount_config")["values"] == {"openai": 0.05}
assert gateway.get("/config/cost_margin_config")["values"] == {
"openai": {"percentage": 0.1, "fixed_amount": 0.001},
"global": 0.05,
}
persisted: Final = _settings_row()
assert persisted["cost_discount_config"] == {"openai": 0.05}, persisted
assert persisted["cost_margin_config"] == {
"openai": {"percentage": 0.1, "fixed_amount": 0.001},
"global": 0.05,
}, persisted
expected: Final = 0.057 * 1.1 + 0.001
estimate: Final = _estimated(gateway, model)
assert estimate.model == model
assert estimate.cost_per_request == pytest.approx(expected), estimate
assert estimate.input_cost_per_request == pytest.approx(0.02), estimate
assert estimate.output_cost_per_request == pytest.approx(0.04), estimate
discounted_subtotal: Final = estimate.cost_per_request - estimate.margin_cost_per_request
assert (
estimate.input_cost_per_request + estimate.output_cost_per_request - discounted_subtotal
== pytest.approx(0.06 * 0.05)
), estimate
assert estimate.margin_cost_per_request == pytest.approx(0.06 * 0.95 * 0.1 + 0.001), estimate
_patch(gateway, "/config/cost_discount_config", {}, "Cost discount configuration updated successfully")
no_discount_request_id: Final = _chat_cost(gateway, upstream, model, key, 0.06 * 1.1 + 0.001)
_patch(
gateway,
"/config/cost_margin_config",
{"global": 1},
"Cost margin configuration updated successfully",
)
global_margin_request_id: Final = _chat_cost(gateway, upstream, model, key, 0.12)
streamed_global_request_id: Final = _streamed_chat(gateway, upstream, model, key)
_assert_logged_spends(
(
(base_request_id, 0.06),
(discounted_request_id, 0.057),
(bare_margin_request_id, 0.06 * 0.95 * 1.05),
(fixed_margin_request_id, 0.057 + 0.001),
(provider_margin_request_id, 0.057 * 1.1 + 0.001),
(streamed_margin_request_id, 0.057 * 1.1 + 0.001),
(no_discount_request_id, 0.06 * 1.1 + 0.001),
(global_margin_request_id, 0.12),
(streamed_global_request_id, 0.12),
)
)
rejections: Final = (
(
"/config/cost_margin_config",
{"not_a_provider": 0.1},
{
"detail": {
"error": "Invalid provider(s): not_a_provider. Must be valid LiteLLM providers or 'global'. "
"See https://docs.litellm.ai/docs/providers for the full list."
}
},
),
(
"/config/cost_discount_config",
{"openai": 1.5},
{"detail": "Discount for openai must be between 0 and 1 (0% to 100%)"},
),
(
"/config/cost_discount_config",
{"global": 0.1},
{
"detail": {
"error": "Invalid provider(s): global. Must be valid LiteLLM providers. "
"See https://docs.litellm.ai/docs/providers for the full list."
}
},
),
)
for path, body, error in rejections:
rejected = gateway.request("PATCH", path, body)
assert rejected.status_code == 400, rejected.text
assert rejected.json() == error, rejected.text
assert gateway.get("/config/cost_discount_config")["values"] == {}
assert gateway.get("/config/cost_margin_config")["values"] == {"global": 1}
final_settings: Final = _settings_row()
assert final_settings["cost_discount_config"] == {}, final_settings
assert final_settings["cost_margin_config"] == {"global": 1}, final_settings
def test_streamed_response_cost_header_matches_billed_spend(gateway: Gateway) -> None:
pytest.skip(
"BUG: a streamed chat with a fixed cost margin returns x-litellm-response-cost equal to the fixed amount "
"(computed before usage is known) while LiteLLM_SpendLogs bills the full margined cost"
)
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
original_margin: Final = gateway.get("/config/cost_margin_config")["values"]
scenario.cleanups.callback(
_patch,
gateway,
"/config/cost_margin_config",
original_margin,
"Cost margin configuration updated successfully",
)
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
key: Final = scenario.key(models=[model])
_patch(
gateway,
"/config/cost_margin_config",
{"openai": {"percentage": 0.1, "fixed_amount": 0.001}},
"Cost margin configuration updated successfully",
)
content: Final = f"margin stream header {uuid.uuid4().hex}"
_observed(upstream)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "max_tokens": 20, "stream": True, "messages": [{"role": "user", "content": content}]},
key=key,
)
assert response.status_code == 200, response.text
events: Final = [
line.removeprefix("data: ") for line in response.text.splitlines() if line.startswith("data: ")
]
assert events[-1] == "[DONE]", response.text
request_id: Final = ChatChunk.model_validate_json(events[0]).id
assert _observed(upstream) == [_upstream_chat(content, streamed=True)], response.text
_assert_logged_spends(((request_id, 0.06 * 1.1 + 0.001),))
advertised: Final = response.headers.get("x-litellm-response-cost")
assert advertised is None or float(advertised) == pytest.approx(0.06 * 1.1 + 0.001), dict(response.headers)

View file

@ -0,0 +1,438 @@
import json
import re
import uuid
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
from integration._support.database import read_rows
from integration._support.upstream import delete_scenario, register_scenario
from integration.cost_calculation.cost_tracking_case import JsonResponse
from pydantic import BaseModel, JsonValue
PROVIDER_AUTHORIZATION: Final = "Bearer integration-provider-key"
CANNED_CONTENT: Final = "Hello! This is a mock response from the fake OpenAI endpoint."
RESPONSES_REPLY: Final = JsonResponse(
content_type="application/json",
body={
"id": "resp_$UNIQUE_ID",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-4o-mini",
"output": [
{
"type": "message",
"id": "msg_$UNIQUE_ID",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "ok", "annotations": []}],
}
],
"usage": {"input_tokens": 20, "output_tokens": 20, "total_tokens": 40},
},
)
class ChatCompletion(BaseModel):
id: str
object: str
class ChatChunkDelta(BaseModel):
content: str | None = None
class ChatChunkChoice(BaseModel):
delta: ChatChunkDelta
class ChatChunk(BaseModel):
object: str
choices: list[ChatChunkChoice]
class ResponsesResult(BaseModel):
id: str
object: str
status: str
class AnthropicMessage(BaseModel):
type: str
role: str
content: list[dict[str, JsonValue]]
@dataclass(frozen=True, slots=True)
class Spelling:
route: str
fields: Callable[[str], dict[str, JsonValue]]
headers: Callable[[str], dict[str, str]]
upstream_fields: Callable[[str], dict[str, JsonValue]]
stream: bool = False
def _no_fields(_: str) -> dict[str, JsonValue]:
return {}
def _no_headers(_: str) -> dict[str, str]:
return {}
SPELLINGS: Final[Mapping[str, Spelling]] = {
"user": Spelling("/v1/chat/completions", lambda user: {"user": user}, _no_headers, lambda user: {"user": user}),
"x-litellm-customer-id": Spelling(
"/v1/chat/completions", _no_fields, lambda user: {"x-litellm-customer-id": user}, _no_fields
),
"x-litellm-end-user-id": Spelling(
"/v1/chat/completions", _no_fields, lambda user: {"x-litellm-end-user-id": user}, _no_fields
),
"litellm_metadata": Spelling(
"/v1/chat/completions", lambda user: {"litellm_metadata": {"user": user}}, _no_headers, _no_fields
),
"litellm_metadata_json": Spelling(
"/v1/chat/completions",
lambda user: {"litellm_metadata": json.dumps({"user": user})},
_no_headers,
_no_fields,
),
"metadata": Spelling("/v1/chat/completions", lambda user: {"metadata": {"user_id": user}}, _no_headers, _no_fields),
"metadata_json": Spelling(
"/v1/chat/completions", lambda user: {"metadata": json.dumps({"user_id": user})}, _no_headers, _no_fields
),
"safety_identifier": Spelling(
"/v1/chat/completions",
lambda user: {"safety_identifier": user},
_no_headers,
lambda user: {"safety_identifier": user},
),
"user_stream": Spelling(
"/v1/chat/completions",
lambda user: {"user": user, "stream": True},
_no_headers,
lambda user: {"user": user, "stream": True, "stream_options": {"include_usage": True}},
stream=True,
),
"litellm_metadata_stream": Spelling(
"/v1/chat/completions",
lambda user: {"litellm_metadata": {"user": user}, "stream": True},
_no_headers,
lambda _: {"stream": True, "stream_options": {"include_usage": True}},
stream=True,
),
"responses_user": Spelling("/v1/responses", lambda user: {"user": user}, _no_headers, lambda user: {"user": user}),
"responses_safety_identifier": Spelling(
"/v1/responses",
lambda user: {"safety_identifier": user},
_no_headers,
lambda user: {"safety_identifier": user},
),
"messages_metadata_user_id": Spelling(
"/v1/messages", lambda user: {"metadata": {"user_id": user}}, _no_headers, lambda user: {"user": user}
),
}
@dataclass(frozen=True, slots=True)
class Deployment:
model: str
upstream_prefix: str
def _deployment(scenario: Scenario, spelling: Spelling) -> Deployment:
if spelling.route == "/v1/chat/completions":
return Deployment(scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002), "")
scenario_id: Final = f"enduser{uuid.uuid4().hex[:12]}"
handle: Final = register_scenario(scenario_id, RESPONSES_REPLY)
scenario.cleanups.callback(delete_scenario, handle)
model: Final = scenario.model(
api_base=f"{handle.api_base()}/v1", input_cost_per_token=0.001, output_cost_per_token=0.002
)
return Deployment(model, f"/{scenario_id}")
def _client_body(spelling: Spelling, model: str, user_id: str, content: str) -> dict[str, JsonValue]:
match spelling.route:
case "/v1/chat/completions":
base: dict[str, JsonValue] = {
"model": model,
"max_tokens": 20,
"messages": [{"role": "user", "content": content}],
}
case "/v1/responses":
base = {"model": model, "max_output_tokens": 20, "input": content}
case _:
base = {"model": model, "max_tokens": 20, "messages": [{"role": "user", "content": content}]}
return {**base, **spelling.fields(user_id)}
def _upstream_request(spelling: Spelling, prefix: str, user_id: str, content: str) -> dict[str, JsonValue]:
match spelling.route:
case "/v1/chat/completions":
path = "/v1/chat/completions"
body: dict[str, JsonValue] = {
"model": "gpt-4o-mini",
"max_tokens": 20,
"messages": [{"role": "user", "content": content}],
}
case "/v1/responses":
path = f"{prefix}/v1/responses"
body = {"model": "gpt-4o-mini", "input": content, "max_output_tokens": 20}
case _:
path = f"{prefix}/v1/responses"
body = {
"model": "gpt-4o-mini",
"input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": content}]}],
"include": ["reasoning.encrypted_content"],
"max_output_tokens": 20,
}
return {
"path": path,
"authorization": PROVIDER_AUTHORIZATION,
"body": {**body, **spelling.upstream_fields(user_id)},
"method": "POST",
"api_key": "",
}
def _send(
gateway: Gateway, spelling: Spelling, deployment: Deployment, key: str, user_id: str, content: str
) -> httpx.Response:
return gateway.request(
"POST",
spelling.route,
_client_body(spelling, deployment.model, user_id, content),
key=key,
headers=spelling.headers(user_id),
)
def _assert_served(spelling: Spelling, response: httpx.Response) -> None:
assert response.status_code == 200, response.text
if spelling.stream:
events: Final = [
line.removeprefix("data: ") for line in response.text.splitlines() if line.startswith("data: ")
]
assert events[-1] == "[DONE]", response.text
chunks: Final = [ChatChunk.model_validate_json(event) for event in events[:-1]]
assert {chunk.object for chunk in chunks} == {"chat.completion.chunk"}, response.text
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == CANNED_CONTENT
return
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(0.06), dict(response.headers)
match spelling.route:
case "/v1/chat/completions":
assert ChatCompletion.model_validate_json(response.content).object == "chat.completion", response.text
case "/v1/responses":
parsed: Final = ResponsesResult.model_validate_json(response.content)
assert (parsed.object, parsed.status) == ("response", "completed"), response.text
case _:
message: Final = AnthropicMessage.model_validate_json(response.content)
assert (message.type, message.role, message.content) == (
"message",
"assistant",
[{"type": "text", "text": "ok"}],
), response.text
def _delete_customer(gateway: Gateway, user_id: str) -> None:
if not _end_user_spend(user_id):
return
response: Final = gateway.request("POST", "/customer/delete", {"user_ids": [user_id]})
assert response.status_code == 200, response.text
def _end_user_spend(user_id: str) -> list[dict[str, JsonValue]]:
return read_rows('SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id = %s', (user_id,))
def _observed(upstream: httpx.Client) -> list[JsonValue]:
response: Final = upstream.get("/__observations")
response.raise_for_status()
requests: Final = object_value(response.json())["requests"]
assert isinstance(requests, list), response.text
return requests
def _denial_error(denied: httpx.Response) -> dict[str, JsonValue]:
assert denied.status_code == 422, denied.text
return object_value(object_value(denied.json())["error"])
def _prove_end_user_budget_lockout(
gateway: Gateway,
upstream: httpx.Client,
spelling: Spelling,
deployment: Deployment,
key: str,
end_user_id: str,
other_end_user_id: str,
expected_budget: Mapping[str, JsonValue],
info_path: str = "/customer/info",
) -> None:
served_content: Final = f"served {end_user_id}"
_observed(upstream)
served: Final = _send(gateway, spelling, deployment, key, end_user_id, served_content)
_assert_served(spelling, served)
assert _observed(upstream) == [
_upstream_request(spelling, deployment.upstream_prefix, end_user_id, served_content)
], served.text
eventually(
lambda: _end_user_spend(end_user_id),
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > 0.05,
seconds=30,
)
persisted_spend: Final = float(str(_end_user_spend(end_user_id)[0]["spend"]))
assert persisted_spend == pytest.approx(0.06), persisted_spend
customer_info: Final = gateway.request("GET", info_path, params={"end_user_id": end_user_id})
assert customer_info.status_code == 200, customer_info.text
customer: Final = object_value(customer_info.json())
assert customer["user_id"] == end_user_id, customer_info.text
assert float(str(customer["spend"])) == pytest.approx(persisted_spend), customer_info.text
assert customer["budget_id"] == expected_budget["budget_id"], customer_info.text
budget: Final = customer["litellm_budget_table"]
expected_max_budget: Final = expected_budget["max_budget"]
if expected_max_budget is None:
assert budget is None, customer_info.text
else:
assert float(str(object_value(budget)["max_budget"])) == float(str(expected_max_budget)), customer_info.text
_observed(upstream)
denied: Final = _send(gateway, spelling, deployment, key, end_user_id, f"denied {end_user_id}")
assert _observed(upstream) == [], denied.text
error: Final = _denial_error(denied)
assert error["type"] == "budget_exceeded", denied.text
matched: Final = re.fullmatch(
rf"ExceededBudget: End User={re.escape(end_user_id)} over budget\. "
rf"Spend=([-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?), Budget=0\.05",
string_value(error["message"]),
)
assert matched is not None, denied.text
assert float(matched.group(1)) == pytest.approx(persisted_spend), denied.text
other_content: Final = f"served {other_end_user_id}"
other: Final = _send(gateway, spelling, deployment, key, other_end_user_id, other_content)
_assert_served(spelling, other)
assert _observed(upstream) == [
_upstream_request(spelling, deployment.upstream_prefix, other_end_user_id, other_content)
], other.text
other_spend: Final = eventually(
lambda: _end_user_spend(other_end_user_id),
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > 0.05,
seconds=30,
)
assert float(str(other_spend[0]["spend"])) == pytest.approx(0.06), other_spend
@pytest.mark.parametrize("spelling", tuple(SPELLINGS))
def test_customer_max_budget_refuses_second_request(gateway: Gateway, spelling: str) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
variant: Final = SPELLINGS[spelling]
deployment: Final = _deployment(scenario, variant)
key: Final = scenario.key(models=[deployment.model])
end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
created: Final = gateway.post("/customer/new", {"user_id": end_user_id, "max_budget": 0.05})
assert created["user_id"] == end_user_id
scenario.cleanups.callback(_delete_customer, gateway, end_user_id)
other_end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
scenario.cleanups.callback(_delete_customer, gateway, other_end_user_id)
_prove_end_user_budget_lockout(
gateway,
upstream,
variant,
deployment,
key,
end_user_id,
other_end_user_id,
{"budget_id": string_value(created["budget_id"]), "max_budget": 0.05},
)
def test_end_user_alias_routes_create_and_read_a_budgeted_customer(gateway: Gateway) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
variant: Final = SPELLINGS["user"]
deployment: Final = _deployment(scenario, variant)
key: Final = scenario.key(models=[deployment.model])
end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
created: Final = gateway.post("/end_user/new", {"user_id": end_user_id, "max_budget": 0.05})
assert created["user_id"] == end_user_id, created
assert float(str(object_value(created["litellm_budget_table"])["max_budget"])) == 0.05, created
scenario.cleanups.callback(_delete_customer, gateway, end_user_id)
other_end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
scenario.cleanups.callback(_delete_customer, gateway, other_end_user_id)
_prove_end_user_budget_lockout(
gateway,
upstream,
variant,
deployment,
key,
end_user_id,
other_end_user_id,
{"budget_id": string_value(created["budget_id"]), "max_budget": 0.05},
info_path="/end_user/info",
)
@pytest.mark.parametrize("spelling", ("litellm_metadata", "litellm_metadata_stream"))
def test_customer_tier_budget_refuses_second_request(gateway: Gateway, spelling: str) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
variant: Final = SPELLINGS[spelling]
deployment: Final = _deployment(scenario, variant)
tier: Final = scenario.budget(max_budget=0.05)
key: Final = scenario.key(models=[deployment.model])
end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
created: Final = gateway.post("/customer/new", {"user_id": end_user_id, "budget_id": tier})
assert created["user_id"] == end_user_id
scenario.cleanups.callback(_delete_customer, gateway, end_user_id)
other_end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
scenario.cleanups.callback(_delete_customer, gateway, other_end_user_id)
_prove_end_user_budget_lockout(
gateway,
upstream,
variant,
deployment,
key,
end_user_id,
other_end_user_id,
{"budget_id": tier, "max_budget": 0.05},
)
def test_key_end_user_budget_id_refuses_second_request(gateway: Gateway) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
variant: Final = SPELLINGS["metadata"]
deployment: Final = _deployment(scenario, variant)
tier: Final = scenario.budget(max_budget=0.05)
key: Final = scenario.key(models=[deployment.model], metadata={"end_user_budget_id": tier})
key_info: Final = object_value(gateway.get("/key/info", {"key": key})["info"])
assert object_value(key_info["metadata"]) == {"end_user_budget_id": tier}, key_info
end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
other_end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
scenario.cleanups.callback(_delete_customer, gateway, end_user_id)
scenario.cleanups.callback(_delete_customer, gateway, other_end_user_id)
_prove_end_user_budget_lockout(
gateway,
upstream,
variant,
deployment,
key,
end_user_id,
other_end_user_id,
{"budget_id": None, "max_budget": None},
)

View file

@ -0,0 +1,159 @@
import uuid
from hashlib import sha256
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, eventually, object_value
from integration._support.database import read_rows
from pydantic import JsonValue
def _chat(gateway: Gateway, model: str, key: str, content: str | None = None) -> httpx.Response:
return gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"max_tokens": 20,
"messages": [{"role": "user", "content": content or f"reset {uuid.uuid4().hex}"}],
},
key=key,
)
def _observed(upstream: httpx.Client) -> list[JsonValue]:
response: Final = upstream.get("/__observations")
response.raise_for_status()
requests: Final = object_value(response.json())["requests"]
assert isinstance(requests, list), response.text
return requests
def _upstream_chat(content: str) -> dict[str, JsonValue]:
return {
"path": "/v1/chat/completions",
"authorization": "Bearer integration-provider-key",
"body": {"model": "gpt-4o-mini", "max_tokens": 20, "messages": [{"role": "user", "content": content}]},
"method": "POST",
"api_key": "",
}
def _token_spend(key: str) -> list[dict[str, object]]:
return read_rows(
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', (sha256(key.encode()).hexdigest(),)
) # fmt: skip
def _team_spend(team_id: str) -> list[dict[str, object]]:
return read_rows('SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,))
def _spend_log_totals(keys: tuple[str, ...]) -> list[dict[str, object]]:
digests: Final = tuple(sha256(key.encode()).hexdigest() for key in keys)
return read_rows(
'SELECT COUNT(*) AS rows, COALESCE(SUM(spend), 0) AS total_spend FROM "LiteLLM_SpendLogs" '
"WHERE api_key = ANY(%s::text[])",
("{" + ",".join(digests) + "}",),
)
def _seed_spend(gateway: Gateway, model: str, team_id: str, key_a: str, key_b: str) -> None:
assert _chat(gateway, model, key_a).status_code == 200
assert _chat(gateway, model, key_b).status_code == 200
for key in (key_a, key_b):
eventually(
lambda key=key: _token_spend(key),
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > 0,
seconds=70,
)
eventually(
lambda: _team_spend(team_id),
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > 0,
seconds=70,
)
def test_global_spend_reset_zeroes_token_and_team_spend_but_keeps_logs(gateway: Gateway) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
team_id: Final = scenario.team()
key_a: Final = scenario.key(team_id=team_id, models=[model], max_budget=0.06)
key_b: Final = scenario.key(team_id=team_id, models=[model])
internal_user: Final = scenario.user(user_role="internal_user")
internal_key: Final = scenario.key(user_id=internal_user, models=[model])
_seed_spend(gateway, model, team_id, key_a, key_b)
before: Final = eventually(
lambda: _spend_log_totals((key_a, key_b)),
lambda rows: (
len(rows) == 1
and int(str(rows[0]["rows"])) == 2
and float(str(rows[0]["total_spend"])) == pytest.approx(0.12)
),
seconds=70,
)
assert float(str(before[0]["total_spend"])) == pytest.approx(0.12), before
refused: Final = gateway.request("POST", "/global/spend/reset", key=internal_key)
assert refused.status_code == 401, refused.text
assert refused.json() == {
"error": {
"message": "Authentication Error, Tried to access route=/global/spend/reset, which is only for MASTER KEY",
"type": "auth_error",
"param": "None",
"code": "401",
}
}, refused.text
assert float(str(_token_spend(key_a)[0]["spend"])) == pytest.approx(0.06), _token_spend(key_a)
assert float(str(_token_spend(key_b)[0]["spend"])) == pytest.approx(0.06), _token_spend(key_b)
assert float(str(_team_spend(team_id)[0]["spend"])) == pytest.approx(0.12), _team_spend(team_id)
_observed(upstream)
denied: Final = _chat(gateway, model, key_a)
assert denied.status_code == 422, denied.text
assert object_value(denied.json()["error"])["type"] == "budget_exceeded", denied.text
assert _observed(upstream) == [], denied.text
reset: Final = gateway.request("POST", "/global/spend/reset")
assert reset.status_code == 200, reset.text
assert reset.json() == {
"message": "Spend for all API Keys and Teams reset successfully",
"status": "success",
}, reset.text
for key in (key_a, key_b):
rows: Final = _token_spend(key)
assert float(str(rows[0]["spend"])) == 0, rows
assert float(str(_team_spend(team_id)[0]["spend"])) == 0
assert _spend_log_totals((key_a, key_b)) == before
def test_reset_exhausted_key_is_served_again(gateway: Gateway) -> None:
pytest.skip(
"BUG: /global/spend/reset zeroes DB spend but leaves Redis spend:key counters, "
"so an exhausted key stays refused with budget_exceeded"
)
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
key: Final = scenario.key(models=[model], max_budget=0.06)
assert _chat(gateway, model, key).status_code == 200
eventually(
lambda: _token_spend(key),
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= 0.06,
seconds=70,
)
denied: Final = _chat(gateway, model, key)
assert denied.status_code == 422, denied.text
reset: Final = gateway.request("POST", "/global/spend/reset")
assert reset.status_code == 200, reset.text
content: Final = f"restored {uuid.uuid4().hex}"
_observed(upstream)
restored: Final = _chat(gateway, model, key, content)
assert restored.status_code == 200, restored.text
assert _observed(upstream) == [_upstream_chat(content)], restored.text

View file

@ -6,6 +6,8 @@ import httpx
import pytest
from integration._support.client import Gateway, eventually, object_value
from integration._support.database import read_rows
from openai.types.chat import ChatCompletionChunk
from pydantic import JsonValue
def test_an_exhausted_key_is_refused_inference_but_can_still_read_its_own_info(gateway: Gateway) -> None:
@ -43,19 +45,46 @@ def test_an_exhausted_key_is_refused_inference_but_can_still_read_its_own_info(g
assert own["max_budget"] == 0.06
def _bounded_chat(gateway: Gateway, model: str, key: str) -> httpx.Response:
def _bounded_chat(gateway: Gateway, model: str, key: str, content: str | None = None) -> httpx.Response:
return gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"max_tokens": 20,
"messages": [{"role": "user", "content": f"key recovery {uuid.uuid4().hex}"}],
"messages": [{"role": "user", "content": content or f"key recovery {uuid.uuid4().hex}"}],
},
key=key,
)
def _observed(upstream: httpx.Client) -> list[JsonValue]:
response: Final = upstream.get("/__observations")
response.raise_for_status()
requests: Final = object_value(response.json())["requests"]
assert isinstance(requests, list), response.text
return requests
def _upstream_chat(content: str) -> dict[str, JsonValue]:
return {
"path": "/v1/chat/completions",
"authorization": "Bearer integration-provider-key",
"body": {"model": "gpt-4o-mini", "max_tokens": 20, "messages": [{"role": "user", "content": content}]},
"method": "POST",
"api_key": "",
}
def _served_chat(gateway: Gateway, upstream: httpx.Client, model: str, key: str) -> None:
content: Final = f"tier served {uuid.uuid4().hex}"
_observed(upstream)
served: Final = _bounded_chat(gateway, model, key, content)
assert served.status_code == 200, served.text
assert float(served.headers["x-litellm-response-cost"]) == 0.06, dict(served.headers)
assert _observed(upstream) == [_upstream_chat(content)], served.text
def test_raising_a_spent_keys_budget_restores_serving(gateway: Gateway) -> None:
with (
gateway.scenario() as scenario,
@ -82,3 +111,156 @@ def test_raising_a_spent_keys_budget_restores_serving(gateway: Gateway) -> None:
served: Final = tuple(_bounded_chat(gateway, model, key) for _ in range(3))
assert [response.status_code for response in served] == [200, 200, 200], [response.text for response in served]
assert len(upstream.get("/__observations").json()["requests"]) == 3
def _denied_with_budget_exceeded(gateway: Gateway, model: str, key: str) -> None:
denied: Final = _bounded_chat(gateway, model, key)
assert denied.status_code == 422, denied.text
assert object_value(denied.json()["error"])["type"] == "budget_exceeded"
@pytest.mark.timeout(149)
def test_raising_a_spent_tier_budget_restores_serving(gateway: Gateway) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
budget_id: Final = scenario.budget(budget_id=f"tier-{uuid.uuid4().hex}", max_budget=0.06, budget_duration="30d")
key: Final = scenario.key(models=[model], budget_id=budget_id)
info: Final = object_value(gateway.get("/key/info", {"key": key})["info"])
assert info["budget_id"] == budget_id, info
tier_rows: Final = read_rows(
'SELECT max_budget, budget_duration FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (budget_id,)
)
assert len(tier_rows) == 1
assert float(str(tier_rows[0]["max_budget"])) == 0.06
assert tier_rows[0]["budget_duration"] == "30d"
_served_chat(gateway, upstream, model, key)
eventually(
lambda: read_rows(
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),)
),
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= 0.06,
seconds=70,
)
upstream.get("/__observations").raise_for_status()
_denied_with_budget_exceeded(gateway, model, key)
assert upstream.get("/__observations").json()["requests"] == []
gateway.post("/budget/update", {"budget_id": budget_id, "max_budget": 1.0})
updated_rows: Final = read_rows(
'SELECT max_budget FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (budget_id,)
)
assert float(str(updated_rows[0]["max_budget"])) == 1.0, updated_rows
restored_content: Final = f"tier restored {uuid.uuid4().hex}"
served: Final = eventually(
lambda: _bounded_chat(gateway, model, key, restored_content),
lambda response: response.status_code == 200,
seconds=125,
)
assert served.status_code == 200, served.text
assert _observed(upstream) == [_upstream_chat(restored_content)], served.text
@pytest.mark.timeout(149)
def test_lowering_a_tier_budget_below_spend_blocks_the_key(gateway: Gateway) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
budget_id: Final = scenario.budget(
budget_id=f"tier-{uuid.uuid4().hex}", max_budget=1000.0, budget_duration="30d"
)
key: Final = scenario.key(models=[model], budget_id=budget_id)
_served_chat(gateway, upstream, model, key)
eventually(
lambda: read_rows(
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),)
),
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= 0.06,
seconds=30,
)
upstream.get("/__observations").raise_for_status()
probe: Final = gateway.request("POST", "/v1/chat/completions", {"model": model}, key=key)
assert probe.status_code == 400, probe.text
assert object_value(probe.json()["error"]) == {
"message": "/chat/completions: Missing required parameter: 'messages'.",
"type": "invalid_request_error",
"param": "messages",
"code": "400",
}, probe.text
assert upstream.get("/__observations").json()["requests"] == []
gateway.post("/budget/update", {"budget_id": budget_id, "max_budget": 0.01})
updated_rows: Final = read_rows(
'SELECT max_budget FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (budget_id,)
)
assert len(updated_rows) == 1
assert float(str(updated_rows[0]["max_budget"])) == 0.01, updated_rows
upstream.get("/__observations").raise_for_status()
lowered_probe: Final = eventually(
lambda: gateway.request("POST", "/v1/chat/completions", {"model": model}, key=key),
lambda response: (
response.status_code == 422 and object_value(response.json()["error"])["type"] == "budget_exceeded"
),
seconds=125,
)
assert "Max budget: 0.01" in lowered_probe.text, lowered_probe.text
denied: Final = _bounded_chat(gateway, model, key)
assert denied.status_code == 422, denied.text
assert object_value(denied.json()["error"])["type"] == "budget_exceeded", denied.text
assert "Max budget: 0.01" in denied.text, denied.text
assert upstream.get("/__observations").json()["requests"] == []
@pytest.mark.timeout(149)
def test_key_linked_to_a_tier_by_update_is_refused_once_streamed_spend_exhausts_it(gateway: Gateway) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
budget_id: Final = scenario.budget(budget_id=f"tier-{uuid.uuid4().hex}", max_budget=0.06, budget_duration="30d")
key: Final = scenario.key(models=[model])
gateway.post("/key/update", {"key": key, "budget_id": budget_id})
info: Final = object_value(gateway.get("/key/info", {"key": key})["info"])
assert info["budget_id"] == budget_id, info
content: Final = f"tier streamed {uuid.uuid4().hex}"
body: Final = {
"model": model,
"max_tokens": 20,
"messages": [{"role": "user", "content": content}],
"stream": True,
}
_observed(upstream)
served: Final = gateway.request("POST", "/v1/chat/completions", body, key=key)
assert served.status_code == 200, served.text
events: Final = tuple(
line.removeprefix("data: ") for line in served.text.splitlines() if line.startswith("data: ")
)
assert events[-1] == "[DONE]", served.text
chunks: Final = tuple(ChatCompletionChunk.model_validate_json(event) for event in events[:-1])
assert {chunk.object for chunk in chunks} == {"chat.completion.chunk"}, served.text
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices), served.text
observations: Final = _observed(upstream)
assert len(observations) == 1, observations
upstream_body: Final = object_value(object_value(observations[0])["body"])
assert upstream_body["stream"] is True, observations
denied: Final = eventually(
lambda: gateway.request("POST", "/v1/chat/completions", body, key=key),
lambda response: (
response.status_code == 422 and object_value(response.json()["error"])["type"] == "budget_exceeded"
),
seconds=125,
)
assert denied.status_code == 422, denied.text
assert object_value(denied.json()["error"])["type"] == "budget_exceeded", denied.text
assert "Max budget: 0.06" in denied.text, denied.text
_observed(upstream)
fresh_denied: Final = gateway.request("POST", "/v1/chat/completions", body, key=key)
assert fresh_denied.status_code == 422, fresh_denied.text
assert object_value(fresh_denied.json()["error"])["type"] == "budget_exceeded", fresh_denied.text
assert "Max budget: 0.06" in fresh_denied.text, fresh_denied.text
assert _observed(upstream) == []

View file

@ -0,0 +1,228 @@
import uuid
from collections.abc import Mapping
from datetime import date, datetime, time, timedelta
from hashlib import sha256
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, eventually, object_value
from integration._support.database import read_rows
from pydantic import BaseModel
class ModelDetail(BaseModel):
model: str
total_cost: float
total_input_tokens: float
total_output_tokens: float
team_id: str | None = None
class ReportRow(BaseModel):
api_key: str
total_cost: float
total_input_tokens: float
total_output_tokens: float
model_details: list[ModelDetail]
def _chat(gateway: Gateway, model: str, key: str) -> None:
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "max_tokens": 20, "messages": [{"role": "user", "content": f"report {uuid.uuid4().hex}"}]},
key=key,
)
assert response.status_code == 200, response.text
def _report(
gateway: Gateway, path: str, *, key: str | None = None, params: Mapping[str, str] | None = None
) -> httpx.Response:
return gateway.request("GET", path, key=key, params=params)
def _rows(
gateway: Gateway, path: str, *, key: str | None = None, params: Mapping[str, str] | None = None
) -> list[ReportRow]:
response: Final = _report(gateway, path, key=key, params=params)
assert response.status_code == 200, response.text
return [ReportRow.model_validate(row) for row in response.json()]
def _logged(scope_column: str, scope_value: str, start: datetime, end: datetime) -> list[dict[str, object]]:
return read_rows(
f'SELECT api_key, model, spend, prompt_tokens, completion_tokens, team_id FROM "LiteLLM_SpendLogs" '
f'WHERE "startTime" >= %s AND "startTime" < %s AND {scope_column} = %s',
(start.isoformat(), end.isoformat(), scope_value),
)
def _logged_keys(keys: tuple[str, ...], start: datetime, end: datetime) -> list[dict[str, object]]:
placeholders: Final = ", ".join("%s" for _ in keys)
return read_rows(
'SELECT api_key, model, spend, prompt_tokens, completion_tokens, team_id FROM "LiteLLM_SpendLogs" '
f'WHERE "startTime" >= %s AND "startTime" < %s AND api_key IN ({placeholders})',
(start.isoformat(), end.isoformat(), *keys),
)
def _assert_report_matches_db(report: list[ReportRow], logged: list[dict[str, object]]) -> None:
assert len(report) == len({row["api_key"] for row in logged}), (report, logged)
by_key: Final[dict[str, dict[str, list[float]]]] = {}
for row in logged:
bucket: Final = by_key.setdefault(str(row["api_key"]), {})
model_bucket: Final = bucket.setdefault(str(row["model"]), [0.0, 0.0, 0.0])
model_bucket[0] += float(str(row["spend"]))
model_bucket[1] += int(str(row["prompt_tokens"]))
model_bucket[2] += int(str(row["completion_tokens"]))
for entry in report:
assert entry.api_key in by_key, (entry, logged)
models: Final = by_key[entry.api_key]
assert entry.total_cost == pytest.approx(sum(v[0] for v in models.values())), entry
assert entry.total_input_tokens == pytest.approx(sum(v[1] for v in models.values())), entry
assert entry.total_output_tokens == pytest.approx(sum(v[2] for v in models.values())), entry
details: Final = {detail.model: detail for detail in entry.model_details}
assert set(details) == set(models), (entry, logged)
for model, detail in details.items():
cost, prompt, completion = models[model]
assert detail.total_cost == pytest.approx(cost), detail
assert detail.total_input_tokens == pytest.approx(prompt), detail
assert detail.total_output_tokens == pytest.approx(completion), detail
def test_scoped_spend_reports_match_sql_and_enforce_caller_scope(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
org1: Final = scenario.organization()
org2: Final = scenario.organization()
team1: Final = scenario.team(organization_id=org1)
team2: Final = scenario.team(organization_id=org2)
user1: Final = scenario.member(team1, role="user")
user2: Final = scenario.member(team2, role="user")
org_admin: Final = scenario.org_member(org1, "org_admin")
k1: Final = scenario.key(user_id=user1, team_id=team1, models=[model])
k2: Final = scenario.key(user_id=user2, team_id=team2, models=[model])
direct_org_key: Final = scenario.key(organization_id=org2, models=[model])
org_admin_key: Final = scenario.key(user_id=org_admin, organization_id=org1, models=[model])
org_admin_info: Final = object_value(gateway.get("/key/info", {"key": org_admin_key})["info"])
assert org_admin_info["organization_id"] == org1, org_admin_info
_chat(gateway, model, k1)
_chat(gateway, model, k2)
_chat(gateway, model, direct_org_key)
k1_hash: Final = sha256(k1.encode()).hexdigest()
k2_hash: Final = sha256(k2.encode()).hexdigest()
direct_org_key_hash: Final = sha256(direct_org_key.encode()).hexdigest()
logged: Final = eventually(
lambda: read_rows(
"SELECT api_key, to_char(\"startTime\", 'YYYY-MM-DD') AS day "
'FROM "LiteLLM_SpendLogs" WHERE api_key IN (%s, %s, %s)',
(k1_hash, k2_hash, direct_org_key_hash),
),
lambda rows: len(rows) == 3,
seconds=30,
)
dates: Final = {date.fromisoformat(str(row["day"])) for row in logged}
start_date: Final[date] = min(dates)
end_date: Final[date] = max(dates)
start: Final = datetime.combine(start_date, time.min)
end: Final = datetime.combine(end_date, time.min) + timedelta(days=1)
window: Final = {"start_date": start_date.isoformat(), "end_date": end_date.isoformat()}
_assert_report_matches_db(
_rows(gateway, "/key/spend/report", key=k1, params=window), _logged("api_key", k1_hash, start, end)
)
_assert_report_matches_db(
_rows(gateway, "/user/spend/report", key=k1, params=window), _logged('"user"', user1, start, end)
)
_assert_report_matches_db(
_rows(gateway, "/team/spend/report", key=k1, params=window), _logged("team_id", team1, start, end)
)
org1_report: Final = _rows(gateway, "/organization/spend/report", key=org_admin_key, params=window)
_assert_report_matches_db(org1_report, _logged_keys((k1_hash,), start, end))
refused_org2: Final = _report(
gateway,
"/organization/spend/report",
key=org_admin_key,
params={**window, "organization_id": org2},
)
assert refused_org2.status_code == 403, refused_org2.text
assert refused_org2.json() == {"detail": "You do not have access to this organization"}, refused_org2.text
own_org: Final = _report(gateway, "/organization/spend/report", key=k1, params=window)
assert own_org.status_code == 403, own_org.text
assert own_org.json() == {"detail": "You do not have access to this organization"}, own_org.text
assert _rows(gateway, "/key/spend/report", key=k1, params={**window, "api_key": k1}) == _rows(
gateway, "/key/spend/report", key=k1, params={**window, "api_key": k1_hash}
)
for path, param, value, detail in (
(
"/key/spend/report",
"api_key",
k2_hash,
"Not authorized to view spend for a api_key other than your own",
),
(
"/user/spend/report",
"internal_user_id",
user2,
"Not authorized to view spend for a internal_user_id other than your own",
),
(
"/team/spend/report",
"team_id",
team2,
"Not authorized to view spend for a team_id other than your own",
),
("/organization/spend/report", "organization_id", org2, "You do not have access to this organization"),
):
refused: Final = _report(gateway, path, key=k1, params={**window, param: value})
assert refused.status_code == 403, refused.text
assert refused.json() == {"detail": detail}, refused.text
_assert_report_matches_db(
_rows(gateway, "/key/spend/report", params={**window, "api_key": k2_hash}),
_logged("api_key", k2_hash, start, end),
)
_assert_report_matches_db(
_rows(gateway, "/user/spend/report", params={**window, "internal_user_id": user2}),
_logged('"user"', user2, start, end),
)
_assert_report_matches_db(
_rows(gateway, "/team/spend/report", params={**window, "team_id": team2}),
_logged("team_id", team2, start, end),
)
org_report: Final = _rows(gateway, "/organization/spend/report", params={**window, "organization_id": org2})
_assert_report_matches_db(org_report, _logged_keys((k2_hash, direct_org_key_hash), start, end))
assert {detail.team_id for entry in org_report for detail in entry.model_details} == {team2, ""}, org_report
@pytest.mark.parametrize(
"path",
(
"/key/spend/report",
"/user/spend/report",
"/team/spend/report",
"/organization/spend/report",
),
)
@pytest.mark.parametrize(
("date_kind", "params", "detail"),
(
("missing", None, "Please provide start_date and end_date"),
(
"malformed",
{"start_date": "2026-13-45", "end_date": "2026-01-01"},
"start_date and end_date must be in YYYY-MM-DD format",
),
),
)
def test_spend_report_rejects_missing_and_malformed_dates(
gateway: Gateway, path: str, date_kind: str, params: Mapping[str, str] | None, detail: str
) -> None:
response: Final = _report(gateway, path, params=params)
assert response.status_code == 400, response.text
assert response.json() == {"detail": detail}, response.text

View file

@ -3,13 +3,15 @@ import uuid
from collections.abc import AsyncIterator
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from hashlib import sha256
from typing import Final
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import psycopg
import pytest
import pytest_asyncio
from integration._support.client import Gateway
from integration._support.client import Gateway, eventually, object_value, string_value
from integration._support.database import read_rows
from prisma import Prisma
from psycopg import sql
from psycopg.types.json import Jsonb
@ -224,3 +226,91 @@ def test_spend_log_routes_preserve_user_and_permitted_team_access(
admin: Final = gateway.request("GET", f"/spend/logs/ui/{rows[2].request_id}")
assert admin.status_code == 200, admin.text
assert admin.json()["messages"] == [{"role": "user", "content": rows[2].request_id + " payload"}], admin.text
def _chat_request_id(gateway: Gateway, model: str, key: str) -> str:
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "max_tokens": 20, "messages": [{"role": "user", "content": f"scope {uuid.uuid4().hex}"}]},
key=key,
)
assert response.status_code == 200, response.text
body: Final = TypeAdapter(dict[str, object]).validate_python(response.json())
return str(body["id"])
def _legacy_spend_rows(gateway: Gateway, key: str, params: dict[str, str]) -> list[dict[str, object]]:
response: Final = gateway.request("GET", "/spend/logs", key=key, params=params)
assert response.status_code == 200, response.text
return TypeAdapter(list[dict[str, object]]).validate_python(response.json())
def test_legacy_spend_logs_clamp_internal_user_filters_to_their_own_user(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
user_a: Final = scenario.user(user_role="internal_user")
user_b: Final = scenario.user(user_role="internal_user")
key_a: Final = scenario.key(user_id=user_a, models=[model])
key_b: Final = scenario.key(user_id=user_b, models=[model])
request_a: Final = _chat_request_id(gateway, model, key_a)
request_b: Final = _chat_request_id(gateway, model, key_b)
for request_id in (request_a, request_b):
eventually(
lambda request_id=request_id: read_rows(
'SELECT request_id, "user" FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)
),
lambda rows: len(rows) == 1,
seconds=70,
)
filtered: Final = _legacy_spend_rows(gateway, key_a, {"user_id": user_b})
assert filtered != [], filtered
assert {string_value(object_value(row)["user"]) for row in filtered} == {user_a}, filtered
assert request_b not in {string_value(object_value(row)["request_id"]) for row in filtered}
assert request_a in {string_value(object_value(row)["request_id"]) for row in filtered}
by_key: Final = _legacy_spend_rows(gateway, key_a, {"api_key": key_b})
assert by_key == [], by_key
hashed_key_b: Final = sha256(key_b.encode()).hexdigest()
by_hashed_key: Final = _legacy_spend_rows(gateway, key_a, {"api_key": hashed_key_b})
assert by_hashed_key == [], by_hashed_key
by_request: Final = _legacy_spend_rows(gateway, key_a, {"request_id": request_b})
assert by_request == [], by_request
now: Final = datetime.now(timezone.utc)
window: Final = {
"start_date": (now - timedelta(days=1)).strftime("%Y-%m-%d"),
"end_date": (now + timedelta(days=1)).strftime("%Y-%m-%d"),
}
summarized: Final = _legacy_spend_rows(gateway, key_a, {**window, "summarize": "true"})
assert user_b not in {name for row in summarized for name in object_value(object_value(row)["users"])}, (
summarized
)
assert sum(float(str(object_value(row)["spend"])) for row in summarized) == pytest.approx(0.06), summarized
unsummarized: Final = _legacy_spend_rows(gateway, key_a, {**window, "summarize": "false"})
assert unsummarized != [], unsummarized
assert {string_value(object_value(row)["user"]) for row in unsummarized} == {user_a}, unsummarized
master: Final = _legacy_spend_rows(gateway, gateway.key, {"user_id": user_b, "request_id": request_b})
assert [string_value(object_value(row)["request_id"]) for row in master] == [request_b], master
for api_key in (key_b, hashed_key_b):
admin_by_key = _legacy_spend_rows(gateway, gateway.key, {"api_key": api_key})
assert [
(string_value(object_value(row)["request_id"]), object_value(row)["user"]) for row in admin_by_key
] == [(request_b, user_b)], admin_by_key
summarized_by_key: Final = _legacy_spend_rows(gateway, key_a, {**window, "summarize": "true", "api_key": key_b})
assert sum(float(str(object_value(row)["spend"])) for row in summarized_by_key) == 0, summarized_by_key
admin_summarized: Final = _legacy_spend_rows(
gateway, gateway.key, {**window, "summarize": "true", "api_key": key_b}
)
assert sum(float(str(object_value(row)["spend"])) for row in admin_summarized) == pytest.approx(0.06), (
admin_summarized
)
admin_users: Final = {
name: float(str(spend))
for row in admin_summarized
for name, spend in object_value(object_value(row)["users"]).items()
}
assert admin_users == pytest.approx({user_b: 0.06}), admin_summarized