address greptile review feedback

litellm/proxy/management_endpoints/common_utils.py

@@ -1,8 +1,6 @@
-from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 from litellm._logging import verbose_proxy_logger
-from litellm.litellm_core_utils.duration_parser import duration_in_seconds
 from litellm.caching import DualCache
 from litellm.proxy._types import (
     KeyRequestBase,
@@ -399,9 +397,11 @@ async def _upsert_budget_and_membership(
     if rpm_limit is not None:
         create_data["rpm_limit"] = rpm_limit
     if budget_duration is not None:
+        from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
+
         create_data["budget_duration"] = budget_duration
-        create_data["budget_reset_at"] = datetime.utcnow() + timedelta(
-            seconds=duration_in_seconds(duration=budget_duration)
+        create_data["budget_reset_at"] = get_budget_reset_time(
+            budget_duration=budget_duration
         )
 
     new_budget = await tx.litellm_budgettable.create(

tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py

@@ -287,13 +287,12 @@ async def test_upsert_with_budget_duration(mock_tx, fake_user):
     from unittest.mock import patch
     from datetime import datetime as dt
 
-    fake_now = dt(2026, 1, 1, 0, 0, 0)
+    fake_reset_at = dt(2026, 1, 31, 0, 0, 0)
 
     with patch(
-        "litellm.proxy.management_endpoints.common_utils.datetime"
-    ) as mock_dt:
-        mock_dt.utcnow.return_value = fake_now
-
+        "litellm.proxy.common_utils.timezone_utils.get_budget_reset_time",
+        return_value=fake_reset_at,
+    ):
         await _upsert_budget_and_membership(
             mock_tx,
             team_id="team-dur",
@@ -306,10 +305,7 @@ async def test_upsert_with_budget_duration(mock_tx, fake_user):
 
     call_data = mock_tx.litellm_budgettable.create.call_args.kwargs["data"]
     assert call_data["budget_duration"] == "30d"
-    assert "budget_reset_at" in call_data
-    # 30 days from fake_now
-    from datetime import timedelta
-    assert call_data["budget_reset_at"] == fake_now + timedelta(days=30)
+    assert call_data["budget_reset_at"] == fake_reset_at
 
     # membership upsert should still happen
     mock_tx.litellm_teammembership.upsert.assert_awaited_once()
@@ -325,13 +321,12 @@ async def test_upsert_budget_duration_only_creates_budget(mock_tx, fake_user):
     from unittest.mock import patch
     from datetime import datetime as dt
 
-    fake_now = dt(2026, 1, 1, 0, 0, 0)
+    fake_reset_at = dt(2026, 1, 8, 0, 0, 0)
 
     with patch(
-        "litellm.proxy.management_endpoints.common_utils.datetime"
-    ) as mock_dt:
-        mock_dt.utcnow.return_value = fake_now
-
+        "litellm.proxy.common_utils.timezone_utils.get_budget_reset_time",
+        return_value=fake_reset_at,
+    ):
         await _upsert_budget_and_membership(
             mock_tx,
             team_id="team-dur-only",
@@ -348,7 +343,7 @@ async def test_upsert_budget_duration_only_creates_budget(mock_tx, fake_user):
     # Should create a budget with budget_duration set
     call_data = mock_tx.litellm_budgettable.create.call_args.kwargs["data"]
     assert call_data["budget_duration"] == "7d"
-    assert "budget_reset_at" in call_data
+    assert call_data["budget_reset_at"] == fake_reset_at
 
     # Should upsert membership
     mock_tx.litellm_teammembership.upsert.assert_awaited_once()

tests/test_litellm/proxy/prompts/test_prompt_endpoints.py

@@ -247,8 +247,12 @@ class TestPromptVersionsEndpoint:
             ),
         }
 
-        # Mock the IN_MEMORY_PROMPT_REGISTRY at the import location
-        with patch("litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY") as mock_registry:
+        # Mock the IN_MEMORY_PROMPT_REGISTRY at the import location.
+        # Also patch prisma_client to None so the function takes the in-memory
+        # path regardless of any test-suite-level state (avoids Python 3.12
+        # "MagicMock can't be used in await expression" errors).
+        with patch("litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY") as mock_registry, \
+             patch("litellm.proxy.proxy_server.prisma_client", None):
             mock_registry.IN_MEMORY_PROMPTS = mock_prompts
 
             # Test with base prompt ID
@@ -293,7 +297,11 @@ class TestPromptVersionsEndpoint:
             user_role=LitellmUserRoles.PROXY_ADMIN
         )
 
-        with patch("litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY") as mock_registry:
+        # Also patch prisma_client to None to force the in-memory path (avoids
+        # Python 3.12 "MagicMock can't be used in await expression" errors when
+        # another test in the same xdist worker sets prisma_client to a MagicMock).
+        with patch("litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY") as mock_registry, \
+             patch("litellm.proxy.proxy_server.prisma_client", None):
             mock_registry.IN_MEMORY_PROMPTS = {}
 
             with pytest.raises(HTTPException) as exc_info: