scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

Merge db8e525c43 into 66a422ea50

Browse source
This commit is contained in:
AxelRay 2026-10-03 16:13:35 +00:00 • committed by GitHub
commit ed33f88e6d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 403 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
litellm/proxy/management_endpoints/common_daily_activity.py
View file

@ -23,6 +23,8 @@ from litellm.proxy.spend_tracking.key_metadata_recovery import (
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
from litellm.proxy.utils import PrismaClient
from litellm.repositories.daily_activity_repository import DailyActivityRepository
from litellm.repositories.table_repositories import DeletedVerificationTokenRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
from litellm.types.proxy.management_endpoints.common_daily_activity import (
BreakdownMetrics,
DailySpendData,
@ -602,6 +604,38 @@ async def get_api_key_metadata(
}
async def _get_deleted_keys_for_user(
prisma_client: PrismaClient,
user_id: str,
) -> Sequence[object]:
return await DeletedVerificationTokenRepository(prisma_client).table.find_many(
where={"user_id": user_id},
order={"deleted_at": "desc"},
)
async def get_user_api_key_filter(
prisma_client: PrismaClient,
user_id: str,
api_key: str | None,
) -> list[str]:
"""Return the key digests that should scope a user activity query."""
active_keys: Final = await VerificationTokenRepository(prisma_client).table.find_many(
where={"user_id": user_id},
)
deleted_keys: Final = await _get_deleted_keys_for_user(prisma_client, user_id)
user_api_keys: Final = list(
dict.fromkeys(
key.token
for key in (*active_keys, *deleted_keys)
if getattr(key, "token", None) and getattr(key, "user_id", None) == user_id
)
)
if api_key is None:
return user_api_keys
return [api_key] if api_key in user_api_keys else []
def _aggregate_spend_records_sync(
*,
records: Sequence[DailySpendRecord],

84
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -55,7 +55,11 @@ from litellm.proxy.management.teams.access import is_team_admin
from litellm.proxy.management_endpoints.common_daily_activity import (
DailySpendRecord,
ScopeDenied,
daily_activity_repository,
daily_activity_scope,
get_daily_activity,
get_daily_activity_aggregated,
get_user_api_key_filter,
raise_public,
)
from litellm.proxy.management_endpoints.common_utils import (
@ -3002,3 +3006,83 @@ async def get_user_daily_activity(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={"error": f"Failed to fetch analytics: {e}"},
)
@router.get(
"/user/daily/activity/aggregated",
tags=["Budget & Spend Tracking", "Internal User management"],
dependencies=[Depends(user_api_key_auth)],
response_model=SpendAnalyticsPaginatedResponse,
)
@management_endpoint_wrapper
async def get_user_daily_activity_aggregated(
start_date: str | None = fastapi.Query(default=None, description="Start date in YYYY-MM-DD format"),
end_date: str | None = fastapi.Query(default=None, description="End date in YYYY-MM-DD format"),
model: str | None = fastapi.Query(default=None, description="Filter by specific model"),
api_key: str | None = fastapi.Query(default=None, description="Filter by specific API key"),
user_id: str | None = fastapi.Query(
default=None,
description="Filter by specific user ID. Admins can filter by any user or omit for global view. "
"Non-admins must provide their own user_id.",
),
timezone: int | None = fastapi.Query(
default=None,
description="Timezone offset in minutes from UTC.",
),
include_current_utc_day: bool = fastapi.Query(
default=False,
description="When the range ends on the caller's current local day, include today's UTC bucket.",
),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection
) -> SpendAnalyticsPaginatedResponse:
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if start_date is None or end_date is None:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={"error": "Please provide start_date and end_date"},
)
typed_prisma_client: Final["PrismaClient"] = prisma_client
try:
resolved_entity_ids: Final = resolve_user_daily_activity_entity_ids(
user_id=user_id,
user_api_key_dict=user_api_key_dict,
)
if isinstance(resolved_entity_ids, ScopeDenied):
raise_public(resolved_entity_ids)
requested_user_id: Final[str | None] = resolved_entity_ids[0] if resolved_entity_ids is not None else None
api_key_filter: Final = (
api_key
if requested_user_id is None
else await get_user_api_key_filter(typed_prisma_client, requested_user_id, api_key)
)
repository = daily_activity_repository(typed_prisma_client)
scope = daily_activity_scope(
table="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id=None,
exclude_entity_ids=None,
api_key=api_key_filter,
start_date=start_date,
end_date=end_date,
model=model,
timezone_offset_minutes=timezone,
include_current_utc_day=include_current_utc_day,
)
return await get_daily_activity_aggregated(repository, scope)
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception("/user/daily/activity/aggregated: Exception occured - %s", e)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={"error": f"Failed to fetch analytics: {e}"},
)

28
tests/unit/proxy/management_endpoints/test_common_daily_activity.py
View file

@ -19,6 +19,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import (
daily_activity_repository,
daily_activity_scope,
get_api_key_metadata,
get_user_api_key_filter,
get_daily_activity,
parse_canonical_date,
parse_canonical_date_range,
@ -2542,3 +2543,30 @@ async def test_get_daily_activity_rejects_non_canonical_dates_before_querying(st
assert error.value.detail == {"error": "start_date and end_date must be valid YYYY-MM-DD dates"}
mock_table.count.assert_not_awaited()
mock_table.find_many.assert_not_awaited()
@pytest.mark.asyncio
async def test_get_user_api_key_filter_scopes_to_active_and_deleted_user_keys():
mock_prisma = MagicMock()
active_key = SimpleNamespace(token="active-key", user_id="target-user")
unrelated_key = SimpleNamespace(token="unrelated-key", user_id="other-user")
deleted_key = SimpleNamespace(token="deleted-key", user_id="target-user")
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[active_key, unrelated_key])
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[deleted_key])
assert await get_user_api_key_filter(mock_prisma, "target-user", None) == ["active-key", "deleted-key"]
assert await get_user_api_key_filter(mock_prisma, "target-user", "unrelated-key") == []
@pytest.mark.asyncio
async def test_get_user_api_key_filter_fails_closed_when_deleted_lookup_fails():
mock_prisma = MagicMock()
active_key = SimpleNamespace(token="active-key", user_id="target-user")
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[active_key])
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(
side_effect=RuntimeError("deleted-key table unavailable")
)
import pytest
with pytest.raises(RuntimeError, match="deleted-key table unavailable"):
await get_user_api_key_filter(mock_prisma, "target-user", "active-key")

292
tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -2143,7 +2143,11 @@ def test_update_internal_user_params_preserves_model_budget_presence_and_neighbo
"user_alias": "Spruce",
}
replacement: Final = {"model-spruce": {"budget_limit": 0, "time_period": "1d"}, "model-birch": 5.0, "model-cedar": 0}
replacement: Final = {
"model-spruce": {"budget_limit": 0, "time_period": "1d"},
"model-birch": 5.0,
"model-cedar": 0,
}
request: Final = UpdateUserRequest(
user_id="user-spruce",
model_max_budget=replacement,
@ -2185,10 +2189,14 @@ async def test_user_model_budget_update_by_email_refreshes_cached_user(mocker: M
prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user)
prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user])
prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user})
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
mocker.patch(
"litellm.proxy.proxy_server.prisma_client", prisma_client
) # test-quality-ok: substitute the database dependency
cache: Final = UserApiKeyCache()
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
mocker.patch(
"litellm.proxy.proxy_server.user_api_key_cache", cache
) # test-quality-ok: exercise a real isolated cache
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new_callable=mocker.AsyncMock,
@ -2224,10 +2232,14 @@ async def test_user_status_update_refreshes_cached_user(
prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user)
prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user])
prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user})
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
mocker.patch(
"litellm.proxy.proxy_server.prisma_client", prisma_client
) # test-quality-ok: substitute the database dependency
cache: Final = UserApiKeyCache()
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
mocker.patch(
"litellm.proxy.proxy_server.user_api_key_cache", cache
) # test-quality-ok: exercise a real isolated cache
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new_callable=mocker.AsyncMock,
@ -2261,10 +2273,14 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke
prisma_client: Final = mocker.MagicMock()
prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user])
prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1)
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
mocker.patch(
"litellm.proxy.proxy_server.prisma_client", prisma_client
) # test-quality-ok: substitute the database dependency
cache: Final = UserApiKeyCache()
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
mocker.patch(
"litellm.proxy.proxy_server.user_api_key_cache", cache
) # test-quality-ok: exercise a real isolated cache
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new_callable=mocker.AsyncMock,
@ -2296,7 +2312,9 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke
@pytest.mark.asyncio
@pytest.mark.parametrize("all_users", [False, True], ids=["single-user", "bulk-all-users"])
async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: MockerFixture, all_users: bool) -> None:
async def test_user_max_budget_update_evicts_cached_user_on_every_worker(
mocker: MockerFixture, all_users: bool
) -> None:
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper, bulk_user_update
@ -2309,10 +2327,14 @@ async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker:
prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1)
prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user])
prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user})
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
mocker.patch(
"litellm.proxy.proxy_server.prisma_client", prisma_client
) # test-quality-ok: substitute the database dependency
cache: Final = UserApiKeyCache()
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
mocker.patch(
"litellm.proxy.proxy_server.user_api_key_cache", cache
) # test-quality-ok: exercise a real isolated cache
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new_callable=mocker.AsyncMock,
@ -2477,6 +2499,204 @@ async def test_get_user_daily_activity_rejects_service_account_caller(monkeypatc
mock_get_daily.assert_not_called()
@pytest.mark.asyncio
async def test_get_user_daily_activity_aggregated_requires_database(monkeypatch):
from litellm.proxy.management_endpoints.internal_user_endpoints import (
get_user_daily_activity_aggregated,
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
with pytest.raises(HTTPException) as exc_info:
await get_user_daily_activity_aggregated(
start_date="2026-01-01",
end_date="2026-01-01",
model=None,
api_key=None,
user_id=None,
timezone=None,
include_current_utc_day=False,
user_api_key_dict=UserAPIKeyAuth(user_id="caller", user_role=LitellmUserRoles.INTERNAL_USER),
)
assert exc_info.value.status_code == 500
assert "database" in str(exc_info.value.detail).lower()
@pytest.mark.asyncio
async def test_get_user_daily_activity_aggregated_requires_dates(monkeypatch):
from litellm.proxy.management_endpoints.internal_user_endpoints import (
get_user_daily_activity_aggregated,
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
with pytest.raises(HTTPException) as exc_info:
await get_user_daily_activity_aggregated(
start_date=None,
end_date="2026-01-01",
model=None,
api_key=None,
user_id=None,
timezone=None,
include_current_utc_day=False,
user_api_key_dict=UserAPIKeyAuth(user_id="caller", user_role=LitellmUserRoles.INTERNAL_USER),
)
assert exc_info.value.status_code == 400
assert "start_date" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_get_user_daily_activity_aggregated_scopes_non_admin(monkeypatch):
import litellm.proxy.management_endpoints.internal_user_endpoints as endpoints
prisma_client = MagicMock()
repository = object()
expected = SimpleNamespace(result="aggregated")
calls = {}
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
async def fake_key_filter(client, user_id, api_key):
calls["filter"] = (client, user_id, api_key)
return ["owned-digest"]
def fake_repository(client):
calls["repository"] = client
return repository
async def fake_aggregate(repository_arg, scope):
calls["aggregate"] = (repository_arg, scope)
return expected
monkeypatch.setattr(endpoints, "get_user_api_key_filter", fake_key_filter)
monkeypatch.setattr(endpoints, "daily_activity_repository", fake_repository)
monkeypatch.setattr(endpoints, "get_daily_activity_aggregated", fake_aggregate)
result = await endpoints.get_user_daily_activity_aggregated(
start_date="2026-01-01",
end_date="2026-01-02",
model="gpt-4o",
api_key="requested-digest",
user_id=None,
timezone=480,
include_current_utc_day=True,
user_api_key_dict=UserAPIKeyAuth(user_id="caller", user_role=LitellmUserRoles.INTERNAL_USER),
)
assert result is expected
assert calls["filter"] == (prisma_client, "caller", "requested-digest")
assert calls["repository"] is prisma_client
assert calls["aggregate"][0] is repository
scope = calls["aggregate"][1]
assert scope.table.value == "litellm_dailyuserspend"
assert scope.entity_ids is None
assert scope.api_keys == ("owned-digest",)
assert scope.start_date == "2026-01-01"
assert scope.end_date == "2026-01-02"
assert scope.model == "gpt-4o"
assert scope.timezone_offset_minutes == 480
assert scope.include_current_utc_day is True
@pytest.mark.asyncio
async def test_get_user_daily_activity_aggregated_admin_keeps_requested_key(monkeypatch):
import litellm.proxy.management_endpoints.internal_user_endpoints as endpoints
prisma_client = MagicMock()
repository = object()
expected = SimpleNamespace(result="admin-aggregated")
captured_scope = None
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
async def unexpected_key_filter(*args):
raise AssertionError("admin global scope must not resolve a user key filter")
def fake_repository(client):
assert client is prisma_client
return repository
async def fake_aggregate(repository_arg, scope):
nonlocal captured_scope
assert repository_arg is repository
captured_scope = scope
return expected
monkeypatch.setattr(endpoints, "get_user_api_key_filter", unexpected_key_filter)
monkeypatch.setattr(endpoints, "daily_activity_repository", fake_repository)
monkeypatch.setattr(endpoints, "get_daily_activity_aggregated", fake_aggregate)
result = await endpoints.get_user_daily_activity_aggregated(
start_date="2026-01-01",
end_date="2026-01-01",
model=None,
api_key="admin-digest",
user_id=None,
timezone=None,
include_current_utc_day=False,
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert result is expected
assert captured_scope is not None
assert captured_scope.entity_ids is None
assert captured_scope.api_keys == ("admin-digest",)
@pytest.mark.asyncio
async def test_get_user_daily_activity_aggregated_rejects_other_non_admin(monkeypatch):
from litellm.proxy.management_endpoints.internal_user_endpoints import (
get_user_daily_activity_aggregated,
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
with pytest.raises(HTTPException) as exc_info:
await get_user_daily_activity_aggregated(
start_date="2026-01-01",
end_date="2026-01-01",
model=None,
api_key=None,
user_id="other-user",
timezone=None,
include_current_utc_day=False,
user_api_key_dict=UserAPIKeyAuth(user_id="caller", user_role=LitellmUserRoles.INTERNAL_USER),
)
assert exc_info.value.status_code == 403
assert "only view their own" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_get_user_daily_activity_aggregated_maps_failures(monkeypatch):
import litellm.proxy.management_endpoints.internal_user_endpoints as endpoints
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
monkeypatch.setattr(endpoints, "daily_activity_repository", lambda client: object())
async def fail_aggregate(repository, scope):
raise RuntimeError("aggregate unavailable")
monkeypatch.setattr(endpoints, "get_daily_activity_aggregated", fail_aggregate)
with pytest.raises(HTTPException) as exc_info:
await endpoints.get_user_daily_activity_aggregated(
start_date="2026-01-01",
end_date="2026-01-01",
model=None,
api_key=None,
user_id=None,
timezone=None,
include_current_utc_day=False,
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert exc_info.value.status_code == 500
assert "aggregate unavailable" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_delete_user_cleans_up_created_by_invitation_links(mocker):
"""
@ -2511,24 +2731,12 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker):
mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[])
# Mock all delete_many calls
mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(
return_value=[]
)
mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(
return_value=0
)
mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(
return_value=1
)
mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(
return_value=0
)
mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(
return_value=0
)
mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(
return_value=1
)
mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(return_value=[])
mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(return_value=0)
mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=1)
mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0)
mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0)
mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
@ -2619,9 +2827,15 @@ async def test_delete_user_evicts_jwt_key_mapping_cache_of_its_keys(mocker):
mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0)
mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # test-quality-ok: substitute the database dependency
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time
mocker.patch(
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
) # test-quality-ok: substitute the database dependency
mocker.patch(
"litellm.proxy.proxy_server.user_api_key_cache", cache
) # test-quality-ok: exercise a real isolated cache
mocker.patch(
"litellm.proxy.proxy_server.proxy_logging_obj", None
) # test-quality-ok: delete_user reads it off proxy_server at call time
await delete_user(
data=DeleteUserRequest(user_ids=["jwt-user"]),
@ -4530,12 +4744,18 @@ async def test_delete_user_evicts_cached_user_rows(mocker: MockerFixture) -> Non
prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0)
prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0)
prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1)
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
mocker.patch(
"litellm.proxy.proxy_server.prisma_client", prisma_client
) # test-quality-ok: substitute the database dependency
cache: Final = UserApiKeyCache()
for row in (deleted, survivor):
await cache.async_set_cache(key=row.user_id, value=row, model_type=LiteLLM_UserTable)
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time
mocker.patch(
"litellm.proxy.proxy_server.user_api_key_cache", cache
) # test-quality-ok: exercise a real isolated cache
mocker.patch(
"litellm.proxy.proxy_server.proxy_logging_obj", None
) # test-quality-ok: delete_user reads it off proxy_server at call time
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new_callable=mocker.AsyncMock,
@ -4641,7 +4861,9 @@ async def test_ui_view_users_answers_a_db_outage_as_503_no_db_connection_not_as_
assert raised.value.code == "503"
assert raised.value.type == ProxyErrorTypes.no_db_connection
assert isinstance(raised.value.__cause__, httpx.ConnectError)
outage_logs: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and "ConnectError" in r.getMessage()]
outage_logs: Final = [
r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and "ConnectError" in r.getMessage()
]
assert outage_logs == ["Database unavailable during user search: ConnectError"]