102 lines
3.7 KiB
Python
102 lines
3.7 KiB
Python
# What is this?
|
|
## Unit tests for the /budget/* endpoints
|
|
from litellm._uuid import uuid
|
|
from datetime import datetime, timezone
|
|
|
|
import aiohttp
|
|
import pytest
|
|
import pytest_asyncio
|
|
|
|
from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time
|
|
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_timezone
|
|
|
|
|
|
def _parse_budget_api_datetime(value: str) -> datetime:
|
|
"""Parse ISO timestamps returned by the proxy JSON API."""
|
|
if value.endswith("Z"):
|
|
value = value[:-1] + "+00:00"
|
|
dt = datetime.fromisoformat(value)
|
|
if dt.tzinfo is None:
|
|
dt = dt.replace(tzinfo=timezone.utc)
|
|
return dt
|
|
|
|
|
|
async def delete_budget(session, budget_id):
|
|
url = "http://0.0.0.0:4000/budget/delete"
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
|
data = {"id": budget_id}
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
assert response.status == 200
|
|
print(f"Deleted Budget {budget_id}")
|
|
|
|
|
|
async def create_budget(session, data):
|
|
url = "http://0.0.0.0:4000/budget/new"
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
assert response.status == 200
|
|
response_data = await response.json()
|
|
budget_id = response_data["budget_id"]
|
|
print(f"Created Budget {budget_id}")
|
|
return response_data
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def budget_setup():
|
|
"""
|
|
Fixture to create a budget for testing and clean it up afterward.
|
|
|
|
This fixture performs the following steps:
|
|
1. Opens an aiohttp ClientSession.
|
|
2. Generates a random budget_id and defines the budget data (duration: 1 day, max_budget: 0.02).
|
|
3. Calls create_budget to create the budget.
|
|
4. Yields the budget_response (a dict) for use in the test.
|
|
5. After the test completes, deletes the created budget by calling delete_budget.
|
|
|
|
Returns:
|
|
dict: The JSON response from create_budget, which includes the created budget's data.
|
|
"""
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
# Generate a unique budget_id and define the budget data.
|
|
budget_id = f"budget-{uuid.uuid4()}"
|
|
data = {"budget_id": budget_id, "budget_duration": "1d", "max_budget": 0.02}
|
|
budget_response = await create_budget(session, data)
|
|
|
|
# Yield the response so the test can use it.
|
|
yield budget_response
|
|
|
|
# After the test, delete the created budget to clean up.
|
|
await delete_budget(session, budget_id)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_budget_with_duration(budget_setup):
|
|
"""
|
|
Test creating a budget with a specified duration and verify that 'budget_reset_at'
|
|
matches the next standardized reset (see get_budget_reset_time / new_budget), not
|
|
necessarily created_at + wall-clock duration.
|
|
"""
|
|
|
|
assert (
|
|
budget_setup["budget_reset_at"] is not None
|
|
), "The budget_reset_at field should not be None"
|
|
|
|
created_at = _parse_budget_api_datetime(budget_setup["created_at"])
|
|
expected_reset_at = get_next_standardized_reset_time(
|
|
duration=budget_setup["budget_duration"],
|
|
current_time=created_at,
|
|
timezone_str=get_budget_reset_timezone(),
|
|
)
|
|
|
|
actual_reset_at = _parse_budget_api_datetime(budget_setup["budget_reset_at"])
|
|
|
|
tolerance_seconds = 3
|
|
time_difference = abs((actual_reset_at - expected_reset_at).total_seconds())
|
|
|
|
assert time_difference <= tolerance_seconds, (
|
|
f"Expected budget_reset_at to be within {tolerance_seconds} seconds of {expected_reset_at}, "
|
|
f"but the difference was {time_difference} seconds."
|
|
)
|