mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
test(lint): clear the two PT011/PT012 violations left on the test tree (#37864)
This commit is contained in:
parent
4307b34aca
commit
ae1eea17bb
2 changed files with 13 additions and 13 deletions
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue