491aa7ea51
1. check_code_and_doc_quality: add PLR0915 ignore for sandbox_executor.py
2. auth_ui_unit_tests: add prisma generate step (entrypoint.sh only runs
migration, not client generation)
3. proxy_store_model_in_db_tests: move does_mcp_server_exist outside
`if MCP_AVAILABLE:` so it's importable on Python < 3.10
4. build_and_test: fix datetime.fromisoformat('...Z') on Python 3.9
(Z suffix support was added in 3.11)
5. proxy_logging_guardrails: fix container name typo my-app-2 -> my-app-3
6. upload-coverage: use `uv tool run` instead of `uv run --with` to avoid
resolving the full workspace (which fails for Python 3.14)
93 lines
3.6 KiB
Python
93 lines
3.6 KiB
Python
# What is this?
|
|
## Unit tests for the /budget/* endpoints
|
|
from litellm._uuid import uuid
|
|
from datetime import datetime, timedelta
|
|
|
|
import aiohttp
|
|
import pytest
|
|
import pytest_asyncio
|
|
|
|
|
|
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 the 'budget_reset_at'
|
|
timestamp is correctly calculated as 'created_at' plus the budget duration (one day).
|
|
|
|
This test uses the budget_setup fixture, which handles both the creation and cleanup of the budget.
|
|
"""
|
|
|
|
# Verify that the response includes a 'budget_reset_at' timestamp.
|
|
assert (
|
|
budget_setup["budget_reset_at"] is not None
|
|
), "The budget_reset_at field should not be None"
|
|
|
|
# Calculate the expected reset time: created_at + 1 day.
|
|
# Replace trailing 'Z' with '+00:00' for Python 3.9 compat (fromisoformat
|
|
# only learned to accept 'Z' in Python 3.11).
|
|
created_at_str = budget_setup["created_at"].replace("Z", "+00:00")
|
|
expected_reset_at_date = datetime.fromisoformat(created_at_str) + timedelta(days=1)
|
|
|
|
# Allow for a small tolerance in seconds for the timestamp calculation.
|
|
tolerance_seconds = 3
|
|
reset_at_str = budget_setup["budget_reset_at"].replace("Z", "+00:00")
|
|
actual_reset_at_date = datetime.fromisoformat(reset_at_str)
|
|
time_difference = abs(
|
|
(actual_reset_at_date - expected_reset_at_date).total_seconds()
|
|
)
|
|
|
|
assert time_difference <= tolerance_seconds, (
|
|
f"Expected budget_reset_at to be within {tolerance_seconds} seconds of {expected_reset_at_date}, "
|
|
f"but the difference was {time_difference} seconds."
|
|
)
|