Incorporate Greptile review- improve tests

tests/test_litellm/proxy/management_endpoints/test_initial_budget_reset_at.py

@@ -14,11 +14,11 @@ sys.path.insert(
     0, os.path.abspath("../../../..")
 )
 
-from litellm.proxy._types import NewUserRequest
-from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
+from litellm.proxy._types import NewUserRequest, UpdateUserRequest, UpdateUserRequestNoUserIDorEmail
+from litellm.proxy.management_endpoints.internal_user_endpoints import new_user, user_update
 
 
-class TestBudgetResetAtValidation:
+class TestInitialBudgetResetAtValidation:
     """Test Pydantic validation for initial_budget_reset_at field"""
 
     def test_past_date_rejection_user_request(self):
@@ -46,12 +46,49 @@ class TestBudgetResetAtValidation:
 
         assert user_request.initial_budget_reset_at == future_date
 
+    def test_past_date_rejection_update_user_request(self):
+        """Test that UpdateUserRequest rejects past dates for initial_budget_reset_at"""
+        past_date = datetime.now(timezone.utc) - timedelta(days=1)
 
-class TestUserCreationWithBudgetResetAt:
+        with pytest.raises(ValueError) as exc_info:
+            UpdateUserRequest(
+                user_id="test_user",
+                budget_duration="10m",
+                initial_budget_reset_at=past_date,
+            )
+
+        assert "initial_budget_reset_at cannot be in the past" in str(exc_info.value)
+
+    def test_past_date_rejection_update_user_request_no_id_or_email(self):
+        """Test that UpdateUserRequestNoUserIDorEmail rejects past dates for initial_budget_reset_at"""
+        past_date = datetime.now(timezone.utc) - timedelta(days=1)
+
+        with pytest.raises(ValueError) as exc_info:
+            UpdateUserRequestNoUserIDorEmail(
+                budget_duration="10m",
+                initial_budget_reset_at=past_date,
+            )
+
+        assert "initial_budget_reset_at cannot be in the past" in str(exc_info.value)
+
+    def test_future_date_accepted_update_user_request(self):
+        """Test that future dates are accepted on UpdateUserRequest"""
+        future_date = datetime.now(timezone.utc) + timedelta(days=10)
+
+        update_request = UpdateUserRequest(
+            user_id="test_user",
+            budget_duration="10m",
+            initial_budget_reset_at=future_date,
+        )
+
+        assert update_request.initial_budget_reset_at == future_date
+
+
+class TestUserCreationWithInitialBudgetResetAt:
     """Integration test for creating users with initial_budget_reset_at"""
 
     @pytest.mark.asyncio
-    async def test_create_user_explicit_budget_reset_at_takes_precedence(self):
+    async def test_create_user_explicit_initial_budget_reset_at_takes_precedence(self):
         """Test that explicit initial_budget_reset_at is honored and takes precedence over duration-based computation"""
         mock_prisma_client = MagicMock()
         mock_prisma_client.insert_data = AsyncMock(
@@ -99,5 +136,53 @@ class TestUserCreationWithBudgetResetAt:
                 raise AssertionError("User insert call not found")
 
 
+class TestUserUpdateWithInitialBudgetResetAt:
+    """Integration test for updating users with initial_budget_reset_at"""
+
+    @pytest.mark.asyncio
+    async def test_update_user_explicit_initial_budget_reset_at_translated(self):
+        """Test that initial_budget_reset_at is translated to budget_reset_at on the update path"""
+        mock_prisma_client = MagicMock()
+        mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
+        mock_prisma_client.update_data = AsyncMock(
+            return_value={"user_id": "existing_user"}
+        )
+
+        explicit_date = datetime.now(timezone.utc) + timedelta(days=30)
+
+        update_request = UpdateUserRequest(
+            user_id="existing_user",
+            budget_duration="10m",
+            initial_budget_reset_at=explicit_date,
+            max_budget=100.0,
+        )
+
+        with patch(
+            "litellm.proxy.proxy_server.prisma_client",
+            mock_prisma_client,
+        ), patch(
+            "litellm.proxy.proxy_server.litellm_proxy_admin_name",
+            "admin",
+        ):
+            result = await user_update(
+                data=update_request,
+                user_api_key_dict=MagicMock(
+                    user_id="existing_user",
+                    user_role="proxy_admin",
+                ),
+            )
+
+            assert result is not None
+            assert mock_prisma_client.update_data.called
+
+            call_kwargs = mock_prisma_client.update_data.call_args.kwargs
+            user_data = call_kwargs["data"]
+            assert "budget_reset_at" in user_data, "budget_reset_at should be set from initial_budget_reset_at"
+            assert "initial_budget_reset_at" not in user_data, "initial_budget_reset_at should be removed"
+            stored_date = user_data["budget_reset_at"]
+            assert isinstance(stored_date, datetime)
+            assert abs((stored_date - explicit_date).total_seconds()) < 1
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])