scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

test(proxy): assert the global rollup split and scheduler through behavior, not SQL text or add_job arguments

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-16 12:13:57 +00:00
parent 7c5fa53d40
commit 834313af4b
2 changed files with 66 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

95
tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py
View file

@ -1647,30 +1647,6 @@ async def test_global_rollup_marker_read_failure_falls_back_to_the_per_key_table
await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
def test_aggregated_sql_splits_the_key_free_arm_at_the_marker_and_keeps_the_key_arm_per_key():
sql, params = _build_aggregated_sql_query(**_unfiltered_user_query(), global_rollup_through="2026-06-01")
marker_param: Final = f"${len(params)}"
assert params[-1] == "2026-06-01"
assert (
f'FROM "LiteLLM_DailyGlobalSpend"\n WHERE date >= $1 AND date <= $2 AND date <= {marker_param}'
in sql
)
assert (
f'FROM "LiteLLM_DailyUserSpend"\n WHERE date >= $1 AND date <= $2 AND date > {marker_param}' in sql
)
key_arm: Final = sql.split("UNION ALL\n (WITH top_api_keys")[1]
assert "LiteLLM_DailyGlobalSpend" not in key_arm
assert marker_param not in key_arm
def test_aggregated_sql_without_a_marker_reads_the_per_key_table_only():
sql, params = _build_aggregated_sql_query(**_unfiltered_user_query())
assert "LiteLLM_DailyGlobalSpend" not in sql
assert params[-1] == PTU_SENTINEL_API_KEY
_GLOBAL_SPEND_MIGRATION: Final = (
pathlib.Path(__file__).resolve().parents[4]
/ "litellm-proxy-extras"
@ -1687,7 +1663,10 @@ async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_
):
"""Day 1 is rolled up and day 2 is still open (never rolled up), so a marker of day 1 must
give the same response as reading everything per-key: day 1 from the global table, day 2
live, one grand total across both. The per-key arm stays on the user table throughout."""
live, one grand total across both. Per-key rows that land after the rollup then tell the
two sources apart: a late day 1 row is invisible to totals until the next reconcile while a
late day 2 row shows up at once, and both keys rank in the key breakdown, which stays
per-key throughout."""
n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 3
rows: Final = [
(
@ -1720,39 +1699,56 @@ async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_
)
_aggregated_postgresql.commit()
async def read(marker: str | None, sql_seen: list[str]):
async def read(marker: str | None):
await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
prisma = _prisma_with_marker(marker)
run_query = _psycopg_query_raw(_aggregated_postgresql, [])
async def query_raw(sql: str, *params: str):
sql_seen.append(sql)
return await run_query(sql, *params)
prisma.db.query_raw = query_raw
prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, [])
return await get_daily_activity_aggregated(
prisma_client=prisma,
entity_metadata_field=None,
**_unfiltered_user_query(),
)
per_key_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim
global_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim
from_per_key = await read(None, per_key_sql)
from_global = await read("2026-06-01", global_sql)
await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
from_per_key = await read(None)
from_global = await read("2026-06-01")
assert per_key_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 0
assert global_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 1
assert global_sql[0].count('FROM "LiteLLM_DailyUserSpend"') == 3
assert from_global.model_dump() == from_per_key.model_dump()
assert from_global.metadata.total_spend == pytest.approx(2 * sum(float(i + 1) for i in range(n_keys)))
seeded_spend: Final = 2 * sum(float(i + 1) for i in range(n_keys))
assert from_global.metadata.total_spend == pytest.approx(seeded_spend)
assert from_global.metadata.total_response_time_ms == 2 * n_keys * 10 * 25
assert from_global.metadata.total_timed_requests == 2 * n_keys
assert {day.date.isoformat() for day in from_global.results} == {"2026-06-01", "2026-06-02"}
assert len(from_global.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT
assert set(from_global.results[0].breakdown.model_groups) == {"gpt-5", "claude"}
with _aggregated_postgresql.cursor() as cur:
cur.executemany(
"""
INSERT INTO "LiteLLM_DailyUserSpend"
(id, user_id, date, api_key, model, model_group, custom_llm_provider,
endpoint, prompt_tokens, spend, api_requests, successful_requests)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
""",
[
("late-1", "user-late", "2026-06-01", "key-late-1", "gpt-5", "", "openai", None, 10, 1000.0, 1, 1),
("late-2", "user-late", "2026-06-02", "key-late-2", "gpt-5", "", "openai", None, 10, 500.0, 1, 1),
],
)
_aggregated_postgresql.commit()
late_per_key = await read(None)
late_global = await read("2026-06-01")
await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
assert late_per_key.metadata.total_spend == pytest.approx(seeded_spend + 1000.0 + 500.0)
assert late_global.metadata.total_spend == pytest.approx(seeded_spend + 500.0)
by_day: Final = {day.date.isoformat(): day for day in late_global.results}
assert by_day["2026-06-01"].metrics.spend == pytest.approx(seeded_spend / 2)
assert by_day["2026-06-02"].metrics.spend == pytest.approx(seeded_spend / 2 + 500.0)
assert by_day["2026-06-01"].breakdown.api_keys["key-late-1"].metrics.spend == pytest.approx(1000.0)
assert by_day["2026-06-02"].breakdown.api_keys["key-late-2"].metrics.spend == pytest.approx(500.0)
assert late_global.metadata.total_api_keys == n_keys + 2
@pytest.mark.asyncio
async def test_get_daily_activity_aggregated_reports_exact_limit_key_count_as_complete(
@ -1810,7 +1806,20 @@ async def test_get_daily_activity_aggregated_model_group_rollups_fall_back_to_mo
"""Rows stored with an empty or NULL model_group must land in the model_groups
breakdown under their model name instead of vanishing from the usage UI."""
rows: Final = [
("row-0", "user-0", "2026-06-01", "key-0", "gpt-5", "gpt-5-eu", "openai", "/v1/chat/completions", 10, 7.0, 1, 1),
(
"row-0",
"user-0",
"2026-06-01",
"key-0",
"gpt-5",
"gpt-5-eu",
"openai",
"/v1/chat/completions",
10,
7.0,
1,
1,
),
("row-1", "user-1", "2026-06-01", "key-1", "gpt-5", "", "openai", "/v1/chat/completions", 10, 3.0, 1, 1),
("row-2", "user-2", "2026-06-01", "key-2", "claude-x", None, "anthropic", "/v1/messages", 10, 2.0, 1, 1),
]

24
tests/test_litellm/proxy/proxy_server/test_lifecycle.py
View file

@ -28,6 +28,7 @@ from typing import List, Optional, Union
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from fastapi import FastAPI
from pydantic import BaseModel
from typing_extensions import TypedDict
@ -1042,8 +1043,8 @@ async def test_spend_report_locks_are_never_released():
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.release_lock.assert_not_awaited()
def _init_daily_global_spend_reconcile_job() -> tuple[MagicMock, MagicMock, MagicMock]:
scheduler = MagicMock()
def _init_daily_global_spend_reconcile_job() -> tuple[AsyncIOScheduler, MagicMock, MagicMock]:
scheduler = AsyncIOScheduler()
proxy_logging_obj = MagicMock()
proxy_logging_obj.alerting_handler = AsyncMock()
prisma_client = MagicMock()
@ -1058,28 +1059,31 @@ def _init_daily_global_spend_reconcile_job() -> tuple[MagicMock, MagicMock, Magi
def test_daily_global_spend_reconcile_job_is_scheduled_nightly_with_an_immediate_catch_up_run():
"""Startup schedules the LiteLLM_DailyGlobalSpend backfill a couple of minutes out, so a
fresh deploy switches usage reads to the global table without waiting for the nightly
run, and replaces any previous registration of the same job id."""
run, and after that it fires once a day at 00:30 UTC, when the previous UTC day is closed."""
from datetime import datetime, timedelta, timezone
from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID
scheduler, _, _ = _init_daily_global_spend_reconcile_job()
job = scheduler.get_job(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID)
assert job is not None
(call,) = scheduler.add_job.call_args_list
assert call.kwargs["id"] == DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID
assert call.kwargs["replace_existing"] is True
assert call.args[1:] == ("cron",)
assert (call.kwargs["hour"], call.kwargs["minute"], call.kwargs["timezone"]) == (0, 30, "UTC")
assert timedelta(0) < call.kwargs["next_run_time"] - datetime.now(timezone.utc) <= timedelta(minutes=2)
assert timedelta(0) < job.next_run_time - datetime.now(timezone.utc) <= timedelta(minutes=2)
after_catch_up = datetime(2026, 9, 16, 12, 0, tzinfo=timezone.utc)
assert job.trigger.get_next_fire_time(None, after_catch_up) == datetime(2026, 9, 17, 0, 30, tzinfo=timezone.utc)
just_after_a_run = datetime(2026, 9, 17, 0, 30, 1, tzinfo=timezone.utc)
assert job.trigger.get_next_fire_time(None, just_after_a_run) == datetime(2026, 9, 18, 0, 30, tzinfo=timezone.utc)
@pytest.mark.asyncio
async def test_daily_global_spend_reconcile_job_runs_under_the_pod_lock_and_alerts_through_the_proxy(monkeypatch):
from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID
scheduler, proxy_logging_obj, prisma_client = _init_daily_global_spend_reconcile_job()
run = AsyncMock()
monkeypatch.setattr(ps, "run_scheduled_daily_global_spend_reconcile", run)
await scheduler.add_job.call_args.args[0]()
await scheduler.get_job(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID).func()
run.assert_awaited_once()
assert run.await_args.args == (prisma_client,)