mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(usage): keep spend when metadata lookups hit database failures (#50253)
* fix(usage): tolerate entity metadata database failures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(usage): keep spend when metadata lookups hit database failures Treat transport errors and the query engine crash on a closed connection as database failures in the key, user and guardrail lookups of the aggregated, keys, users, model_top_keys and cache_leakage_keys routes, and keep prisma out of module-level imports so a proxy install without prisma still imports the usage endpoints. * fix(usage): fall back only on database outages, on every daily activity route Data-layer Prisma errors such as a rejected query or a missing table return 500 again, a key whose details failed to load reports key_exists null instead of false, and the non-aggregated daily activity routes keep spend when the key or entity metadata lookup hits an outage * fix(proxy): skip the key metadata lookup when the daily rows carry no keys * fix(usage): keep the guardrail breakdown fallback for any Prisma error A missing guardrail usage table or column made the Usage call return 500 after the outage-only narrowing. The guardrail lookup now falls back on outages and on any Prisma error, as it did before this PR * fix(usage): keep spend when the team, organization, customer, or agent name lookup drops --------- Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
e60e642290
commit
5d66f2d8ff
9 changed files with 579 additions and 68 deletions
|
|
@ -69,7 +69,7 @@ from litellm.proxy.agent_endpoints.kill_switch import (
|
|||
from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity, unless_database_fails
|
||||
from litellm.proxy.utils import PrismaClient, get_custom_url
|
||||
from litellm.repositories.chunked_in import find_many_in
|
||||
from litellm.types.agents import (
|
||||
|
|
@ -394,10 +394,12 @@ async def resolve_agent_daily_activity_scope(
|
|||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
agent_records: Final = (
|
||||
await agents_table(prisma_client).find_many(where={})
|
||||
agent_records: Final = await unless_database_fails(
|
||||
agents_table(prisma_client).find_many(where={})
|
||||
if resolved_agent_ids is None
|
||||
else await find_many_in(agents_table(prisma_client), "agent_id", resolved_agent_ids)
|
||||
else find_many_in(agents_table(prisma_client), "agent_id", resolved_agent_ids),
|
||||
(),
|
||||
"its agent names",
|
||||
)
|
||||
agent_metadata: Final[Mapping[str, dict[str, object]]] = MappingProxyType(
|
||||
{agent.agent_id: {"agent_name": agent.agent_name} for agent in agent_records}
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ from dataclasses import dataclass, replace
|
|||
from datetime import date, datetime, timedelta
|
||||
from itertools import groupby
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, NoReturn, Protocol
|
||||
from typing import Final, Literal, NoReturn, Protocol, TypeVar
|
||||
|
||||
from fastapi import HTTPException, status
|
||||
from typing_extensions import ReadOnly, TypedDict, assert_never
|
||||
|
|
@ -14,6 +14,7 @@ from litellm import constants
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY, SUBSCRIPTION_BACKED_PROVIDERS
|
||||
from litellm.proxy._types import CommonProxyErrors
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.spend_tracking.key_metadata_recovery import (
|
||||
attach_user_details,
|
||||
recover_cli_session_key_metadata,
|
||||
|
|
@ -56,6 +57,8 @@ from litellm.types.repositories.daily_activity import (
|
|||
SpendLogsWindow,
|
||||
)
|
||||
|
||||
_T: Final = TypeVar("_T")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ScopeDenied:
|
||||
|
|
@ -188,9 +191,11 @@ class _AggregatedSpendData(TypedDict):
|
|||
totals: ReadOnly[SpendMetrics]
|
||||
|
||||
|
||||
def _key_metadata(api_key_metadata: Mapping[str, KeyMetadataRow], api_key: str | None) -> KeyMetadata:
|
||||
def _key_metadata(api_key_metadata: Mapping[str, KeyMetadataRow] | None, api_key: str | None) -> KeyMetadata:
|
||||
if not api_key:
|
||||
return KeyMetadata(key_alias="Unattributed", key_exists=False)
|
||||
if api_key_metadata is None:
|
||||
return KeyMetadata()
|
||||
meta: Final = api_key_metadata.get(api_key)
|
||||
return KeyMetadata(
|
||||
key_alias=meta.key_alias if meta is not None else None,
|
||||
|
|
@ -306,7 +311,7 @@ def update_breakdown_metrics(
|
|||
record: DailySpendRecord,
|
||||
model_metadata: Mapping[str, dict[str, object]],
|
||||
provider_metadata: Mapping[str, dict[str, object]],
|
||||
api_key_metadata: Mapping[str, KeyMetadataRow],
|
||||
api_key_metadata: Mapping[str, KeyMetadataRow] | None,
|
||||
entity_id_field: str | None = None,
|
||||
entity_metadata_field: Mapping[str, dict[str, object]] | None = None,
|
||||
) -> BreakdownMetrics:
|
||||
|
|
@ -626,7 +631,7 @@ async def get_api_key_metadata(
|
|||
def _aggregate_spend_records_sync(
|
||||
*,
|
||||
records: Sequence[DailySpendRecord],
|
||||
api_key_metadata: Mapping[str, KeyMetadataRow],
|
||||
api_key_metadata: Mapping[str, KeyMetadataRow] | None,
|
||||
entity_id_field: str | None,
|
||||
entity_metadata_field: Mapping[str, dict[str, object]] | None,
|
||||
) -> _AggregatedSpendData:
|
||||
|
|
@ -685,22 +690,10 @@ async def _aggregate_spend_records(
|
|||
The per-row loop is offloaded to a worker thread via asyncio.to_thread so
|
||||
a large result set doesn't peg the event loop.
|
||||
"""
|
||||
api_keys: Final[set[str]] = {
|
||||
record.api_key for record in records if record.api_key and record.api_key != PTU_SENTINEL_API_KEY
|
||||
}
|
||||
|
||||
api_key_metadata: Final[Mapping[str, KeyMetadataRow]] = (
|
||||
await repository.key_metadata(
|
||||
frozenset(api_keys), spend_logs_window(frozenset(record.date for record in records))
|
||||
)
|
||||
if api_keys
|
||||
else MappingProxyType({})
|
||||
)
|
||||
|
||||
return await asyncio.to_thread(
|
||||
_aggregate_spend_records_sync,
|
||||
records=records,
|
||||
api_key_metadata=api_key_metadata,
|
||||
api_key_metadata=await _key_metadata_unless_database_fails(repository, records),
|
||||
entity_id_field=entity_id_field,
|
||||
entity_metadata_field=entity_metadata_field,
|
||||
)
|
||||
|
|
@ -783,7 +776,7 @@ def _record_to_spend_metrics(record: RollupMetricsRow) -> SpendMetrics:
|
|||
def _aggregate_grouping_sets_records_sync(
|
||||
*,
|
||||
records: Sequence[GroupingSetsRow],
|
||||
api_key_metadata: Mapping[str, KeyMetadataRow],
|
||||
api_key_metadata: Mapping[str, KeyMetadataRow] | None,
|
||||
) -> _AggregatedSpendData:
|
||||
"""Build the response from rollup rows produced by the GROUPING SETS query.
|
||||
|
||||
|
|
@ -919,24 +912,40 @@ def _aggregate_grouping_sets_records_sync(
|
|||
return {"results": results, "totals": total_metrics}
|
||||
|
||||
|
||||
async def unless_database_fails(lookup: Awaitable[_T], fallback: _T, left_out: str) -> _T:
|
||||
try:
|
||||
return await lookup
|
||||
except Exception as exc:
|
||||
if not PrismaDBExceptionHandler.is_database_service_unavailable_error(exc):
|
||||
raise
|
||||
verbose_proxy_logger.warning("Returning daily activity without %s: %s", left_out, exc)
|
||||
return fallback
|
||||
|
||||
|
||||
async def _key_metadata_unless_database_fails(
|
||||
repository: DailyActivityRepository,
|
||||
rows: Sequence[GroupingSetsRow] | Sequence[EntityRollupRow] | Sequence[DailySpendRecord],
|
||||
) -> Mapping[str, KeyMetadataRow] | None:
|
||||
api_keys: Final = frozenset(row.api_key for row in rows if row.api_key and row.api_key != PTU_SENTINEL_API_KEY)
|
||||
if not api_keys:
|
||||
return MappingProxyType({})
|
||||
return await unless_database_fails(
|
||||
repository.key_metadata(api_keys, spend_logs_window(frozenset(row.date for row in rows))),
|
||||
None,
|
||||
"its key metadata",
|
||||
)
|
||||
|
||||
|
||||
async def _aggregate_grouping_sets_records(
|
||||
*,
|
||||
repository: DailyActivityRepository,
|
||||
records: Sequence[GroupingSetsRow],
|
||||
) -> _AggregatedSpendData:
|
||||
"""Async wrapper: fetch api_key_metadata, then dispatch on a worker thread."""
|
||||
api_keys: Final[set[str]] = {r.api_key for r in records if r.api_key and r.api_key != PTU_SENTINEL_API_KEY}
|
||||
|
||||
api_key_metadata: Final[Mapping[str, KeyMetadataRow]] = (
|
||||
await repository.key_metadata(frozenset(api_keys), spend_logs_window(frozenset(r.date for r in records)))
|
||||
if api_keys
|
||||
else MappingProxyType({})
|
||||
)
|
||||
|
||||
return await asyncio.to_thread(
|
||||
_aggregate_grouping_sets_records_sync,
|
||||
records=records,
|
||||
api_key_metadata=api_key_metadata,
|
||||
api_key_metadata=await _key_metadata_unless_database_fails(repository, records),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -991,7 +1000,13 @@ async def get_daily_activity(
|
|||
if resolve_entity_metadata is not None:
|
||||
resolved_entity_metadata = {
|
||||
**(entity_metadata_field or {}),
|
||||
**(await resolve_entity_metadata(daily_spend_data)),
|
||||
**(
|
||||
await unless_database_fails(
|
||||
resolve_entity_metadata(daily_spend_data),
|
||||
MappingProxyType[str, dict[str, object]]({}),
|
||||
"its entity metadata",
|
||||
)
|
||||
),
|
||||
}
|
||||
aggregated: Final = await _aggregate_spend_records(
|
||||
repository=repository,
|
||||
|
|
@ -1038,7 +1053,7 @@ def _fold_entity_rollups_sync(
|
|||
*,
|
||||
results: Sequence[DailySpendData],
|
||||
entity_rows: Sequence[EntityRollupRow],
|
||||
api_key_metadata: Mapping[str, KeyMetadataRow],
|
||||
api_key_metadata: Mapping[str, KeyMetadataRow] | None,
|
||||
entity_metadata_field: Mapping[str, Mapping[str, object]] | None,
|
||||
) -> None:
|
||||
"""Write breakdown.entities onto the already-built per-day results."""
|
||||
|
|
@ -1107,7 +1122,7 @@ async def _guardrail_spend_or_empty(
|
|||
from prisma.errors import PrismaError
|
||||
|
||||
try:
|
||||
return await repository.guardrail_spend(scope)
|
||||
return await unless_database_fails(repository.guardrail_spend(scope), (), "its guardrail breakdown")
|
||||
except PrismaError as exc:
|
||||
verbose_proxy_logger.warning("Returning daily activity without its guardrail breakdown: %s", exc)
|
||||
return ()
|
||||
|
|
@ -1117,7 +1132,9 @@ async def _user_entity_metadata(
|
|||
repository: DailyActivityRepository, entity_rows: Sequence[EntityRollupRow]
|
||||
) -> Mapping[str, Mapping[str, object]]:
|
||||
user_ids: Final = frozenset(row.entity_id for row in entity_rows if row.entity_id)
|
||||
users: Final = await repository.user_metadata(user_ids)
|
||||
users: Final = await unless_database_fails(
|
||||
repository.user_metadata(user_ids), MappingProxyType({}), "its entity user metadata"
|
||||
)
|
||||
return MappingProxyType(
|
||||
{
|
||||
user_id: MappingProxyType({"user_email": user.user_email, "user_alias": user.user_alias})
|
||||
|
|
@ -1160,16 +1177,7 @@ async def get_daily_activity_aggregated(
|
|||
)
|
||||
if aggregated_rows.entity_rows:
|
||||
entity_records: Final = aggregated_rows.entity_rows
|
||||
entity_api_keys: Final = frozenset(
|
||||
row.api_key for row in entity_records if row.api_key and row.api_key != PTU_SENTINEL_API_KEY
|
||||
)
|
||||
entity_key_metadata: Final[Mapping[str, KeyMetadataRow]] = (
|
||||
await repository.key_metadata(
|
||||
entity_api_keys, spend_logs_window(frozenset(row.date for row in entity_records))
|
||||
)
|
||||
if entity_api_keys
|
||||
else MappingProxyType({})
|
||||
)
|
||||
entity_key_metadata: Final = await _key_metadata_unless_database_fails(repository, entity_records)
|
||||
user_entity_metadata: Final[Mapping[str, Mapping[str, object]]] = (
|
||||
await _user_entity_metadata(repository, entity_records)
|
||||
if scope.table is DailyActivityTable.USER
|
||||
|
|
|
|||
|
|
@ -35,7 +35,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
end_user_cache_key,
|
||||
end_user_restricted_registry_cache_key,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity, unless_database_fails
|
||||
from litellm.proxy.management_endpoints.common_utils import validate_budget_duration
|
||||
from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports
|
||||
_set_object_permission, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
|
|
@ -522,10 +522,12 @@ async def resolve_customer_daily_activity_scope(
|
|||
prisma_client: "PrismaClient",
|
||||
) -> _CustomerDailyActivityScope:
|
||||
end_user_table: Final = _typed_table(EndUserRepository(prisma_client))
|
||||
end_user_aliases: Final = (
|
||||
await find_many_in(end_user_table, "user_id", end_user_ids)
|
||||
end_user_aliases: Final = await unless_database_fails(
|
||||
find_many_in(end_user_table, "user_id", end_user_ids)
|
||||
if end_user_ids is not None
|
||||
else await end_user_table.find_many(where={})
|
||||
else end_user_table.find_many(where={}),
|
||||
(),
|
||||
"its customer aliases",
|
||||
)
|
||||
metadata: Final = MappingProxyType({end_user.user_id: {"alias": end_user.alias} for end_user in end_user_aliases})
|
||||
return _CustomerDailyActivityScope(end_user_ids, metadata)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import json
|
|||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from dataclasses import asdict, fields, replace
|
||||
from datetime import date, datetime
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Final, Literal
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
|
|
@ -23,6 +24,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import (
|
|||
raise_public,
|
||||
spend_logs_window,
|
||||
sum_spend_metrics,
|
||||
unless_database_fails,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.daily_activity_scopes import (
|
||||
AGENT_RESOLVER,
|
||||
|
|
@ -149,9 +151,10 @@ async def _key_activity_rows(
|
|||
spend_window: Final[tuple[datetime, datetime] | None] = spend_logs_window(
|
||||
frozenset((resolved_scope.scope.start_date, resolved_scope.scope.end_date))
|
||||
)
|
||||
metadata: Final[Mapping[str, KeyMetadataRow]] = await repository.key_metadata(
|
||||
frozenset(row.api_key for row in rows if row.api_key),
|
||||
spend_window,
|
||||
metadata: Final[Mapping[str, KeyMetadataRow]] = await unless_database_fails(
|
||||
repository.key_metadata(frozenset(row.api_key for row in rows if row.api_key), spend_window),
|
||||
MappingProxyType({}),
|
||||
"its key metadata",
|
||||
)
|
||||
return [_key_activity_row(row, metadata) for row in rows]
|
||||
|
||||
|
|
@ -176,8 +179,10 @@ async def _user_activity_rows(
|
|||
repository: DailyActivityRepository,
|
||||
rows: Sequence[UserSpendRow],
|
||||
) -> Sequence[UserActivityRow]:
|
||||
metadata: Final[Mapping[str, UserMetadataRow]] = await repository.user_metadata(
|
||||
frozenset(user_id for row in rows if (user_id := row.user_id) is not None)
|
||||
metadata: Final[Mapping[str, UserMetadataRow]] = await unless_database_fails(
|
||||
repository.user_metadata(frozenset(user_id for row in rows if (user_id := row.user_id) is not None)),
|
||||
MappingProxyType({}),
|
||||
"its user metadata",
|
||||
)
|
||||
return [_user_activity_row(row, metadata) for row in rows]
|
||||
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ from litellm.proxy.management_endpoints.budget_management_endpoints import (
|
|||
new_budget,
|
||||
update_budget,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity, unless_database_fails
|
||||
from litellm.proxy.management_endpoints.common_utils import ( # noqa: F401 # legacy module exports
|
||||
_set_object_metadata_field, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
_user_has_admin_view, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
|
|
@ -655,10 +655,12 @@ async def resolve_organization_daily_activity_scope(
|
|||
)
|
||||
|
||||
organization_table: Final = _table(OrganizationRepository(prisma_client))
|
||||
organization_rows: Final = (
|
||||
await find_many_in(organization_table, "organization_id", resolved_organization_ids)
|
||||
organization_rows: Final = await unless_database_fails(
|
||||
find_many_in(organization_table, "organization_id", resolved_organization_ids)
|
||||
if resolved_organization_ids is not None
|
||||
else await organization_table.find_many(where={})
|
||||
else organization_table.find_many(where={}),
|
||||
(),
|
||||
"its organization aliases",
|
||||
)
|
||||
metadata: Final = MappingProxyType(
|
||||
{
|
||||
|
|
|
|||
|
|
@ -155,6 +155,7 @@ from litellm.proxy.management.teams.dependencies import get_team_access
|
|||
from litellm.proxy.management_endpoints.common_daily_activity import (
|
||||
InvalidDateRange,
|
||||
parse_canonical_date_range,
|
||||
unless_database_fails,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.common_utils import ( # noqa: F401 # legacy module exports
|
||||
_check_disable_global_guardrails_caller_permission, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
|
|
@ -7303,7 +7304,12 @@ async def resolve_team_daily_activity_scope(
|
|||
where_condition: Final[_TeamIdInFilter] = {}
|
||||
if team_ids_list:
|
||||
where_condition["team_id"] = {"in": list(team_ids_list)}
|
||||
team_aliases: Final = await _team_db(prisma_client).find_many(where=where_condition)
|
||||
team_alias_lookup: Final = _team_db(prisma_client).find_many(where=where_condition)
|
||||
team_aliases: Final = (
|
||||
await unless_database_fails(team_alias_lookup, [], "its team aliases")
|
||||
if user_api_key_has_admin_view(user_api_key_dict)
|
||||
else await team_alias_lookup
|
||||
)
|
||||
team_alias_metadata: Final = {t.team_id: {"team_alias": _as_object(t.team_alias)} for t in team_aliases}
|
||||
|
||||
# Check if user is team admin or has /team/daily/activity permission
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ lint.extend-select = [
|
|||
"PLW0133", "PYI030", "PYI041", "PYI064", "RET501", "RUF010", "RUF022", "RUF023", "RUF051", "S113",
|
||||
"SIM114", "SIM118", "SIM201", "SIM211", "SIM222", "TC005", "UP006", "UP007", "UP008",
|
||||
"UP012", "UP018", "UP024", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045",
|
||||
"C404", "C419",
|
||||
"C404", "C419", "TID253",
|
||||
]
|
||||
# RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip
|
||||
# `# noqa` directives that protect rules enforced elsewhere. List those codes as external
|
||||
|
|
@ -35,6 +35,9 @@ format.exclude = ["**/enterprise/**"]
|
|||
lint.exclude = ["litellm/__init__.py", "litellm/proxy/example_config_yaml/*", "tests/*"]
|
||||
|
||||
|
||||
[lint.flake8-tidy-imports]
|
||||
banned-module-level-imports = ["prisma"]
|
||||
|
||||
[lint.per-file-ignores]
|
||||
"litellm/main.py" = ["F401"]
|
||||
"litellm/types/caching.py" = ["F401"]
|
||||
|
|
|
|||
|
|
@ -1,11 +1,14 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import asdict, fields
|
||||
from datetime import date, datetime
|
||||
from types import SimpleNamespace
|
||||
from types import FunctionType, SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from prisma.errors import DataError, PrismaError, TableNotFoundError
|
||||
|
||||
import litellm.proxy.management_endpoints.common_daily_activity as common_daily_activity_module
|
||||
from litellm.constants import USAGE_TOP_API_KEYS_DEFAULT
|
||||
|
|
@ -32,10 +35,301 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_E
|
|||
from litellm.proxy.utils import PrismaClient, hash_token
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import (
|
||||
DailySpendMetadata,
|
||||
KeyMetadata,
|
||||
SpendAnalyticsPaginatedResponse,
|
||||
SpendMetrics,
|
||||
)
|
||||
from litellm.types.repositories.daily_activity import GroupingSetsRow, KeyMetadataRow
|
||||
from litellm.types.repositories.daily_activity import (
|
||||
AggregatedRows,
|
||||
DailyActivityScope,
|
||||
EntityRollupRow,
|
||||
GroupingSetsRow,
|
||||
KeyMetadataRow,
|
||||
)
|
||||
|
||||
|
||||
def _raise_missing_error_meta() -> None:
|
||||
raise AttributeError("'NoneType' object has no attribute 'get'")
|
||||
|
||||
|
||||
def _query_engine_crash_on_closed_connection() -> AttributeError:
|
||||
raise_inside_query_engine: Final = FunctionType(
|
||||
_raise_missing_error_meta.__code__, {"__name__": "prisma.engine.utils"}
|
||||
)
|
||||
with pytest.raises(AttributeError) as crash:
|
||||
raise_inside_query_engine()
|
||||
return crash.value
|
||||
|
||||
|
||||
DATABASE_FAILURES: Final = pytest.mark.parametrize(
|
||||
"failure",
|
||||
[
|
||||
DataError(
|
||||
{"user_facing_error": {"error_code": "P1001", "message": "Can't reach database server at `db:5432`"}}
|
||||
),
|
||||
PrismaError("metadata unavailable"),
|
||||
httpx.ReadTimeout("the query engine stopped answering"),
|
||||
httpx.ConnectError("the query engine refused the connection"),
|
||||
_query_engine_crash_on_closed_connection(),
|
||||
],
|
||||
ids=[
|
||||
"unreachable-database",
|
||||
"prisma-error",
|
||||
"engine-read-timeout",
|
||||
"engine-connect-error",
|
||||
"engine-crash-on-closed-connection",
|
||||
],
|
||||
)
|
||||
|
||||
CODE_BUGS: Final = (
|
||||
pytest.param(TypeError("not a database failure"), id="type-error"),
|
||||
pytest.param(AttributeError("not a database failure"), id="attribute-error"),
|
||||
)
|
||||
REJECTED_QUERIES: Final = (
|
||||
pytest.param(
|
||||
DataError({"user_facing_error": {"error_code": "P2010", "message": "not a database failure"}}),
|
||||
id="rejected-query",
|
||||
),
|
||||
pytest.param(
|
||||
TableNotFoundError(
|
||||
{
|
||||
"user_facing_error": {
|
||||
"error_code": "P2021",
|
||||
"message": "not a database failure",
|
||||
"meta": {"table": "LiteLLM_VerificationToken"},
|
||||
}
|
||||
}
|
||||
),
|
||||
id="missing-table",
|
||||
),
|
||||
)
|
||||
NOT_DATABASE_OUTAGES: Final = pytest.mark.parametrize("bug", (*CODE_BUGS, *REJECTED_QUERIES))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@DATABASE_FAILURES
|
||||
@pytest.mark.parametrize(
|
||||
("key_failure", "user_failure", "empty"),
|
||||
[(False, False, False), (True, False, False), (False, True, False), (True, True, False), (False, False, True)],
|
||||
)
|
||||
async def test_entity_metadata_outages_preserve_daily_spend_and_other_metadata(
|
||||
failure: Exception, key_failure: bool, user_failure: bool, empty: bool
|
||||
) -> None:
|
||||
prisma: Final = MagicMock()
|
||||
key: Final = SimpleNamespace(token="key", key_alias="Alias", team_id=None, user_id=None, metadata=None)
|
||||
user: Final = SimpleNamespace(user_id="user", user_email="user@example.com", user_alias="User")
|
||||
main_rows: Final = tuple(
|
||||
asdict(_grouping_row(level, api_key=api_key, spend=2.0))
|
||||
for level, api_key in ((127, None), (63, None), (31, "key"))
|
||||
)
|
||||
entity_rows: Final = tuple(
|
||||
{**asdict(_grouping_row(63, api_key=api_key, spend=2.0)), "entity_id": "user", "api_key_rolled": rolled}
|
||||
for rolled, api_key in ((1, None), (0, "key"))
|
||||
)
|
||||
prisma.db.query_raw = AsyncMock(side_effect=[main_rows, entity_rows])
|
||||
prisma.db.litellm_verificationtoken.find_many = AsyncMock(
|
||||
side_effect=[[key], failure if key_failure else ([] if empty else [key])]
|
||||
)
|
||||
prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
|
||||
prisma.db.litellm_usertable.find_many = AsyncMock(
|
||||
side_effect=failure if user_failure else None,
|
||||
return_value=[] if empty else [user],
|
||||
)
|
||||
|
||||
result: Final = await get_daily_activity_aggregated(
|
||||
prisma_client=prisma,
|
||||
table_name="litellm_dailyuserspend",
|
||||
entity_id_field="user_id",
|
||||
entity_id=None,
|
||||
start_date="2024-01-01",
|
||||
end_date="2024-01-01",
|
||||
model=None,
|
||||
api_key=None,
|
||||
include_entity_breakdown=True,
|
||||
)
|
||||
|
||||
day: Final = result.results[0]
|
||||
entity: Final = day.breakdown.entities["user"]
|
||||
assert result.metadata.total_spend == day.metrics.spend == entity.metrics.spend == 2.0
|
||||
assert day.breakdown.api_keys["key"].metadata.key_alias == "Alias"
|
||||
assert entity.metadata == (
|
||||
{} if user_failure or empty else {"user_email": "user@example.com", "user_alias": "User"}
|
||||
)
|
||||
assert entity.api_key_breakdown["key"].metrics.spend == 2.0
|
||||
assert entity.api_key_breakdown["key"].metadata == (
|
||||
KeyMetadata() if key_failure else KeyMetadata(key_alias=None if empty else "Alias", key_exists=not empty)
|
||||
)
|
||||
|
||||
|
||||
def _usage_scope() -> DailyActivityScope:
|
||||
return daily_activity_scope(
|
||||
"litellm_dailyuserspend", "user_id", None, None, None, "2024-01-01", "2024-01-01", None, None
|
||||
)
|
||||
|
||||
|
||||
def _entity_row() -> EntityRollupRow:
|
||||
grouping: Final = asdict(_grouping_row(63, api_key="key", spend=2.0))
|
||||
return EntityRollupRow(
|
||||
**{
|
||||
**{field.name: grouping[field.name] for field in fields(EntityRollupRow) if field.name in grouping},
|
||||
"entity_id": "user",
|
||||
"api_key_rolled": 0,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _usage_repository() -> MagicMock:
|
||||
repository: Final = MagicMock()
|
||||
repository.aggregated = AsyncMock(
|
||||
return_value=AggregatedRows(
|
||||
grouping_rows=tuple(
|
||||
_grouping_row(level, api_key=api_key, spend=2.0)
|
||||
for level, api_key in ((127, None), (63, None), (31, "key"))
|
||||
),
|
||||
entity_rows=(_entity_row(),),
|
||||
distinct_api_keys=1,
|
||||
)
|
||||
)
|
||||
repository.key_metadata = AsyncMock(return_value={})
|
||||
repository.user_metadata = AsyncMock(return_value={})
|
||||
repository.guardrail_spend = AsyncMock(return_value=())
|
||||
return repository
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@DATABASE_FAILURES
|
||||
async def test_usage_keeps_spend_when_the_first_key_lookup_hits_a_database_failure(failure: Exception) -> None:
|
||||
repository: Final = _usage_repository()
|
||||
repository.key_metadata = AsyncMock(side_effect=failure)
|
||||
|
||||
result: Final = await _get_daily_activity_aggregated(repository, _usage_scope())
|
||||
|
||||
day: Final = result.results[0]
|
||||
assert result.metadata.total_spend == day.metrics.spend == 2.0
|
||||
assert day.breakdown.api_keys["key"].metrics.spend == 2.0
|
||||
assert day.breakdown.api_keys["key"].metadata == KeyMetadata()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@DATABASE_FAILURES
|
||||
async def test_usage_keeps_spend_when_the_guardrail_lookup_hits_a_database_failure(failure: Exception) -> None:
|
||||
repository: Final = _usage_repository()
|
||||
repository.guardrail_spend = AsyncMock(side_effect=failure)
|
||||
|
||||
result: Final = await _get_daily_activity_aggregated(repository, _usage_scope(), include_guardrail_breakdown=True)
|
||||
|
||||
assert result.metadata.total_spend == result.results[0].metrics.spend == 2.0
|
||||
assert result.results[0].breakdown.guardrails == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure", REJECTED_QUERIES)
|
||||
async def test_usage_keeps_spend_when_the_guardrail_table_rejects_the_query(failure: Exception) -> None:
|
||||
repository: Final = _usage_repository()
|
||||
repository.guardrail_spend = AsyncMock(side_effect=failure)
|
||||
|
||||
result: Final = await _get_daily_activity_aggregated(repository, _usage_scope(), include_guardrail_breakdown=True)
|
||||
|
||||
assert result.metadata.total_spend == result.results[0].metrics.spend == 2.0
|
||||
assert result.results[0].breakdown.guardrails == {}
|
||||
|
||||
|
||||
async def _assert_usage_reports_the_bug(repository: MagicMock) -> None:
|
||||
with pytest.raises(HTTPException) as error:
|
||||
await _get_daily_activity_aggregated(
|
||||
repository, _usage_scope(), include_entity_breakdown=True, include_guardrail_breakdown=True
|
||||
)
|
||||
|
||||
assert error.value.status_code == 500
|
||||
assert error.value.detail == {"error": "Failed to fetch analytics: not a database failure"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("lookup", ["key_metadata", "user_metadata"])
|
||||
@NOT_DATABASE_OUTAGES
|
||||
async def test_usage_still_reports_a_lookup_bug_that_is_not_a_database_failure(lookup: str, bug: Exception) -> None:
|
||||
repository: Final = _usage_repository()
|
||||
setattr(repository, lookup, AsyncMock(side_effect=bug))
|
||||
|
||||
await _assert_usage_reports_the_bug(repository)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("bug", CODE_BUGS)
|
||||
async def test_usage_still_reports_a_guardrail_lookup_bug_that_is_not_a_database_error(bug: Exception) -> None:
|
||||
repository: Final = _usage_repository()
|
||||
repository.guardrail_spend = AsyncMock(side_effect=bug)
|
||||
|
||||
await _assert_usage_reports_the_bug(repository)
|
||||
|
||||
|
||||
def _daily_user_spend_prisma(key_lookup: AsyncMock) -> MagicMock:
|
||||
prisma: Final = MagicMock()
|
||||
records: Final = [_daily_user_spend_record(user_id="user", api_key="key", spend=7.0)]
|
||||
prisma.db.litellm_dailyuserspend.count = AsyncMock(return_value=len(records))
|
||||
prisma.db.litellm_dailyuserspend.find_many = AsyncMock(return_value=records)
|
||||
prisma.db.litellm_verificationtoken.find_many = key_lookup
|
||||
prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
|
||||
return prisma
|
||||
|
||||
|
||||
async def _daily_user_activity(prisma: MagicMock, user_emails: AsyncMock) -> SpendAnalyticsPaginatedResponse:
|
||||
return await get_daily_activity(
|
||||
prisma_client=prisma,
|
||||
table_name="litellm_dailyuserspend",
|
||||
entity_id_field="user_id",
|
||||
entity_id=None,
|
||||
entity_metadata_field=None,
|
||||
start_date="2024-01-01",
|
||||
end_date="2024-01-01",
|
||||
model=None,
|
||||
api_key=None,
|
||||
page=1,
|
||||
page_size=1000,
|
||||
resolve_entity_metadata=user_emails,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@DATABASE_FAILURES
|
||||
@pytest.mark.parametrize(("key_failure", "email_failure"), [(True, False), (False, True)])
|
||||
async def test_daily_activity_keeps_spend_when_a_metadata_lookup_hits_a_database_failure(
|
||||
failure: Exception, key_failure: bool, email_failure: bool
|
||||
) -> None:
|
||||
key: Final = SimpleNamespace(token="key", key_alias="Alias", team_id=None, user_id=None, metadata=None)
|
||||
prisma: Final = _daily_user_spend_prisma(
|
||||
AsyncMock(side_effect=failure if key_failure else None, return_value=[key])
|
||||
)
|
||||
user_emails: Final = AsyncMock(
|
||||
side_effect=failure if email_failure else None, return_value={"user": {"user_email": "user@example.com"}}
|
||||
)
|
||||
|
||||
result: Final = await _daily_user_activity(prisma, user_emails)
|
||||
|
||||
day: Final = result.results[0]
|
||||
assert result.metadata.total_spend == day.metrics.spend == 7.0
|
||||
assert day.breakdown.api_keys["key"].metadata == (
|
||||
KeyMetadata() if key_failure else KeyMetadata(key_alias="Alias", key_exists=True)
|
||||
)
|
||||
assert day.breakdown.entities["user"].metadata == ({} if email_failure else {"user_email": "user@example.com"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("lookup", ["key metadata", "user emails"])
|
||||
@NOT_DATABASE_OUTAGES
|
||||
async def test_daily_activity_still_reports_a_metadata_lookup_bug_that_is_not_a_database_failure(
|
||||
lookup: str, bug: Exception
|
||||
) -> None:
|
||||
prisma: Final = _daily_user_spend_prisma(
|
||||
AsyncMock(side_effect=bug if lookup == "key metadata" else None, return_value=[])
|
||||
)
|
||||
user_emails: Final = AsyncMock(side_effect=bug if lookup == "user emails" else None, return_value={})
|
||||
|
||||
with pytest.raises(HTTPException) as error:
|
||||
await _daily_user_activity(prisma, user_emails)
|
||||
|
||||
assert error.value.status_code == 500
|
||||
assert error.value.detail == {"error": "Failed to fetch analytics: not a database failure"}
|
||||
|
||||
|
||||
async def _run_aggregated_daily_activity(
|
||||
|
|
|
|||
|
|
@ -1,15 +1,18 @@
|
|||
import csv
|
||||
import io
|
||||
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, fields
|
||||
from itertools import chain
|
||||
from types import MappingProxyType, SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.testclient import TestClient
|
||||
from prisma.errors import DataError
|
||||
|
||||
from litellm import constants
|
||||
from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY
|
||||
|
|
@ -520,6 +523,14 @@ class _PrismaTable:
|
|||
return next((row for row in self._rows if _matches(row, where)), None)
|
||||
|
||||
|
||||
class _FailingPrismaTable:
|
||||
def __init__(self, failure: Exception) -> None:
|
||||
self._failure: Final = failure
|
||||
|
||||
async def find_many(self, **kwargs: object) -> tuple[object, ...]:
|
||||
raise self._failure
|
||||
|
||||
|
||||
def _matches(row: object, where: Mapping[str, object]) -> bool:
|
||||
return all(
|
||||
getattr(row, field_name, None) in value["in"]
|
||||
|
|
@ -529,7 +540,7 @@ def _matches(row: object, where: Mapping[str, object]) -> bool:
|
|||
)
|
||||
|
||||
|
||||
def _prisma_client() -> object:
|
||||
def _prisma_client(**table_overrides: object) -> object:
|
||||
user: Final = LiteLLM_UserTable(
|
||||
user_id="user-a",
|
||||
user_email="user@example.test",
|
||||
|
|
@ -552,7 +563,7 @@ def _prisma_client() -> object:
|
|||
team_alias="Team B",
|
||||
members_with_roles=[Member(user_id="user-b", role="user")],
|
||||
)
|
||||
db: Final = SimpleNamespace(
|
||||
default_tables: Final = SimpleNamespace(
|
||||
litellm_usertable=_PrismaTable((user, other_user)),
|
||||
litellm_teamtable=_PrismaTable((team, other_team)),
|
||||
litellm_verificationtoken=_PrismaTable((SimpleNamespace(token="key-alpha", user_id="user-a"),)),
|
||||
|
|
@ -581,13 +592,12 @@ def _prisma_client() -> object:
|
|||
)
|
||||
),
|
||||
)
|
||||
db: Final = SimpleNamespace(**{**vars(default_tables), **table_overrides})
|
||||
return SimpleNamespace(db=db, writer_db=db)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def daily_activity_client() -> Iterator[tuple[TestClient, _FakeRepository]]:
|
||||
repository: Final = _FakeRepository(_seeded_activity())
|
||||
prisma_client: Final = _prisma_client()
|
||||
@contextmanager
|
||||
def _daily_activity_test_client(repository: _FakeRepository, prisma_client: object) -> Iterator[TestClient]:
|
||||
app: Final = FastAPI()
|
||||
app.include_router(router)
|
||||
|
||||
|
|
@ -607,6 +617,13 @@ def daily_activity_client() -> Iterator[tuple[TestClient, _FakeRepository]]:
|
|||
app.dependency_overrides[get_daily_activity_llm_router] = lambda: None
|
||||
app.dependency_overrides[user_api_key_auth] = resolve_auth
|
||||
with TestClient(app) as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def daily_activity_client() -> Iterator[tuple[TestClient, _FakeRepository]]:
|
||||
repository: Final = _FakeRepository(_seeded_activity())
|
||||
with _daily_activity_test_client(repository, _prisma_client()) as client:
|
||||
yield client, repository
|
||||
|
||||
|
||||
|
|
@ -1015,6 +1032,178 @@ def test_user_page_route_ranks_users_by_spend_and_attaches_metadata(
|
|||
assert repository.user_metadata.call_args.args[0] == frozenset(("user-a", "user-b"))
|
||||
|
||||
|
||||
_LOOKUP_DATABASE_FAILURES: Final = pytest.mark.parametrize(
|
||||
"failure",
|
||||
[
|
||||
DataError({"user_facing_error": {"error_code": "P1001", "message": "Can't reach database server"}}),
|
||||
httpx.ConnectError("the query engine refused the connection"),
|
||||
],
|
||||
ids=["unreachable-database", "engine-connect-error"],
|
||||
)
|
||||
_KEY_ROWS_ROUTES: Final = pytest.mark.parametrize(
|
||||
("suffix", "extra_params"),
|
||||
[
|
||||
("keys", {}),
|
||||
("search", {"search": "key"}),
|
||||
("model_top_keys", {"model_group": "rare-group"}),
|
||||
("cache_leakage_keys", {}),
|
||||
],
|
||||
ids=["keys", "search", "model_top_keys", "cache_leakage_keys"],
|
||||
)
|
||||
|
||||
|
||||
@_LOOKUP_DATABASE_FAILURES
|
||||
@_KEY_ROWS_ROUTES
|
||||
def test_key_routes_keep_spend_when_the_key_lookup_hits_a_database_failure(
|
||||
daily_activity_client: tuple[TestClient, _FakeRepository],
|
||||
suffix: str,
|
||||
extra_params: Mapping[str, str],
|
||||
failure: Exception,
|
||||
) -> None:
|
||||
client, repository = daily_activity_client
|
||||
repository.key_metadata.side_effect = failure
|
||||
|
||||
response: Final = client.get(
|
||||
f"/user/daily/activity/aggregated/{suffix}",
|
||||
params={**_entity_params("user_id", "user-a"), **extra_params},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
rows: Final = response.json()["api_keys"]
|
||||
assert rows
|
||||
assert all(row["metrics"]["spend"] > 0 for row in rows)
|
||||
assert {row["metadata"]["key_alias"] for row in rows} == {None}
|
||||
|
||||
|
||||
@_KEY_ROWS_ROUTES
|
||||
def test_key_routes_still_report_a_lookup_bug_that_is_not_a_database_failure(
|
||||
daily_activity_client: tuple[TestClient, _FakeRepository],
|
||||
suffix: str,
|
||||
extra_params: Mapping[str, str],
|
||||
) -> None:
|
||||
client, repository = daily_activity_client
|
||||
repository.key_metadata.side_effect = AttributeError("not a database failure")
|
||||
|
||||
response: Final = client.get(
|
||||
f"/user/daily/activity/aggregated/{suffix}",
|
||||
params={**_entity_params("user_id", "user-a"), **extra_params},
|
||||
)
|
||||
|
||||
assert response.status_code == 500, response.text
|
||||
assert "not a database failure" in response.text
|
||||
|
||||
|
||||
@_LOOKUP_DATABASE_FAILURES
|
||||
def test_user_page_route_keeps_spend_when_the_user_lookup_hits_a_database_failure(
|
||||
daily_activity_client: tuple[TestClient, _FakeRepository],
|
||||
failure: Exception,
|
||||
) -> None:
|
||||
client, repository = daily_activity_client
|
||||
repository.user_metadata.side_effect = failure
|
||||
|
||||
response: Final = client.get("/user/daily/activity/aggregated/users", params=_DATE_PARAMS)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert [(user["user_id"], user["spend"], user["user_email"]) for user in response.json()["users"]] == [
|
||||
("user-b", 100.0, None),
|
||||
("user-a", 14.5, None),
|
||||
]
|
||||
|
||||
|
||||
def test_user_page_route_still_reports_a_lookup_bug_that_is_not_a_database_failure(
|
||||
daily_activity_client: tuple[TestClient, _FakeRepository],
|
||||
) -> None:
|
||||
client, repository = daily_activity_client
|
||||
repository.user_metadata.side_effect = AttributeError("not a database failure")
|
||||
|
||||
response: Final = client.get("/user/daily/activity/aggregated/users", params=_DATE_PARAMS)
|
||||
|
||||
assert response.status_code == 500, response.text
|
||||
assert "not a database failure" in response.text
|
||||
|
||||
|
||||
_ENTITY_ALIAS_LOOKUPS: Final = pytest.mark.parametrize(
|
||||
("prefix", "query_name", "entity_id", "alias_table", "alias"),
|
||||
[
|
||||
("/team", "team_ids", "team-a", "litellm_teamtable", {"team_alias": "Team A"}),
|
||||
("/organization", "organization_ids", "org-a", "litellm_organizationtable", {"organization_alias": "Org A"}),
|
||||
("/customer", "end_user_ids", "customer-a", "litellm_endusertable", {"alias": "Customer A"}),
|
||||
("/agent", "agent_ids", "agent-a", "litellm_agentstable", {"agent_name": "Agent A"}),
|
||||
],
|
||||
ids=["team", "organization", "customer", "agent"],
|
||||
)
|
||||
|
||||
|
||||
def _aggregated_entity_metadata(
|
||||
prefix: str, query_name: str, entity_id: str, prisma_client: object
|
||||
) -> tuple[httpx.Response, tuple[Mapping[str, object], ...]]:
|
||||
with _daily_activity_test_client(_FakeRepository(_seeded_activity()), prisma_client) as client:
|
||||
response: Final = client.get(f"{prefix}/daily/activity/aggregated", params=_entity_params(query_name, entity_id))
|
||||
if response.status_code != 200:
|
||||
return response, ()
|
||||
entities: Final = chain.from_iterable(
|
||||
result["breakdown"]["entities"].values() for result in response.json()["results"]
|
||||
)
|
||||
return response, tuple(entity["metadata"] for entity in entities)
|
||||
|
||||
|
||||
@_LOOKUP_DATABASE_FAILURES
|
||||
@_ENTITY_ALIAS_LOOKUPS
|
||||
def test_entity_routes_keep_spend_and_blank_the_alias_when_the_alias_lookup_hits_a_database_failure(
|
||||
prefix: str,
|
||||
query_name: str,
|
||||
entity_id: str,
|
||||
alias_table: str,
|
||||
alias: Mapping[str, object],
|
||||
failure: Exception,
|
||||
) -> None:
|
||||
healthy, healthy_metadata = _aggregated_entity_metadata(prefix, query_name, entity_id, _prisma_client())
|
||||
failed, failed_metadata = _aggregated_entity_metadata(
|
||||
prefix, query_name, entity_id, _prisma_client(**{alias_table: _FailingPrismaTable(failure)})
|
||||
)
|
||||
|
||||
assert failed.status_code == 200, failed.text
|
||||
assert failed.json()["metadata"] == healthy.json()["metadata"]
|
||||
assert healthy_metadata and set(map(str, healthy_metadata)) == {str(alias)}
|
||||
assert failed_metadata and set(map(str, failed_metadata)) == {str({})}
|
||||
|
||||
|
||||
@_ENTITY_ALIAS_LOOKUPS
|
||||
def test_entity_routes_still_report_an_alias_lookup_bug_that_is_not_a_database_failure(
|
||||
prefix: str,
|
||||
query_name: str,
|
||||
entity_id: str,
|
||||
alias_table: str,
|
||||
alias: Mapping[str, object],
|
||||
) -> None:
|
||||
failed, _ = _aggregated_entity_metadata(
|
||||
prefix,
|
||||
query_name,
|
||||
entity_id,
|
||||
_prisma_client(**{alias_table: _FailingPrismaTable(AttributeError("not a database failure"))}),
|
||||
)
|
||||
|
||||
assert failed.status_code == 500, failed.text
|
||||
assert "not a database failure" in failed.text
|
||||
|
||||
|
||||
@_LOOKUP_DATABASE_FAILURES
|
||||
def test_team_route_fails_closed_for_a_team_member_when_the_team_lookup_hits_a_database_failure(
|
||||
failure: Exception,
|
||||
) -> None:
|
||||
repository: Final = _FakeRepository(_seeded_activity())
|
||||
prisma_client: Final = _prisma_client(litellm_teamtable=_FailingPrismaTable(failure))
|
||||
with _daily_activity_test_client(repository, prisma_client) as client:
|
||||
response: Final = client.get(
|
||||
"/team/daily/activity/aggregated",
|
||||
params=_entity_params("team_ids", "team-a"),
|
||||
headers={"x-user-role": LitellmUserRoles.INTERNAL_USER.value, "x-user-id": "user-a"},
|
||||
)
|
||||
|
||||
assert response.status_code == 500, response.text
|
||||
repository.aggregated.assert_not_called()
|
||||
|
||||
|
||||
def test_user_page_route_forwards_offset_and_limit(
|
||||
daily_activity_client: tuple[TestClient, _FakeRepository],
|
||||
) -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue