mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
test(spend): cover the cursor check, representative guard, filters and ties
Adds a session with rows on both sides of the cursor, a row whose request_id equals another session's key, and a filtered page to the small integration seed, plus a unit test that a walk reads past a tied timestamp before settling.
This commit is contained in:
shrey-berri
parent
1f3f512e50
commit
9adfe67c34
3 changed files with 123 additions and 25 deletions
|
|
@ -304,8 +304,10 @@ def _small_rows(marker: str) -> tuple[_SeedRow, ...]:
|
|||
_SeedRow(f"{marker}-shared-b", f"{marker}-shared", at(60), api_key=f"{marker}-key-b"),
|
||||
_SeedRow(f"{marker}-tie-1", f"{marker}-tie", at(70)),
|
||||
_SeedRow(f"{marker}-tie-2", f"{marker}-tie", at(70)),
|
||||
_SeedRow(f"{marker}-mcp-old", f"{marker}-mcp", at(45)),
|
||||
_SeedRow(f"{marker}-mcp-llm", f"{marker}-mcp", at(80)),
|
||||
_SeedRow(f"{marker}-mcp-tool", f"{marker}-mcp", at(90), call_type="call_mcp_tool"),
|
||||
_SeedRow(f"{marker}-shared", f"{marker}-other", at(95), api_key=f"{marker}-key-b"),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -346,6 +348,7 @@ def _tie_as_one(request_id: str) -> str:
|
|||
|
||||
def _newest_first(marker: str) -> list[tuple[str, str]]:
|
||||
return [
|
||||
(f"{marker}-shared", f"{marker}-key-b"),
|
||||
(f"{marker}-mcp-llm", "K"),
|
||||
(f"{marker}-tie", "K"),
|
||||
(f"{marker}-shared-b", f"{marker}-key-b"),
|
||||
|
|
@ -360,10 +363,10 @@ def test_session_view_lists_each_session_once(gateway: Gateway, small_seed: str)
|
|||
|
||||
assert _sessions(page) == _newest_first(small_seed), (
|
||||
"each session once, represented by its newest non-MCP row: a row without a session_id newer than the rows "
|
||||
"adopting its request_id joins them, two rows at one startTime are one session, and one session_id under two "
|
||||
"api_keys is two sessions"
|
||||
"adopting its request_id joins them, two rows at one startTime are one session, one session_id under two "
|
||||
"api_keys is two sessions, and a row whose request_id equals another session's key stays in its own session"
|
||||
)
|
||||
assert page["total"] == 6
|
||||
assert page["total"] == 7
|
||||
assert page["has_more"] is False
|
||||
|
||||
|
||||
|
|
@ -382,9 +385,11 @@ def test_session_view_cursor_walk_counts_sessions_per_api_key(gateway: Gateway,
|
|||
)
|
||||
walked: Final = [session for page in cursor_pages for session in _sessions(page)]
|
||||
|
||||
assert first["total"] == 6, "the capped count must count one session_id under two api_keys as two sessions"
|
||||
assert walked == _newest_first(small_seed)
|
||||
assert [page["has_more"] for page in cursor_pages] == [True, True, False]
|
||||
assert first["total"] == 7, "the capped count must count one session_id under two api_keys as two sessions"
|
||||
assert walked == _newest_first(small_seed), (
|
||||
"a session whose older rows fall below the cursor was already shown and must not come back on a later page"
|
||||
)
|
||||
assert [page["has_more"] for page in cursor_pages] == [True, True, True, False]
|
||||
|
||||
|
||||
def test_session_view_ascending_and_offset_pages_match_the_newest_first_order(
|
||||
|
|
@ -394,6 +399,16 @@ def test_session_view_ascending_and_offset_pages_match_the_newest_first_order(
|
|||
offset_page: Final = _small_page(gateway, sort_order="desc", page_size="2", page="2")
|
||||
|
||||
assert _sessions(ascending) == _newest_first(small_seed)[::-1]
|
||||
assert ascending["total"] == 6
|
||||
assert ascending["total"] == 7
|
||||
assert _sessions(offset_page) == _newest_first(small_seed)[2:4]
|
||||
assert offset_page["has_more"] is True
|
||||
|
||||
|
||||
def test_session_view_filtered_page_only_lists_matching_sessions(gateway: Gateway, small_seed: str) -> None:
|
||||
page: Final = _small_page(gateway, sort_order="desc", page_size="50", page="1", api_key=f"{small_seed}-key-b")
|
||||
|
||||
assert _sessions(page) == [
|
||||
(f"{small_seed}-shared", f"{small_seed}-key-b"),
|
||||
(f"{small_seed}-shared-b", f"{small_seed}-key-b"),
|
||||
]
|
||||
assert page["total"] == 2
|
||||
|
|
|
|||
|
|
@ -13,7 +13,6 @@ from fastapi.testclient import TestClient
|
|||
|
||||
import litellm
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy.spend_tracking.spend_management_endpoints import _SESSION_HEAD_WALK_ROWS
|
||||
|
||||
|
||||
def _default_date_range():
|
||||
|
|
@ -7212,10 +7211,32 @@ SESSION_KEY_EXPR = "COALESCE(NULLIF(session_id, ''), request_id)"
|
|||
SESSION_GROUP_KEY_SQL = f"{SESSION_KEY_EXPR}, api_key"
|
||||
|
||||
|
||||
def _session_walk_batch(page_rows, rows_read):
|
||||
"""Walk-batch rows for ``page_rows`` (newest first): one first row per session plus the batch's oldest row."""
|
||||
first_rows = [
|
||||
{
|
||||
"session_key": row["session_key"],
|
||||
"api_key": row["api_key"],
|
||||
"last_activity": row["last_activity"],
|
||||
"sort_time": row["last_activity"],
|
||||
"request_id": f"req-{row['session_key']}",
|
||||
"is_edge": False,
|
||||
"batch_rows": 0,
|
||||
}
|
||||
for row in page_rows
|
||||
]
|
||||
edge = [{**first_rows[-1], "is_edge": True, "batch_rows": rows_read}] if first_rows else []
|
||||
return first_rows + edge
|
||||
|
||||
|
||||
def _session_grouped_mock_prisma(session_page_rows, session_total, representative_rows):
|
||||
"""Mock prisma dispatching the raw queries the keyset grouped path emits."""
|
||||
|
||||
async def mock_query_raw(sql_query, *params):
|
||||
if "AS newer_row" in sql_query:
|
||||
return []
|
||||
if "WITH batch AS MATERIALIZED" in sql_query:
|
||||
return _session_walk_batch(session_page_rows, 0)
|
||||
if "COUNT(*) AS total_count" in sql_query:
|
||||
return [{"total_count": session_total}]
|
||||
if "CROSS JOIN LATERAL" in sql_query or "DISTINCT ON" in sql_query:
|
||||
|
|
@ -7268,6 +7289,10 @@ def _session_grouped_paginating_prisma(sessions, counted_total=None):
|
|||
"""Mock prisma serving the grouped page query out of ``sessions``, honoring the LIMIT and OFFSET it asks for."""
|
||||
|
||||
async def mock_query_raw(sql_query, *params):
|
||||
if "AS newer_row" in sql_query:
|
||||
return []
|
||||
if "WITH batch AS MATERIALIZED" in sql_query:
|
||||
return _session_walk_batch([_session_page_row(key, activity) for key, activity in sessions], 0)
|
||||
if "COUNT(*) AS total_count" in sql_query:
|
||||
return [{"total_count": min(len(sessions) if counted_total is None else counted_total, params[-1])}]
|
||||
if "CROSS JOIN LATERAL" in sql_query or "DISTINCT ON" in sql_query:
|
||||
|
|
@ -7335,7 +7360,6 @@ async def test_ui_view_spend_logs_group_by_session_first_page(client, monkeypatc
|
|||
page_query_sql = emitted[0][0]
|
||||
assert "OFFSET" not in page_query_sql
|
||||
assert "HAVING" not in page_query_sql
|
||||
assert emitted[0][-1] == 3, "page query fetches page_size + 1 sessions to detect has_more"
|
||||
|
||||
count_sql = emitted[1][0]
|
||||
assert "LIMIT" in count_sql and "FROM (" in count_sql, "the grouped count must stay bounded"
|
||||
|
|
@ -7394,20 +7418,18 @@ async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypat
|
|||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
def _session_walk_mock_prisma(walked_rows, walk_window_rows, grouped_rows):
|
||||
"""Mock prisma serving the head walk, its window-size check and the grouped fallback page separately."""
|
||||
def _session_walk_mock_prisma(walked_rows, window_rows_per_batch, grouped_rows):
|
||||
"""Mock prisma whose walk batches keep returning ``walked_rows`` and report ``window_rows_per_batch`` rows read."""
|
||||
|
||||
async def mock_query_raw(sql_query, *params):
|
||||
if "AS walked_rows" in sql_query:
|
||||
return [{"total_count": walk_window_rows}]
|
||||
if "WITH batch AS MATERIALIZED" in sql_query:
|
||||
return _session_walk_batch(walked_rows, min(window_rows_per_batch, params[-1]))
|
||||
if "COUNT(*) AS total_count" in sql_query:
|
||||
return [{"total_count": 3}]
|
||||
if "CROSS JOIN LATERAL" in sql_query:
|
||||
return [_session_representative_row(f"req-{session_key}", session_key) for session_key in params[-2]]
|
||||
if "COALESCE(SUM(spend)" in sql_query:
|
||||
return []
|
||||
if "LEFT JOIN LATERAL" in sql_query:
|
||||
return walked_rows
|
||||
return grouped_rows
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
|
|
@ -7418,21 +7440,19 @@ def _session_walk_mock_prisma(walked_rows, walk_window_rows, grouped_rows):
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("walk_window_rows", "expected_request_ids", "expected_has_more"),
|
||||
("window_rows_per_batch", "expected_request_ids", "expected_has_more"),
|
||||
[
|
||||
pytest.param(
|
||||
_SESSION_HEAD_WALK_ROWS, ["req-grouped-1", "req-grouped-2"], True, id="walk-cut-short-groups-the-window"
|
||||
),
|
||||
pytest.param(_SESSION_HEAD_WALK_ROWS - 1, ["req-walked-1"], False, id="walk-covered-the-window"),
|
||||
pytest.param(10**9, ["req-grouped-1", "req-grouped-2"], True, id="walk-that-cannot-settle-groups-the-window"),
|
||||
pytest.param(10, ["req-walked-1"], False, id="walk-that-read-the-whole-window-is-trusted"),
|
||||
],
|
||||
)
|
||||
async def test_ui_view_spend_logs_group_by_session_falls_back_when_the_walk_is_cut_short(
|
||||
client, monkeypatch, walk_window_rows, expected_request_ids, expected_has_more
|
||||
async def test_ui_view_spend_logs_group_by_session_falls_back_when_the_walk_cannot_settle(
|
||||
client, monkeypatch, window_rows_per_batch, expected_request_ids, expected_has_more
|
||||
):
|
||||
"""A head walk that found fewer sessions than the page needs is only trusted when it read the whole window."""
|
||||
"""A walk that keeps finding too few sessions serves the grouped page; one that ran out of rows is exact."""
|
||||
mock_prisma = _session_walk_mock_prisma(
|
||||
walked_rows=[_session_page_row("walked-1", "2026-08-29 10:00:00")],
|
||||
walk_window_rows=walk_window_rows,
|
||||
window_rows_per_batch=window_rows_per_batch,
|
||||
grouped_rows=[
|
||||
_session_page_row("grouped-1", "2026-08-29 10:00:00"),
|
||||
_session_page_row("grouped-2", "2026-08-29 09:00:00"),
|
||||
|
|
@ -7462,6 +7482,64 @@ async def test_ui_view_spend_logs_group_by_session_falls_back_when_the_walk_is_c
|
|||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
def _walk_row(session_key, sort_time, is_edge=False, batch_rows=0):
|
||||
return {
|
||||
"session_key": session_key,
|
||||
"api_key": "hashed-key",
|
||||
"last_activity": sort_time,
|
||||
"sort_time": sort_time,
|
||||
"request_id": f"req-{session_key}",
|
||||
"is_edge": is_edge,
|
||||
"batch_rows": batch_rows,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_group_by_session_walk_finishes_a_tied_timestamp_before_settling(client, monkeypatch):
|
||||
"""Sessions tied on their newest startTime order by session key, so the walk reads past the tie before settling."""
|
||||
tied = "2026-08-29 10:00:00"
|
||||
|
||||
async def mock_query_raw(sql_query, *params):
|
||||
if "WITH batch AS MATERIALIZED" in sql_query:
|
||||
if '("startTime", request_id) <' not in sql_query:
|
||||
return [
|
||||
_walk_row("b", tied),
|
||||
_walk_row("a", tied),
|
||||
_walk_row("a", tied, is_edge=True, batch_rows=params[-1]),
|
||||
]
|
||||
return [_walk_row("c", tied), _walk_row("c", "2026-08-29 09:00:00", is_edge=True, batch_rows=1)]
|
||||
if "COUNT(*) AS total_count" in sql_query:
|
||||
return [{"total_count": 3}]
|
||||
if "CROSS JOIN LATERAL" in sql_query:
|
||||
return [_session_representative_row(f"req-{session_key}", session_key) for session_key in params[-2]]
|
||||
return []
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe",
|
||||
lambda user_api_key_dict: True,
|
||||
)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
|
||||
)
|
||||
try:
|
||||
start_date, end_date = _default_date_range()
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={"start_date": start_date, "end_date": end_date, "group_by_session": "true", "page_size": 1},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
data = response.json()
|
||||
assert [row["request_id"] for row in data["data"]] == ["req-c"]
|
||||
assert data["next_session_cursor"] == f"{tied}|hashed-key|c"
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_group_by_session_jumps_to_page_without_cursor(client, monkeypatch):
|
||||
"""page > 1 with no session_cursor (the UI's last-page jump) serves the sessions that page starts at."""
|
||||
|
|
|
|||
|
|
@ -535,6 +535,12 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch):
|
|||
]
|
||||
|
||||
async def mock_query_raw(sql_query, *params):
|
||||
if "WITH batch AS MATERIALIZED" in sql_query:
|
||||
first_rows = [
|
||||
{**row, "sort_time": row["last_activity"], "request_id": row["session_key"], "is_edge": False}
|
||||
for row in session_rows
|
||||
]
|
||||
return [*first_rows, {**first_rows[-1], "is_edge": True, "batch_rows": 0}]
|
||||
if "COUNT(*) AS total_count" in sql_query:
|
||||
return [{"total_count": 60}]
|
||||
if "CROSS JOIN LATERAL" in sql_query or "DISTINCT ON" in sql_query:
|
||||
|
|
@ -570,7 +576,6 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch):
|
|||
emitted = [call[0] for call in mock_prisma.db.query_raw.call_args_list]
|
||||
page_sql = emitted[0][0]
|
||||
assert "OFFSET" not in page_sql, "the startTime page must be keyset-selected, not offset-selected"
|
||||
assert emitted[0][-1] == 51, "the page query fetches page_size + 1 sessions to detect has_more"
|
||||
|
||||
count_call = emitted[1]
|
||||
count_sql = count_call[0]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue