test(lint): clear the two PT011/PT012 violations left on the test tree (#37864)

This commit is contained in:
tin-berri 2026-08-21 11:15:06 -07:00 • committed by GitHub
parent 4307b34aca
commit ae1eea17bb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 13 additions and 13 deletions

View file

@ -1210,18 +1210,19 @@ async def test_a_pre_upgrade_counter_keyed_on_the_request_model_still_enforces(e
model_max_budget = {"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}}
await limiter.dual_cache.async_set_cache(key=f"{prefix}:entity-1:openai/gpt-4:1d", value=25.0, ttl=86400)
if entity_type == Litellm_EntityType.KEY:
budget_check = limiter.is_key_within_model_budget(
user_api_key_dict=UserAPIKeyAuth(token="entity-1", model_max_budget=model_max_budget),
model="openai/gpt-4",
)
else:
budget_check = limiter.is_end_user_within_model_budget(
end_user_id="entity-1",
end_user_model_max_budget=model_max_budget,
model="openai/gpt-4",
)
with pytest.raises(litellm.BudgetExceededError) as exc_info:
if entity_type == Litellm_EntityType.KEY:
await limiter.is_key_within_model_budget(
user_api_key_dict=UserAPIKeyAuth(token="entity-1", model_max_budget=model_max_budget),
model="openai/gpt-4",
)
else:
await limiter.is_end_user_within_model_budget(
end_user_id="entity-1",
end_user_model_max_budget=model_max_budget,
model="openai/gpt-4",
)
await budget_check
assert exc_info.value.current_cost == 25.0

View file

@ -1357,9 +1357,8 @@ async def test_user_model_budget_is_enforced_through_user_api_key_auth(over_budg
new=fake_get_user_object,
):
if expect_refusal:
with pytest.raises(Exception) as exc:
with pytest.raises(Exception, match=r"(?i)budget") as exc:
await user_api_key_auth(request=request, api_key="Bearer " + key)
assert "budget" in str(exc.value).lower()
assert user_id in str(exc.value)
else:
result = await user_api_key_auth(request=request, api_key="Bearer " + key)