mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(proxy): cancel /team/spend/by_user in Postgres after 30s (#49922)
* fix(proxy): let the Admin UI cancel /team/spend/by_user at 60s with an opt-in statement timeout GET /team/spend/by_user takes an optional timeout_seconds (1 to 600). When set, the raw LiteLLM_SpendLogs aggregation runs under SET LOCAL statement_timeout and Postgres cancels it, and the route returns 503 instead of letting the query run for hours. Callers that omit the param are unchanged. The Admin UI sends timeout_seconds=60 and the Spend Per User Within Team card shows the server's error instead of an empty table * fix(proxy): cancel /team/spend/by_user in Postgres after 30s Prisma's HTTP client already gives up on the statement at 30s, so the query now runs under SET LOCAL statement_timeout = 30000 for every caller and the opt-in timeout_seconds param is gone. The Admin UI card shows the request error instead of the empty-state text. * fix(proxy): keep /team/spend/by_user on the read replica inside its timeout transaction --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
6db7001c05
commit
391c8013c9
6 changed files with 180 additions and 24 deletions
|
|
@ -34,6 +34,8 @@ _CONNECTION_CAPACITY_PHRASES: Final = (
|
||||||
_PRISMA_POOL_TIMEOUT_CODE: Final = "P2024"
|
_PRISMA_POOL_TIMEOUT_CODE: Final = "P2024"
|
||||||
_LOCK_TIMEOUT_SQLSTATE: Final = "55P03"
|
_LOCK_TIMEOUT_SQLSTATE: Final = "55P03"
|
||||||
_LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout"
|
_LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout"
|
||||||
|
_STATEMENT_TIMEOUT_SQLSTATE: Final = "57014"
|
||||||
|
_STATEMENT_TIMEOUT_PHRASE: Final = "canceling statement due to statement timeout"
|
||||||
|
|
||||||
|
|
||||||
def exception_chain(e: BaseException) -> Iterator[BaseException]:
|
def exception_chain(e: BaseException) -> Iterator[BaseException]:
|
||||||
|
|
@ -302,6 +304,17 @@ class PrismaDBExceptionHandler:
|
||||||
return True
|
return True
|
||||||
return _LOCK_TIMEOUT_PHRASE in str(e).lower()
|
return _LOCK_TIMEOUT_PHRASE in str(e).lower()
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def is_statement_timeout_error(e: Exception) -> bool:
|
||||||
|
"""True iff Postgres cancelled a statement for exceeding ``statement_timeout`` (SQLSTATE 57014)."""
|
||||||
|
import prisma
|
||||||
|
|
||||||
|
if not isinstance(e, _exception_types(prisma.errors.PrismaError)):
|
||||||
|
return False
|
||||||
|
if PrismaDBExceptionHandler.postgres_sqlstate(e) == _STATEMENT_TIMEOUT_SQLSTATE:
|
||||||
|
return True
|
||||||
|
return _STATEMENT_TIMEOUT_PHRASE in str(e).lower()
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def postgres_sqlstate(e: Exception) -> str | None:
|
def postgres_sqlstate(e: Exception) -> str | None:
|
||||||
"""The SQLSTATE Postgres attached to a failed statement, as prisma surfaces it, or None."""
|
"""The SQLSTATE Postgres attached to a failed statement, as prisma surfaces it, or None."""
|
||||||
|
|
|
||||||
|
|
@ -140,6 +140,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
||||||
team_membership_reservation_cache_key,
|
team_membership_reservation_cache_key,
|
||||||
)
|
)
|
||||||
from litellm.proxy.db.db_span import db_span
|
from litellm.proxy.db.db_span import db_span
|
||||||
|
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||||
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
|
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
|
||||||
from litellm.proxy.hooks.model_max_budget_limiter import (
|
from litellm.proxy.hooks.model_max_budget_limiter import (
|
||||||
RoutedModels,
|
RoutedModels,
|
||||||
|
|
@ -7473,6 +7474,19 @@ class _TeamUserSpendDbRow(TypedDict):
|
||||||
failed_requests: ReadOnly[int]
|
failed_requests: ReadOnly[int]
|
||||||
|
|
||||||
|
|
||||||
|
_TEAM_USER_SPEND_STATEMENT_TIMEOUT_SECONDS: Final = 30
|
||||||
|
|
||||||
|
|
||||||
|
async def _team_user_spend_rows(
|
||||||
|
prisma_client: PrismaClient, sql: str, params: tuple[str, ...]
|
||||||
|
) -> Sequence[_TeamUserSpendDbRow]:
|
||||||
|
async with prisma_client.read_db.tx(
|
||||||
|
timeout=timedelta(seconds=_TEAM_USER_SPEND_STATEMENT_TIMEOUT_SECONDS + 5)
|
||||||
|
) as tx:
|
||||||
|
await tx.execute_raw(f"SET LOCAL statement_timeout = {_TEAM_USER_SPEND_STATEMENT_TIMEOUT_SECONDS * 1000}")
|
||||||
|
return await tx.query_raw(sql, *params)
|
||||||
|
|
||||||
|
|
||||||
@router.get(
|
@router.get(
|
||||||
"/team/spend/by_user",
|
"/team/spend/by_user",
|
||||||
response_model=TeamUserSpendResponse,
|
response_model=TeamUserSpendResponse,
|
||||||
|
|
@ -7522,14 +7536,23 @@ async def get_team_spend_by_user(
|
||||||
|
|
||||||
own_user_only: Final = scope.api_key_filter is not None
|
own_user_only: Final = scope.api_key_filter is not None
|
||||||
user_param: Final = (user_api_key_dict.user_id or "",) if own_user_only else ()
|
user_param: Final = (user_api_key_dict.user_id or "",) if own_user_only else ()
|
||||||
async with db_span("team_user_spend", "LiteLLM_SpendLogs"):
|
try:
|
||||||
rows: Final[Sequence[_TeamUserSpendDbRow]] = await prisma_client.db.query_raw(
|
async with db_span("team_user_spend", "LiteLLM_SpendLogs"):
|
||||||
_team_user_spend_sql(team_count=len(scoped_team_ids), restrict_to_user=own_user_only),
|
rows: Final = await _team_user_spend_rows(
|
||||||
start_date,
|
prisma_client,
|
||||||
end_date,
|
_team_user_spend_sql(team_count=len(scoped_team_ids), restrict_to_user=own_user_only),
|
||||||
*scoped_team_ids,
|
(start_date, end_date, *scoped_team_ids, *user_param),
|
||||||
*user_param,
|
)
|
||||||
)
|
except Exception as e:
|
||||||
|
if PrismaDBExceptionHandler.is_statement_timeout_error(e):
|
||||||
|
raise _daily_activity_error(
|
||||||
|
status_code=503,
|
||||||
|
message=(
|
||||||
|
f"Spend per user took longer than {_TEAM_USER_SPEND_STATEMENT_TIMEOUT_SECONDS}s. "
|
||||||
|
"Pick a shorter date range or fewer teams"
|
||||||
|
),
|
||||||
|
) from e
|
||||||
|
raise
|
||||||
results: Final = tuple(
|
results: Final = tuple(
|
||||||
TeamUserSpendRow(
|
TeamUserSpendRow(
|
||||||
team_id=row["team_id"],
|
team_id=row["team_id"],
|
||||||
|
|
|
||||||
|
|
@ -953,6 +953,21 @@ def test_is_lock_timeout_error_excludes_other_failures(error: Exception) -> None
|
||||||
assert PrismaDBExceptionHandler.is_lock_timeout_error(error) is False
|
assert PrismaDBExceptionHandler.is_lock_timeout_error(error) is False
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("error", "expected"),
|
||||||
|
[
|
||||||
|
(_lock_timeout_error(message="canceling statement due to statement timeout", code="57014"), True),
|
||||||
|
(_lock_timeout_error(message="canceling statement due to statement timeout", code=None), True),
|
||||||
|
(_lock_timeout_error(), False),
|
||||||
|
(_capacity_error("sorry, too many clients already"), False),
|
||||||
|
(RuntimeError("canceling statement due to statement timeout"), False),
|
||||||
|
],
|
||||||
|
ids=["meta_57014", "message_only", "lock_timeout_55P03", "capacity_53300", "not_prisma"],
|
||||||
|
)
|
||||||
|
def test_is_statement_timeout_error(error: Exception, expected: bool) -> None:
|
||||||
|
assert PrismaDBExceptionHandler.is_statement_timeout_error(error) is expected
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
("body", "retryable"),
|
("body", "retryable"),
|
||||||
[
|
[
|
||||||
|
|
|
||||||
|
|
@ -13,6 +13,7 @@ import litellm
|
||||||
import pytest
|
import pytest
|
||||||
from fastapi import HTTPException, Request
|
from fastapi import HTTPException, Request
|
||||||
from fastapi.testclient import TestClient
|
from fastapi.testclient import TestClient
|
||||||
|
from prisma.errors import DataError
|
||||||
from pydantic import ValidationError
|
from pydantic import ValidationError
|
||||||
|
|
||||||
from litellm._uuid import uuid
|
from litellm._uuid import uuid
|
||||||
|
|
@ -17278,6 +17279,28 @@ def _team_spend_by_user_db_row(team_id: str, user_id: str, spend: float, request
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class _RecordingTransaction:
|
||||||
|
def __init__(self, query_raw: AsyncMock) -> None:
|
||||||
|
self.statements: tuple[str, ...] = ()
|
||||||
|
self.query_raw: AsyncMock = query_raw
|
||||||
|
|
||||||
|
async def execute_raw(self, sql: str) -> int:
|
||||||
|
self.statements = (*self.statements, sql)
|
||||||
|
return 0
|
||||||
|
|
||||||
|
async def __aenter__(self) -> "_RecordingTransaction":
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __aexit__(self, *exc_info: object) -> None:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def _team_spend_by_user_query(mock_db_client: MagicMock, query_raw: AsyncMock) -> _RecordingTransaction:
|
||||||
|
transaction: Final = _RecordingTransaction(query_raw)
|
||||||
|
mock_db_client.read_db.tx = MagicMock(return_value=transaction)
|
||||||
|
return transaction
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(mock_db_client):
|
async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(mock_db_client):
|
||||||
from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user
|
from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user
|
||||||
|
|
@ -17286,12 +17309,15 @@ async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(m
|
||||||
alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), [])
|
alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), [])
|
||||||
beta = _team_spend_by_user_team("team-beta", "Team Beta", Member(user_id="alice", role="user"), [])
|
beta = _team_spend_by_user_team("team-beta", "Team Beta", Member(user_id="alice", role="user"), [])
|
||||||
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha, beta])
|
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha, beta])
|
||||||
mock_db_client.db.query_raw = AsyncMock(
|
transaction = _team_spend_by_user_query(
|
||||||
return_value=[
|
mock_db_client,
|
||||||
_team_spend_by_user_db_row("team-alpha", "alice", 0.5, 3),
|
AsyncMock(
|
||||||
_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2),
|
return_value=[
|
||||||
_team_spend_by_user_db_row("team-beta", "alice", 0.1, 1),
|
_team_spend_by_user_db_row("team-alpha", "alice", 0.5, 3),
|
||||||
]
|
_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2),
|
||||||
|
_team_spend_by_user_db_row("team-beta", "alice", 0.1, 1),
|
||||||
|
]
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
response = await get_team_spend_by_user(
|
response = await get_team_spend_by_user(
|
||||||
|
|
@ -17301,7 +17327,9 @@ async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(m
|
||||||
end_date="2026-09-04",
|
end_date="2026-09-04",
|
||||||
)
|
)
|
||||||
|
|
||||||
sql, *params = mock_db_client.db.query_raw.call_args.args
|
assert any("statement_timeout = 30000" in statement for statement in transaction.statements)
|
||||||
|
mock_db_client.db.tx.assert_not_called()
|
||||||
|
sql, *params = transaction.query_raw.call_args.args
|
||||||
assert params == ["2026-09-01", "2026-09-04", "team-alpha", "team-beta"]
|
assert params == ["2026-09-01", "2026-09-04", "team-alpha", "team-beta"]
|
||||||
assert 'FROM "LiteLLM_SpendLogs" sl' in sql
|
assert 'FROM "LiteLLM_SpendLogs" sl' in sql
|
||||||
assert 'sl."startTime" >= $1::timestamp' in sql
|
assert 'sl."startTime" >= $1::timestamp' in sql
|
||||||
|
|
@ -17328,7 +17356,7 @@ async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_clien
|
||||||
caller = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER)
|
caller = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER)
|
||||||
alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), [])
|
alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), [])
|
||||||
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha])
|
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha])
|
||||||
mock_db_client.db.query_raw = AsyncMock(return_value=[])
|
transaction = _team_spend_by_user_query(mock_db_client, AsyncMock(return_value=[]))
|
||||||
mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock(
|
mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock(
|
||||||
return_value=_team_spend_by_user_caller("alice", ["team-alpha"])
|
return_value=_team_spend_by_user_caller("alice", ["team-alpha"])
|
||||||
)
|
)
|
||||||
|
|
@ -17337,7 +17365,7 @@ async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_clien
|
||||||
user_api_key_dict=caller, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04"
|
user_api_key_dict=caller, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04"
|
||||||
)
|
)
|
||||||
|
|
||||||
sql, *params = mock_db_client.db.query_raw.call_args.args
|
sql, *params = transaction.query_raw.call_args.args
|
||||||
assert params == ["2026-09-01", "2026-09-04", "team-alpha"]
|
assert params == ["2026-09-01", "2026-09-04", "team-alpha"]
|
||||||
assert 'sl."user" = $' not in sql
|
assert 'sl."user" = $' not in sql
|
||||||
|
|
||||||
|
|
@ -17350,7 +17378,9 @@ async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_cli
|
||||||
alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="bob", role="user"), ["/key/info"])
|
alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="bob", role="user"), ["/key/info"])
|
||||||
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha])
|
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha])
|
||||||
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||||
mock_db_client.db.query_raw = AsyncMock(return_value=[_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2)])
|
transaction = _team_spend_by_user_query(
|
||||||
|
mock_db_client, AsyncMock(return_value=[_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2)])
|
||||||
|
)
|
||||||
mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock(
|
mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock(
|
||||||
return_value=_team_spend_by_user_caller("bob", ["team-alpha"])
|
return_value=_team_spend_by_user_caller("bob", ["team-alpha"])
|
||||||
)
|
)
|
||||||
|
|
@ -17359,7 +17389,7 @@ async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_cli
|
||||||
user_api_key_dict=caller, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04"
|
user_api_key_dict=caller, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04"
|
||||||
)
|
)
|
||||||
|
|
||||||
sql, *params = mock_db_client.db.query_raw.call_args.args
|
sql, *params = transaction.query_raw.call_args.args
|
||||||
assert params == ["2026-09-01", "2026-09-04", "team-alpha", "bob"]
|
assert params == ["2026-09-01", "2026-09-04", "team-alpha", "bob"]
|
||||||
assert "sl.team_id IN ($3)" in sql
|
assert "sl.team_id IN ($3)" in sql
|
||||||
assert 'AND sl."user" = $4' in sql
|
assert 'AND sl."user" = $4' in sql
|
||||||
|
|
@ -17371,7 +17401,7 @@ async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_clie
|
||||||
from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user
|
from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user
|
||||||
|
|
||||||
caller = UserAPIKeyAuth(user_id="bob", user_role=LitellmUserRoles.INTERNAL_USER)
|
caller = UserAPIKeyAuth(user_id="bob", user_role=LitellmUserRoles.INTERNAL_USER)
|
||||||
mock_db_client.db.query_raw = AsyncMock(return_value=[])
|
transaction = _team_spend_by_user_query(mock_db_client, AsyncMock(return_value=[]))
|
||||||
mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock(
|
mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock(
|
||||||
return_value=_team_spend_by_user_caller("bob", ["team-alpha"])
|
return_value=_team_spend_by_user_caller("bob", ["team-alpha"])
|
||||||
)
|
)
|
||||||
|
|
@ -17382,7 +17412,7 @@ async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_clie
|
||||||
)
|
)
|
||||||
|
|
||||||
assert exc_info.value.status_code == 404
|
assert exc_info.value.status_code == 404
|
||||||
mock_db_client.db.query_raw.assert_not_called()
|
transaction.query_raw.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
@ -17400,7 +17430,7 @@ async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_clie
|
||||||
async def test_get_team_spend_by_user_rejects_bad_input(mock_db_client, team_ids, start_date, end_date, expected_error):
|
async def test_get_team_spend_by_user_rejects_bad_input(mock_db_client, team_ids, start_date, end_date, expected_error):
|
||||||
from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user
|
from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user
|
||||||
|
|
||||||
mock_db_client.db.query_raw = AsyncMock(return_value=[])
|
transaction = _team_spend_by_user_query(mock_db_client, AsyncMock(return_value=[]))
|
||||||
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||||
|
|
||||||
with pytest.raises(HTTPException) as exc_info:
|
with pytest.raises(HTTPException) as exc_info:
|
||||||
|
|
@ -17410,7 +17440,41 @@ async def test_get_team_spend_by_user_rejects_bad_input(mock_db_client, team_ids
|
||||||
|
|
||||||
assert exc_info.value.status_code == 400
|
assert exc_info.value.status_code == 400
|
||||||
assert expected_error in str(exc_info.value.detail)
|
assert expected_error in str(exc_info.value.detail)
|
||||||
mock_db_client.db.query_raw.assert_not_called()
|
transaction.query_raw.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
def _statement_timeout_error() -> DataError:
|
||||||
|
message = "canceling statement due to statement timeout"
|
||||||
|
return DataError(
|
||||||
|
data={
|
||||||
|
"user_facing_error": {
|
||||||
|
"is_panic": False,
|
||||||
|
"message": f"Error querying the database: {message}",
|
||||||
|
"meta": {"code": "57014", "message": message},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_team_spend_by_user_statement_timeout_returns_503(mock_db_client):
|
||||||
|
from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user
|
||||||
|
|
||||||
|
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||||
|
alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), [])
|
||||||
|
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha])
|
||||||
|
_team_spend_by_user_query(mock_db_client, AsyncMock(side_effect=_statement_timeout_error()))
|
||||||
|
|
||||||
|
with pytest.raises(HTTPException) as exc_info:
|
||||||
|
await get_team_spend_by_user(
|
||||||
|
user_api_key_dict=admin,
|
||||||
|
team_ids="team-alpha",
|
||||||
|
start_date="2026-09-01",
|
||||||
|
end_date="2026-09-04",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert exc_info.value.status_code == 503
|
||||||
|
assert "longer than 30s" in str(exc_info.value.detail)
|
||||||
|
|
||||||
|
|
||||||
class _TeamRowWithOrganization(LiteLLM_TeamTable):
|
class _TeamRowWithOrganization(LiteLLM_TeamTable):
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,39 @@
|
||||||
|
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||||
|
import { render, screen } from "@testing-library/react";
|
||||||
|
import React from "react";
|
||||||
|
import { describe, expect, it, vi } from "vitest";
|
||||||
|
|
||||||
|
import { teamSpendByUserCall } from "@/components/networking";
|
||||||
|
|
||||||
|
import TeamUserSpendCard from "./TeamUserSpendCard";
|
||||||
|
|
||||||
|
vi.mock("@/components/networking", async (importOriginal) => ({
|
||||||
|
...(await importOriginal<typeof import("@/components/networking")>()),
|
||||||
|
teamSpendByUserCall: vi.fn(),
|
||||||
|
}));
|
||||||
|
|
||||||
|
const renderCard = () => {
|
||||||
|
const client = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } });
|
||||||
|
return render(
|
||||||
|
<QueryClientProvider client={client}>
|
||||||
|
<TeamUserSpendCard
|
||||||
|
accessToken="sk-test"
|
||||||
|
startTime={new Date("2026-09-01T00:00:00Z")}
|
||||||
|
endTime={new Date("2026-09-30T00:00:00Z")}
|
||||||
|
teamIds={["team-alpha"]}
|
||||||
|
/>
|
||||||
|
</QueryClientProvider>,
|
||||||
|
);
|
||||||
|
};
|
||||||
|
|
||||||
|
describe("TeamUserSpendCard", () => {
|
||||||
|
it("shows the server's error instead of an empty table when the request fails", async () => {
|
||||||
|
const timeoutMessage = "Spend per user took longer than 30s. Pick a shorter date range or fewer teams";
|
||||||
|
vi.mocked(teamSpendByUserCall).mockRejectedValueOnce(new Error(timeoutMessage));
|
||||||
|
|
||||||
|
renderCard();
|
||||||
|
|
||||||
|
expect(await screen.findByText(timeoutMessage)).toBeInTheDocument();
|
||||||
|
expect(screen.queryByText("No user spend in this range")).not.toBeInTheDocument();
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
@ -73,7 +73,7 @@ const columns: ColumnDef<TeamUserSpendRow>[] = [
|
||||||
|
|
||||||
const TeamUserSpendCard: React.FC<TeamUserSpendCardProps> = ({ accessToken, startTime, endTime, teamIds }) => {
|
const TeamUserSpendCard: React.FC<TeamUserSpendCardProps> = ({ accessToken, startTime, endTime, teamIds }) => {
|
||||||
const hasTeams = teamIds.length > 0;
|
const hasTeams = teamIds.length > 0;
|
||||||
const { data, isLoading } = useQuery({
|
const { data, isLoading, error } = useQuery({
|
||||||
queryKey: ["teamSpendByUser", startTime?.toISOString(), endTime?.toISOString(), teamIds],
|
queryKey: ["teamSpendByUser", startTime?.toISOString(), endTime?.toISOString(), teamIds],
|
||||||
queryFn: () =>
|
queryFn: () =>
|
||||||
accessToken && startTime && endTime ? teamSpendByUserCall(accessToken, startTime, endTime, teamIds) : null,
|
accessToken && startTime && endTime ? teamSpendByUserCall(accessToken, startTime, endTime, teamIds) : null,
|
||||||
|
|
@ -105,6 +105,8 @@ const TeamUserSpendCard: React.FC<TeamUserSpendCardProps> = ({ accessToken, star
|
||||||
data={rows}
|
data={rows}
|
||||||
getRowId={teamUserSpendRowId}
|
getRowId={teamUserSpendRowId}
|
||||||
isLoading={isLoading}
|
isLoading={isLoading}
|
||||||
|
isError={error !== null}
|
||||||
|
errorMessage={error instanceof Error ? error.message : "Spend per user failed to load"}
|
||||||
maxBodyHeight={320}
|
maxBodyHeight={320}
|
||||||
noDataMessage={teamIds.length === 0 ? "Select a team to see spend per user" : "No user spend in this range"}
|
noDataMessage={teamIds.length === 0 ? "Select a team to see spend per user" : "No user spend in this range"}
|
||||||
size="compact"
|
size="compact"
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue