scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-05 02:41:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 48

fix greptile review

Browse source
This commit is contained in:
mubashir1osmani 2026-06-09 13:55:46 -07:00
parent 5f954f1f4a
commit 5f3541c68b
No known key found for this signature in database
GPG key ID: AB055FF67D0B4D9A
1 changed files with 44 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -1,5 +1,6 @@
import asyncio
import json
import math
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
@ -2411,8 +2412,8 @@ async def test_virtual_key_budget_check_fallback_no_counter():
@pytest.mark.parametrize(
"non_finite_max_budget",
[float("nan"), float("inf")],
ids=["nan", "positive_infinity"],
[float("nan"), float("inf"), float("-inf")],
ids=["nan", "positive_infinity", "negative_infinity"],
)
@pytest.mark.asyncio
async def test_virtual_key_max_budget_check_non_finite_max_budget_does_not_bypass(
@ -2434,11 +2435,47 @@ async def test_virtual_key_max_budget_check_non_finite_max_budget_does_not_bypas
async def mock_get_current_spend(counter_key, fallback_spend):
return 999.0
with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend):
await _virtual_key_max_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
with patch(
"litellm.proxy.auth.auth_checks.math.isfinite", wraps=math.isfinite
) as mock_isfinite:
with patch(
"litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend
):
await _virtual_key_max_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
mock_isfinite.assert_called_once_with(non_finite_max_budget)
@pytest.mark.asyncio
async def test_virtual_key_max_budget_check_negative_infinity_blocks_without_isfinite_guard():
"""Removing math.isfinite would wrongly block keys with max_budget=-inf."""
from litellm.proxy.utils import ProxyLogging
valid_token = UserAPIKeyAuth(
token="test-hashed-token",
spend=0.0,
max_budget=float("-inf"),
user_id="test-user",
)
proxy_logging_obj = ProxyLogging(user_api_key_cache=None)
proxy_logging_obj.budget_alerts = AsyncMock()
async def mock_get_current_spend(counter_key, fallback_spend):
return 999.0
with patch("litellm.proxy.auth.auth_checks.math.isfinite", return_value=True):
with patch(
"litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend
):
with pytest.raises(litellm.BudgetExceededError):
await _virtual_key_max_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
@pytest.mark.asyncio