From 5f3541c68bde09b76359e66aa9e03e94939a3560 Mon Sep 17 00:00:00 2001
From: mubashir1osmani <mubashir.osmani777@gmail.com>
Date: Tue, 9 Jun 2026 13:55:46 -0700
Subject: [PATCH] fix greptile review

---
 .../proxy/auth/test_auth_checks.py            | 51 ++++++++++++++++---
 1 file changed, 44 insertions(+), 7 deletions(-)

diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py
index 6360d4a4e1b..cc05a1073ac 100644
--- a/tests/test_litellm/proxy/auth/test_auth_checks.py
+++ b/tests/test_litellm/proxy/auth/test_auth_checks.py
@@ -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