diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index 7ac2299b7d8..fe8b71e7623 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -166,6 +166,37 @@ def get_next_standardized_reset_time( return base_midnight + timedelta(days=1) +def _subtract_months(moment: datetime, months: int) -> datetime: + total_months: Final = moment.year * 12 + moment.month - 1 - months + year, month_index = divmod(total_months, 12) + month: Final = month_index + 1 + return moment.replace(year=year, month=month, day=min(moment.day, get_last_day_of_month(year, month))) + + +def get_budget_window_start(duration: str, reset_at: datetime) -> datetime: + """Start of the budget period that ends at `reset_at`, under the same rules + `get_next_standardized_reset_time` used to pick it: `30d` and `Nmo` span calendar + months and a duration it does not recognize resets at the next midnight.""" + value, unit = _parse_duration(_normalize_duration(duration)) + if value is None: + return reset_at - timedelta(days=1) + match unit: + case "mo": + return _subtract_months(reset_at, value) + case "d": + return _subtract_months(reset_at, 1) if value == 30 else reset_at - timedelta(days=value) + case "w": + return reset_at - timedelta(weeks=value) + case "h": + return reset_at - timedelta(hours=value) + case "m": + return reset_at - timedelta(minutes=value) + case "s": + return reset_at - timedelta(seconds=value) + case _: + return reset_at - timedelta(days=1) + + def _setup_timezone(current_time: datetime, timezone_str: str = "UTC") -> tuple[datetime, tzinfo]: """Set up timezone and normalize current time to that timezone.""" try: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 000ffa9c305..90e760093b4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -887,7 +887,7 @@ from litellm.proxy.utils import ( # noqa: F401, RUF100 # legacy module exports hash_password, hash_token, invalidate_config_param, - is_projected_spend_over_limit, + is_projected_spend_over_limit, # pyright: ignore[reportUnusedImport] # backwards-compatible package export is_valid_team_configs, litellm_config_cache, migrate_passwords_to_scrypt_async, @@ -4165,21 +4165,18 @@ async def update_cache( new_spend: Final = existing_spend + response_cost ## CHECK IF USER PROJECTED SPEND > SOFT LIMIT - if ( - existing_spend_obj.soft_budget_cooldown is False - and existing_spend_obj.soft_budget is not None - and ( - is_projected_spend_over_limit( - current_spend=new_spend, - soft_budget_limit=existing_spend_obj.soft_budget, - ) - is True - ) - ): - projected_spend, projected_exceeded_date = get_projected_spend_over_limit( + projection: Final = ( + get_projected_spend_over_limit( current_spend=new_spend, soft_budget_limit=existing_spend_obj.soft_budget, + budget_duration=existing_spend_obj.budget_duration, + budget_reset_at=existing_spend_obj.budget_reset_at, ) + if existing_spend_obj.soft_budget_cooldown is False and existing_spend_obj.soft_budget is not None + else None + ) + if projection is not None: + projected_spend, projected_exceeded_date = projection soft_limit: Final = existing_spend_obj.soft_budget call_info: Final = CallInfo( token=existing_spend_obj.token or "", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 3e73e2dee40..8f796f065a6 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -147,6 +147,7 @@ from litellm.litellm_core_utils.core_helpers import ( independent_snapshot, is_expected_client_error, ) +from litellm.litellm_core_utils.duration_parser import get_budget_window_start from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -8101,66 +8102,94 @@ def _get_month_end_date(today: date) -> date: return date(today.year, today.month + 1, 1) - timedelta(days=1) -def is_projected_spend_over_limit(current_spend: float, soft_budget_limit: float | None) -> bool: - if soft_budget_limit is None: - # If there's no limit, we can't exceed it. - return False +MIN_ELAPSED_WINDOW_FRACTION: Final = 1 / 24 + +def _as_aware(moment: datetime) -> datetime: + return moment if moment.tzinfo is not None else moment.replace(tzinfo=timezone.utc) + + +def _reset_window(budget_duration: str, budget_reset_at: datetime) -> tuple[datetime, datetime] | None: + window_end: Final = _as_aware(budget_reset_at) + window_start: Final = get_budget_window_start(budget_duration, window_end) + return (window_start, window_end) if window_start < window_end else None + + +def _project_within_window( + current_spend: float, + soft_budget_limit: float, + window: tuple[datetime, datetime], + now: datetime | None, +) -> tuple[float, date] | None: + window_start, window_end = window + moment: Final = (_as_aware(now) if now is not None else datetime.now(timezone.utc)).astimezone(window_end.tzinfo) + if moment < window_start: + return None + elapsed: Final = max(moment - window_start, (window_end - window_start) * MIN_ELAPSED_WINDOW_FRACTION) + remaining: Final = max(window_end - moment, timedelta(0)) + spend_per_second: Final = current_spend / elapsed.total_seconds() + projected_spend: Final = current_spend + spend_per_second * remaining.total_seconds() + if projected_spend <= soft_budget_limit: + return None + remaining_budget: Final = soft_budget_limit - current_spend + if spend_per_second <= 0 or remaining_budget <= 0: + return projected_spend, moment.date() + exceed_at: Final = min(moment + timedelta(seconds=remaining_budget / spend_per_second), window_end) + return projected_spend, exceed_at.date() + + +def _project_to_month_end(current_spend: float, soft_budget_limit: float) -> tuple[float, date] | None: today: Final = date.today() + remaining_days: Final = (_get_month_end_date(today) - today).days + daily_spend: Final = current_spend / max(today.day - 1, 1) + projected_spend: Final = current_spend + daily_spend * remaining_days + if projected_spend <= soft_budget_limit: + return None + remaining_budget: Final = soft_budget_limit - current_spend + if daily_spend <= 0 or remaining_budget <= 0: + return projected_spend, today + return projected_spend, today + timedelta(days=remaining_budget / daily_spend) - # Finding the first day of the next month, then subtracting one day to get the end of the current month. - end_month: Final = _get_month_end_date(today) - remaining_days: Final = (end_month - today).days - - # Check for the start of the month to avoid division by zero - if today.day == 1: - daily_spend_estimate = current_spend - else: - daily_spend_estimate = current_spend / (today.day - 1) - - # Total projected spend for the month - projected_spend: Final = current_spend + (daily_spend_estimate * remaining_days) - - if projected_spend > soft_budget_limit: - print_verbose("Projected spend exceeds soft budget limit!") - return True - return False +def is_projected_spend_over_limit( + current_spend: float, + soft_budget_limit: float | None, + budget_duration: str | None = None, + budget_reset_at: datetime | None = None, + now: datetime | None = None, +) -> bool: + return ( + get_projected_spend_over_limit( + current_spend=current_spend, + soft_budget_limit=soft_budget_limit, + budget_duration=budget_duration, + budget_reset_at=budget_reset_at, + now=now, + ) + is not None + ) _is_projected_spend_over_limit: Final = is_projected_spend_over_limit -def get_projected_spend_over_limit(current_spend: float, soft_budget_limit: float | None) -> tuple | None: +def get_projected_spend_over_limit( + current_spend: float, + soft_budget_limit: float | None, + budget_duration: str | None = None, + budget_reset_at: datetime | None = None, + now: datetime | None = None, +) -> tuple[float, date] | None: if soft_budget_limit is None: return None - - today: Final = date.today() - end_month: Final = _get_month_end_date(today) - remaining_days: Final = (end_month - today).days - - # assuming the current spend till today (not including today) - if today.day == 1: - daily_spend = current_spend - else: - daily_spend = current_spend / (today.day - 1) - projected_spend: Final = current_spend + (daily_spend * remaining_days) - - if projected_spend > soft_budget_limit: - if daily_spend <= 0: - limit_exceed_date = today - else: - remaining_budget: Final = soft_budget_limit - current_spend - if remaining_budget <= 0: - limit_exceed_date = today - else: - approx_days: Final = remaining_budget / daily_spend - limit_exceed_date = today + timedelta(days=approx_days) - - # return the projected spend and the date it will exceeded - return projected_spend, limit_exceed_date - - return None + window: Final = ( + _reset_window(budget_duration, budget_reset_at) + if budget_duration is not None and budget_reset_at is not None + else None + ) + if window is None: + return _project_to_month_end(current_spend, soft_budget_limit) + return _project_within_window(current_spend, soft_budget_limit, window, now) _get_projected_spend_over_limit: Final = get_projected_spend_over_limit diff --git a/tests/unit/litellm_core_utils/test_duration_parser.py b/tests/unit/litellm_core_utils/test_duration_parser.py index cb9f273a0a7..f5e50d67e71 100644 --- a/tests/unit/litellm_core_utils/test_duration_parser.py +++ b/tests/unit/litellm_core_utils/test_duration_parser.py @@ -1,11 +1,13 @@ import unittest from datetime import datetime, time, timezone +from typing import Final from unittest.mock import patch from zoneinfo import ZoneInfo import litellm.litellm_core_utils.duration_parser as duration_parser from litellm.litellm_core_utils.duration_parser import ( duration_in_seconds, + get_budget_window_start, get_next_standardized_reset_time, ) @@ -23,9 +25,7 @@ class TestStandardizedResetTime(unittest.TestCase): # Weekly reset (7d) - should reset on next Monday wednesday = datetime(2023, 5, 17, 15, 45, 0, tzinfo=timezone.utc) # A Wednesday - weekly_expected = datetime( - 2023, 5, 22, 0, 0, 0, tzinfo=timezone.utc - ) # Next Monday + weekly_expected = datetime(2023, 5, 22, 0, 0, 0, tzinfo=timezone.utc) # Next Monday weekly_result = get_next_standardized_reset_time("7d", wednesday, "UTC") self.assertEqual(weekly_result, weekly_expected) @@ -105,17 +105,13 @@ class TestStandardizedResetTime(unittest.TestCase): # Europe/London (UTC+1): 11:30 PM, so next 15m reset is 11:45 PM london = ZoneInfo("Europe/London") london_expected = datetime(2023, 5, 15, 23, 45, 0, tzinfo=london) - london_result = get_next_standardized_reset_time( - "15m", base_time, "Europe/London" - ) + london_result = get_next_standardized_reset_time("15m", base_time, "Europe/London") self.assertEqual(london_result, london_expected) # Test Bangkok timezone (UTC+7): 5:30 AM next day, so next reset is midnight the day after bangkok = ZoneInfo("Asia/Bangkok") bangkok_expected = datetime(2023, 5, 17, 0, 0, 0, tzinfo=bangkok) - bangkok_result = get_next_standardized_reset_time( - "1d", base_time, "Asia/Bangkok" - ) + bangkok_result = get_next_standardized_reset_time("1d", base_time, "Asia/Bangkok") self.assertEqual(bangkok_result, bangkok_expected) def test_edge_cases(self): @@ -137,16 +133,12 @@ class TestStandardizedResetTime(unittest.TestCase): # 30m near midnight - should roll over to next day midnight_minute_expected = datetime(2023, 5, 16, 0, 0, 0, tzinfo=timezone.utc) - midnight_minute_result = get_next_standardized_reset_time( - "30m", near_midnight, "UTC" - ) + midnight_minute_result = get_next_standardized_reset_time("30m", near_midnight, "UTC") self.assertEqual(midnight_minute_result, midnight_minute_expected) # Invalid timezone - should fall back to UTC invalid_tz_expected = datetime(2023, 5, 16, 0, 0, 0, tzinfo=timezone.utc) - invalid_tz_result = get_next_standardized_reset_time( - "1d", on_hour, "NonExistentTimeZone" - ) + invalid_tz_result = get_next_standardized_reset_time("1d", on_hour, "NonExistentTimeZone") self.assertEqual(invalid_tz_result, invalid_tz_expected) def test_iana_timezones_previously_unsupported(self): @@ -164,17 +156,13 @@ class TestStandardizedResetTime(unittest.TestCase): sydney = ZoneInfo("Australia/Sydney") # At 15:00 UTC it's 01:00 AEST May 16 → next midnight is May 17 00:00 AEST sydney_expected = datetime(2023, 5, 17, 0, 0, 0, tzinfo=sydney) - sydney_result = get_next_standardized_reset_time( - "1d", base_time, "Australia/Sydney" - ) + sydney_result = get_next_standardized_reset_time("1d", base_time, "Australia/Sydney") self.assertEqual(sydney_result, sydney_expected) # America/Chicago (UTC-5): at 15:00 UTC it's 10:00 CDT → next midnight is May 16 00:00 CDT chicago = ZoneInfo("America/Chicago") chicago_expected = datetime(2023, 5, 16, 0, 0, 0, tzinfo=chicago) - chicago_result = get_next_standardized_reset_time( - "1d", base_time, "America/Chicago" - ) + chicago_result = get_next_standardized_reset_time("1d", base_time, "America/Chicago") self.assertEqual(chicago_result, chicago_expected) def test_dst_fall_back(self): @@ -209,107 +197,77 @@ class TestResetTimeOfDay(unittest.TestCase): def test_daily_reset_before_offset_is_today(self): now = datetime(2023, 5, 15, 8, 0, 0, tzinfo=timezone.utc) - result = get_next_standardized_reset_time( - "1d", now, "UTC", reset_time_of_day=time(12, 0) - ) + result = get_next_standardized_reset_time("1d", now, "UTC", reset_time_of_day=time(12, 0)) self.assertEqual(result, datetime(2023, 5, 15, 12, 0, 0, tzinfo=timezone.utc)) def test_daily_reset_after_offset_is_tomorrow(self): now = datetime(2023, 5, 15, 14, 0, 0, tzinfo=timezone.utc) - result = get_next_standardized_reset_time( - "1d", now, "UTC", reset_time_of_day=time(12, 0) - ) + result = get_next_standardized_reset_time("1d", now, "UTC", reset_time_of_day=time(12, 0)) self.assertEqual(result, datetime(2023, 5, 16, 12, 0, 0, tzinfo=timezone.utc)) def test_daily_reset_exactly_at_offset_rolls_forward(self): now = datetime(2023, 5, 15, 12, 0, 0, tzinfo=timezone.utc) - result = get_next_standardized_reset_time( - "1d", now, "UTC", reset_time_of_day=time(12, 0) - ) + result = get_next_standardized_reset_time("1d", now, "UTC", reset_time_of_day=time(12, 0)) self.assertEqual(result, datetime(2023, 5, 16, 12, 0, 0, tzinfo=timezone.utc)) def test_daily_reset_with_seconds_offset(self): now = datetime(2023, 5, 15, 8, 0, 0, tzinfo=timezone.utc) - result = get_next_standardized_reset_time( - "1d", now, "UTC", reset_time_of_day=time(9, 30, 15) - ) + result = get_next_standardized_reset_time("1d", now, "UTC", reset_time_of_day=time(9, 30, 15)) self.assertEqual(result, datetime(2023, 5, 15, 9, 30, 15, tzinfo=timezone.utc)) def test_offset_applies_in_configured_timezone(self): # 2023-05-15 22:30 UTC == 2023-05-16 01:30 in Jerusalem (IDT, UTC+3), # so the next noon-Jerusalem reset is 2023-05-16 12:00 IDT. now = datetime(2023, 5, 15, 22, 30, 0, tzinfo=timezone.utc) - result = get_next_standardized_reset_time( - "1d", now, "Asia/Jerusalem", reset_time_of_day=time(12, 0) - ) + result = get_next_standardized_reset_time("1d", now, "Asia/Jerusalem", reset_time_of_day=time(12, 0)) jerusalem = result.astimezone(ZoneInfo("Asia/Jerusalem")) - self.assertEqual( - (jerusalem.year, jerusalem.month, jerusalem.day), (2023, 5, 16) - ) + self.assertEqual((jerusalem.year, jerusalem.month, jerusalem.day), (2023, 5, 16)) self.assertEqual(jerusalem.hour, 12) self.assertEqual(jerusalem.minute, 0) def test_weekly_reset_lands_on_monday_at_offset(self): wednesday = datetime(2023, 5, 17, 15, 45, 0, tzinfo=timezone.utc) - result = get_next_standardized_reset_time( - "7d", wednesday, "UTC", reset_time_of_day=time(12, 0) - ) + result = get_next_standardized_reset_time("7d", wednesday, "UTC", reset_time_of_day=time(12, 0)) self.assertEqual(result, datetime(2023, 5, 22, 12, 0, 0, tzinfo=timezone.utc)) def test_weekly_reset_today_is_monday_before_offset_is_today(self): monday_morning = datetime(2023, 5, 22, 9, 0, 0, tzinfo=timezone.utc) - result = get_next_standardized_reset_time( - "7d", monday_morning, "UTC", reset_time_of_day=time(12, 0) - ) + result = get_next_standardized_reset_time("7d", monday_morning, "UTC", reset_time_of_day=time(12, 0)) self.assertEqual(result, datetime(2023, 5, 22, 12, 0, 0, tzinfo=timezone.utc)) def test_weekly_reset_today_is_monday_after_offset_is_next_week(self): monday_afternoon = datetime(2023, 5, 22, 15, 0, 0, tzinfo=timezone.utc) - result = get_next_standardized_reset_time( - "7d", monday_afternoon, "UTC", reset_time_of_day=time(12, 0) - ) + result = get_next_standardized_reset_time("7d", monday_afternoon, "UTC", reset_time_of_day=time(12, 0)) self.assertEqual(result, datetime(2023, 5, 29, 12, 0, 0, tzinfo=timezone.utc)) def test_monthly_30d_lands_on_first_at_offset(self): now = datetime(2023, 5, 15, 10, 30, 0, tzinfo=timezone.utc) - result = get_next_standardized_reset_time( - "30d", now, "UTC", reset_time_of_day=time(12, 0) - ) + result = get_next_standardized_reset_time("30d", now, "UTC", reset_time_of_day=time(12, 0)) self.assertEqual(result, datetime(2023, 6, 1, 12, 0, 0, tzinfo=timezone.utc)) def test_monthly_1mo_today_is_first_before_offset_is_today(self): now = datetime(2023, 5, 1, 9, 0, 0, tzinfo=timezone.utc) - result = get_next_standardized_reset_time( - "1mo", now, "UTC", reset_time_of_day=time(12, 0) - ) + result = get_next_standardized_reset_time("1mo", now, "UTC", reset_time_of_day=time(12, 0)) self.assertEqual(result, datetime(2023, 5, 1, 12, 0, 0, tzinfo=timezone.utc)) def test_monthly_year_rollover_at_offset(self): now = datetime(2023, 12, 15, 9, 0, 0, tzinfo=timezone.utc) - result = get_next_standardized_reset_time( - "1mo", now, "UTC", reset_time_of_day=time(12, 0) - ) + result = get_next_standardized_reset_time("1mo", now, "UTC", reset_time_of_day=time(12, 0)) self.assertEqual(result, datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)) def test_custom_day_reset_applies_offset(self): now = datetime(2023, 5, 15, 10, 30, 0, tzinfo=timezone.utc) - result = get_next_standardized_reset_time( - "3d", now, "UTC", reset_time_of_day=time(12, 0) - ) + result = get_next_standardized_reset_time("3d", now, "UTC", reset_time_of_day=time(12, 0)) self.assertEqual(result, datetime(2023, 5, 18, 12, 0, 0, tzinfo=timezone.utc)) def test_sub_day_durations_ignore_offset(self): base = datetime(2023, 5, 15, 15, 20, 30, tzinfo=timezone.utc) self.assertEqual( - get_next_standardized_reset_time( - "2h", base, "UTC", reset_time_of_day=time(12, 0) - ), + get_next_standardized_reset_time("2h", base, "UTC", reset_time_of_day=time(12, 0)), datetime(2023, 5, 15, 16, 0, 0, tzinfo=timezone.utc), ) self.assertEqual( - get_next_standardized_reset_time( - "30m", base, "UTC", reset_time_of_day=time(12, 0) - ), + get_next_standardized_reset_time("30m", base, "UTC", reset_time_of_day=time(12, 0)), datetime(2023, 5, 15, 15, 30, 0, tzinfo=timezone.utc), ) @@ -385,5 +343,35 @@ class TestWordFormBudgetDurations(unittest.TestCase): self.assertIn("garbage", mock_warning.call_args.args) +class TestGetBudgetWindowStart(unittest.TestCase): + def test_window_start_is_the_previous_reset_on_the_same_schedule(self): + created_at: Final = datetime(2024, 10, 1, 0, 30, tzinfo=timezone.utc) + durations: Final = "1d 24h daily 7d weekly 2w 10d 30d monthly 1mo 4h 5h 30m 45s 1hr fortnightly".split() + for duration in durations: + with self.subTest(duration=duration): + reset_at = get_next_standardized_reset_time(duration, created_at, "UTC") + window_start = get_budget_window_start(duration, reset_at) + self.assertLessEqual(window_start, created_at) + self.assertEqual(get_next_standardized_reset_time(duration, window_start, "UTC"), reset_at) + + def test_thirty_days_spans_the_calendar_month_it_resets_on(self): + self.assertEqual( + get_budget_window_start("30d", datetime(2024, 3, 1, tzinfo=timezone.utc)), + datetime(2024, 2, 1, tzinfo=timezone.utc), + ) + + def test_months_clamp_to_the_shorter_month(self): + self.assertEqual( + get_budget_window_start("1mo", datetime(2024, 3, 31, tzinfo=timezone.utc)), + datetime(2024, 2, 29, tzinfo=timezone.utc), + ) + + def test_unrecognized_duration_is_the_one_day_window_the_scheduler_falls_back_to(self): + self.assertEqual( + get_budget_window_start("1hr", datetime(2024, 1, 16, tzinfo=timezone.utc)), + datetime(2024, 1, 15, tzinfo=timezone.utc), + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/proxy/utils/helpers/test_month_end_projection.py b/tests/unit/proxy/utils/helpers/test_month_end_projection.py index 8c1a4ae746d..1f2a88b209d 100644 --- a/tests/unit/proxy/utils/helpers/test_month_end_projection.py +++ b/tests/unit/proxy/utils/helpers/test_month_end_projection.py @@ -1,4 +1,5 @@ -from datetime import date, timedelta +from datetime import date, datetime, timedelta, timezone +from typing import Final import pytest @@ -211,3 +212,181 @@ def test_get_projected_spend_over_limit_raises_when_today_missing(monkeypatch): monkeypatch.setattr("litellm.proxy.utils.date", _Broken) with pytest.raises(RuntimeError): get_projected_spend_over_limit(current_spend=1.0, soft_budget_limit=1.0) + + +UTC_RESET_JAN_16: Final = datetime(2024, 1, 16, tzinfo=timezone.utc) + + +def test_daily_budget_projects_pace_to_the_reset_not_month_end(): + result: Final = get_projected_spend_over_limit( + current_spend=8.0, + soft_budget_limit=10.0, + budget_duration="1d", + budget_reset_at=UTC_RESET_JAN_16, + now=datetime(2024, 1, 15, 12, tzinfo=timezone.utc), + ) + assert result == (16.0, date(2024, 1, 15)) + + +def test_daily_budget_late_in_the_day_does_not_double_todays_spend(): + assert ( + get_projected_spend_over_limit( + current_spend=6.0, + soft_budget_limit=10.0, + budget_duration="1d", + budget_reset_at=UTC_RESET_JAN_16, + now=datetime(2024, 1, 15, 23, tzinfo=timezone.utc), + ) + is None + ) + + +def test_naive_reset_time_is_read_as_utc(): + result: Final = get_projected_spend_over_limit( + current_spend=8.0, + soft_budget_limit=10.0, + budget_duration="1d", + budget_reset_at=datetime(2024, 1, 16), + now=datetime(2024, 1, 15, 12), + ) + assert result == (16.0, date(2024, 1, 15)) + + +def test_weekly_budget_measures_pace_from_the_previous_reset(): + result: Final = get_projected_spend_over_limit( + current_spend=8.0, + soft_budget_limit=10.0, + budget_duration="7d", + budget_reset_at=datetime(2024, 1, 18, tzinfo=timezone.utc), + now=datetime(2024, 1, 15, tzinfo=timezone.utc), + ) + assert result == (14.0, date(2024, 1, 16)) + + +def test_monthly_budget_window_starts_on_the_previous_reset_day(): + result: Final = get_projected_spend_over_limit( + current_spend=8.0, + soft_budget_limit=10.0, + budget_duration="1mo", + budget_reset_at=datetime(2024, 3, 1, tzinfo=timezone.utc), + now=datetime(2024, 2, 5, tzinfo=timezone.utc), + ) + assert result == (pytest.approx(58.0), date(2024, 2, 6)) + + +def test_sub_day_budget_is_measured_in_hours_not_days(): + result: Final = get_projected_spend_over_limit( + current_spend=3.0, + soft_budget_limit=5.0, + budget_duration="4h", + budget_reset_at=datetime(2024, 1, 15, 12, tzinfo=timezone.utc), + now=datetime(2024, 1, 15, 10, tzinfo=timezone.utc), + ) + assert result == (pytest.approx(6.0), date(2024, 1, 15)) + + +def test_first_minutes_after_a_reset_do_not_project_a_burst_across_the_window(): + assert ( + get_projected_spend_over_limit( + current_spend=0.10, + soft_budget_limit=10.0, + budget_duration="1d", + budget_reset_at=UTC_RESET_JAN_16, + now=datetime(2024, 1, 15, 0, 1, tzinfo=timezone.utc), + ) + is None + ) + + +def test_exceed_date_is_reported_in_the_reset_timezone(): + result: Final = get_projected_spend_over_limit( + current_spend=9.5, + soft_budget_limit=10.0, + budget_duration="1d", + budget_reset_at=datetime(2024, 1, 17, tzinfo=timezone(timedelta(hours=2))), + now=datetime(2024, 1, 15, 22, 10, tzinfo=timezone.utc), + ) + assert result is not None + assert result[1] == date(2024, 1, 16) + + +@pytest.mark.parametrize("soft_budget_limit", [9.99, 13.9, 13.99]) +def test_exceed_date_never_lands_after_the_reset(soft_budget_limit): + result: Final = get_projected_spend_over_limit( + current_spend=8.0, + soft_budget_limit=soft_budget_limit, + budget_duration="7d", + budget_reset_at=datetime(2024, 1, 18, tzinfo=timezone.utc), + now=datetime(2024, 1, 15, tzinfo=timezone.utc), + ) + assert result is not None + assert result[1] <= date(2024, 1, 18) + + +@pytest.mark.parametrize("current_spend, expected", [(4.0, False), (8.0, True)]) +def test_is_projected_spend_over_limit_follows_the_reset_window(current_spend, expected): + assert ( + is_projected_spend_over_limit( + current_spend=current_spend, + soft_budget_limit=10.0, + budget_duration="1d", + budget_reset_at=UTC_RESET_JAN_16, + now=datetime(2024, 1, 15, 12, tzinfo=timezone.utc), + ) + is expected + ) + + +def test_without_duration_keeps_month_end_behavior(monkeypatch): + _freeze_today(monkeypatch, date(2024, 1, 15)) + result: Final = get_projected_spend_over_limit(current_spend=8.0, soft_budget_limit=10.0) + assert result is not None + assert result[0] == pytest.approx(8.0 + (8.0 / 14) * 16) + + +def test_thirty_day_budget_measures_pace_from_the_first_of_the_month(): + result: Final = get_projected_spend_over_limit( + current_spend=1.0, + soft_budget_limit=10.0, + budget_duration="30d", + budget_reset_at=datetime(2024, 11, 1, tzinfo=timezone.utc), + now=datetime(2024, 10, 4, tzinfo=timezone.utc), + ) + assert result == (pytest.approx(1.0 + 672 / 72), date(2024, 10, 31)) + + +def test_hour_spelling_the_scheduler_resets_at_midnight_does_not_alert_before_midnight(): + assert ( + get_projected_spend_over_limit( + current_spend=1.0, + soft_budget_limit=10.0, + budget_duration="1hr", + budget_reset_at=UTC_RESET_JAN_16, + now=datetime(2024, 1, 15, 22, tzinfo=timezone.utc), + ) + is None + ) + + +def test_unrecognized_duration_projects_within_the_daily_window_it_resets_on(): + result: Final = get_projected_spend_over_limit( + current_spend=8.0, + soft_budget_limit=10.0, + budget_duration="fortnightly", + budget_reset_at=UTC_RESET_JAN_16, + now=datetime(2024, 1, 15, 12, tzinfo=timezone.utc), + ) + assert result == (16.0, date(2024, 1, 15)) + + +def test_reset_more_than_one_window_ahead_does_not_project(): + assert ( + get_projected_spend_over_limit( + current_spend=9.0, + soft_budget_limit=10.0, + budget_duration="1d", + budget_reset_at=UTC_RESET_JAN_16, + now=datetime(2024, 1, 14, 12, tzinfo=timezone.utc), + ) + is None + )