From 09d73480403be6d6da78ebf7f2e72fe4f93943b1 Mon Sep 17 00:00:00 2001 From: daanhendrio <255322319+daanhendrio@users.noreply.github.com> Date: Sat, 11 Apr 2026 15:05:56 +0000 Subject: [PATCH] address greptile review feedback --- .../management_endpoints/common_utils.py | 8 +++--- .../test_upsert_budget_membership.py | 25 ++++++++----------- .../proxy/prompts/test_prompt_endpoints.py | 14 ++++++++--- 3 files changed, 25 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 37e76b85a2c..bf80984efef 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/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( diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index facbf3efb0f..80d90492e5e 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/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() diff --git a/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py b/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py index 6c2e5fa7667..21ec98b88cb 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py +++ b/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: