Schedule budget resets at expectable times (#10331) (#10333)

* Schedule budget resets at expectable times (#10331)

* Enhance budget reset functionality with timezone support and standardized reset times

- Added `get_next_standardized_reset_time` function to calculate budget reset times based on specified durations and timezones.
- Introduced `timezone_utils.py` to manage timezone retrieval and budget reset time calculations.
- Updated budget reset logic in `reset_budget_job.py`, `internal_user_endpoints.py`, `key_management_endpoints.py`, and `team_endpoints.py` to utilize the new timezone-aware reset time calculations.
- Added unit tests for the new reset time functionality in `test_duration_parser.py`.
- Updated `.gitignore` to include `test.py` and made minor formatting adjustments in `docker-compose.yml` for consistency.

* Fixed linting

* Fix for mypy

* Fixed testcase for reset

* fix(duration_parser.py): move off zoneinfo - doesn't work with python 3.8

* test: update test

* refactor: improve budget reset time calculation and update related tests for accuracy

* clean up imports in team_endpoints.py

* test: update budget remaining hours assertions to reflect new reset time logic

* build(model_prices_and_context_window.json): update model

---------

Co-authored-by: Prathamesh Saraf <pratamesh1867@gmail.com>

tests/proxy_unit_tests/test_key_generate_prisma.py

@@ -1361,8 +1361,19 @@ def test_generate_and_update_key(prisma_client):
             )
             current_time = datetime.now(timezone.utc)
 
-            # assert budget_reset_at is 30 days from now
-            assert 31 >= (budget_reset_at - current_time).days >= 27
+            # Calculate days until end of current month
+            if current_time.month == 12:
+                end_of_month = datetime(current_time.year + 1, 1, 1, tzinfo=timezone.utc)
+            else:
+                end_of_month = datetime(current_time.year, current_time.month + 1, 1, tzinfo=timezone.utc)
+            days_until_end_of_month = (end_of_month - current_time).days - 1
+
+            # assert budget_reset_at is at the end of the current month
+            # assert days_until_end_of_month >= (budget_reset_at - current_time).days >= days_until_end_of_month - 2 # handle time zone differences
+            # Check that budget_reset_at is on the first day of next month
+            next_month_first_day = end_of_month 
+            # Assert that the reset date is the 1st of next month (0 or 1 day difference)
+            assert abs((budget_reset_at - next_month_first_day).days) <= 1
 
             # cleanup - delete key
             delete_key_request = KeyRequest(keys=[generated_key])
@@ -1924,7 +1935,7 @@ async def test_call_with_key_never_over_budget(prisma_client):
         pytest.fail(f"This should have not failed!. They key uses max_budget=None. {e}")
 
 
-@pytest.mark.asyncio()
+@pytest.mark.asyncio
 async def test_call_with_key_over_budget_stream(prisma_client):
     # 14. Make a call with a key over budget, expect to fail
     setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
@@ -2729,10 +2740,14 @@ async def test_create_update_team(prisma_client):
     budget_reset_at = _updated_info["budget_reset_at"].replace(tzinfo=timezone.utc)
     current_time = datetime.datetime.now(timezone.utc)
 
-    # assert budget_reset_at is 2 days from now
-    assert (
-        abs((budget_reset_at - current_time).total_seconds() - 2 * 24 * 60 * 60) <= 10
-    )
+    # Verify that budget_reset_at is at midnight (hour, minute, second are all 0)
+    assert budget_reset_at.hour == 0
+    assert budget_reset_at.minute == 0
+    assert budget_reset_at.second == 0
+
+    # Calculate days difference - should be close to 2 days (within 1 day to account for time of test execution)
+    days_diff = (budget_reset_at.date() - current_time.date()).days
+    assert 1 <= days_diff <= 2
 
     # now hit team_info
     try:
@@ -2859,12 +2874,18 @@ async def test_update_user_unit_test(prisma_client):
     assert _user_info["rpm_limit"] == 100
     assert _user_info["metadata"] == {"very-new-metadata": "something"}
 
-    # budget reset at should be 10 days from now
+    # budget_reset_at should be at midnight 10 days from now
     budget_reset_at = _user_info["budget_reset_at"].replace(tzinfo=timezone.utc)
     current_time = datetime.now(timezone.utc)
-    assert (
-        abs((budget_reset_at - current_time).total_seconds() - 10 * 24 * 60 * 60) <= 10
-    )
+    
+    # Verify that budget_reset_at is at midnight (hour, minute, second are all 0)
+    assert budget_reset_at.hour == 0
+    assert budget_reset_at.minute == 0
+    assert budget_reset_at.second == 0
+    
+    # Calculate days difference - should be close to 10 days (within 1 day to account for time of test execution)
+    days_diff = (budget_reset_at.date() - current_time.date()).days
+    assert 9 <= days_diff <= 10
 
 
 @pytest.mark.asyncio()