fix: pydantic validation for budget info, dont allow explicit type cast

This commit is contained in:
mubashir1osmani 2026-06-18 23:15:38 -07:00
parent a143d74e76
commit 6843ca0325
No known key found for this signature in database
GPG key ID: AB055FF67D0B4D9A

View file

@ -29,10 +29,10 @@ def test_budget_crud_roundtrip(client: BudgetClient, resources: ResourceManager)
# Attach the budget to a key and confirm the key reflects it.
key = client.generate_key(extra_params={"budget_id": budget_id})
resources.defer(lambda: client.delete_key(key))
info = client.key_info(key)
linked = info.get("litellm_budget_table") or {}
assert info.get("budget_id") == budget_id or linked.get("max_budget") == 12.5, (
f"key does not reflect attached budget: {info.get('budget_id')}, {linked}"
info = client.gateway.key_info(key)
linked = info.litellm_budget_table
assert info.budget_id == budget_id or (linked is not None and linked.max_budget == 12.5), (
f"key does not reflect attached budget: {info.budget_id}, {linked}"
)
@ -46,7 +46,7 @@ def test_budget_duration_schedules_reset_on_key(client: BudgetClient, resources:
key = client.generate_key(max_budget=10.0, extra_params={"budget_duration": "30d"})
resources.defer(lambda: client.delete_key(key))
reset_at = client.key_info(key).get("budget_reset_at")
reset_at = client.gateway.key_info(key).budget_reset_at
assert reset_at, "budget_duration did not set budget_reset_at on the key"
# budget_duration schedules a FUTURE reset. Don't assume now+30d exactly: the