mirror of
https://github.com/tiennm99/litellm.git
synced 2026-06-18 09:32:08 +00:00
Krish Dholakia
ef42461c1e
* test: add __init__.py files * refactor: rename test folder to avoid naming conflict * test: update workflows * test: update tests * test: update imports * test: update tests * test: remove unused import * ci(test-litellm.yml): add pytest retry to github workflow * test: fix test
36 lines
992 B
Python
36 lines
992 B
Python
import asyncio
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../../..")
|
|
) # Adds the parent directory to the system path
|
|
|
|
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
|
|
|
|
|
def test_get_budget_reset_time():
|
|
"""
|
|
Test that the budget reset time is set to the first of the next month
|
|
"""
|
|
# Get the current date
|
|
now = datetime.now(timezone.utc)
|
|
|
|
# Calculate expected reset date (first of next month)
|
|
if now.month == 12:
|
|
expected_month = 1
|
|
expected_year = now.year + 1
|
|
else:
|
|
expected_month = now.month + 1
|
|
expected_year = now.year
|
|
expected_reset_at = datetime(expected_year, expected_month, 1, tzinfo=timezone.utc)
|
|
|
|
# Verify budget_reset_at is set to first of next month
|
|
assert get_budget_reset_time(budget_duration="1mo") == expected_reset_at
|