diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index ac761315a27..fc3ede88aa9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -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), ] diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index ee72e98ffa9..6121608b658 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -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,)