mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
test: cover aggregated user activity endpoint
This commit is contained in:
parent
f580e48422
commit
db8e525c43
2 changed files with 261 additions and 37 deletions
|
|
@ -3049,6 +3049,8 @@ async def get_user_daily_activity_aggregated(
|
|||
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,
|
||||
|
|
@ -3060,9 +3062,9 @@ async def get_user_daily_activity_aggregated(
|
|||
api_key_filter: Final = (
|
||||
api_key
|
||||
if requested_user_id is None
|
||||
else await get_user_api_key_filter(prisma_client, requested_user_id, api_key)
|
||||
else await get_user_api_key_filter(typed_prisma_client, requested_user_id, api_key)
|
||||
)
|
||||
repository = daily_activity_repository(prisma_client)
|
||||
repository = daily_activity_repository(typed_prisma_client)
|
||||
scope = daily_activity_scope(
|
||||
table="litellm_dailyuserspend",
|
||||
entity_id_field="user_id",
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue