From 2e65c464ade1a0dd5e86a2902f2e2fe1b01f1168 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 4 Dec 2025 12:36:15 -0800 Subject: [PATCH] Adding tests --- .../proxy/db/test_db_spend_update_writer.py | 75 +++++++++++++- .../test_customer_endpoints.py | 98 ++++++++++++++++++- 2 files changed, 171 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 181d21b44f..db6c318357 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -572,4 +572,77 @@ async def test_add_spend_log_transaction_to_daily_org_transaction_skips_when_org org_id=None, ) - writer.daily_org_spend_update_queue.add_update.assert_not_called() \ No newline at end of file + writer.daily_org_spend_update_queue.add_update.assert_not_called() + + +@pytest.mark.asyncio +async def test_add_spend_log_transaction_to_daily_end_user_transaction_injects_end_user_id_and_queues_update(): + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + mock_prisma.get_request_status = MagicMock(return_value="success") + + end_user_id = "end-user-xyz" + payload = { + "request_id": "req-1", + "user": "test-user", + "end_user": end_user_id, + "startTime": "2024-01-01T12:00:00", + "api_key": "test-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "model_group": "gpt-4-group", + "prompt_tokens": 10, + "completion_tokens": 5, + "spend": 0.2, + "metadata": '{"usage_object": {}}', + } + + writer.daily_end_user_spend_update_queue.add_update = AsyncMock() + + await writer.add_spend_log_transaction_to_daily_end_user_transaction( + payload=payload, + prisma_client=mock_prisma, + ) + + writer.daily_end_user_spend_update_queue.add_update.assert_called_once() + + call_args = writer.daily_end_user_spend_update_queue.add_update.call_args[1] + update_dict = call_args["update"] + assert len(update_dict) == 1 + for key, transaction in update_dict.items(): + assert key == f"{end_user_id}_2024-01-01_test-key_gpt-4_openai" + assert transaction["end_user_id"] == end_user_id + assert transaction["date"] == "2024-01-01" + assert transaction["api_key"] == "test-key" + assert transaction["model"] == "gpt-4" + assert transaction["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_add_spend_log_transaction_to_daily_end_user_transaction_skips_when_end_user_id_missing(): + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + mock_prisma.get_request_status = MagicMock(return_value="success") + + payload = { + "request_id": "req-2", + "user": "test-user", + "startTime": "2024-01-01T12:00:00", + "api_key": "test-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "model_group": "gpt-4-group", + "prompt_tokens": 10, + "completion_tokens": 5, + "spend": 0.2, + "metadata": '{"usage_object": {}}', + } + + writer.daily_end_user_spend_update_queue.add_update = AsyncMock() + + await writer.add_spend_log_transaction_to_daily_end_user_transaction( + payload=payload, + prisma_client=mock_prisma, + ) + + writer.daily_end_user_spend_update_queue.add_update.assert_not_called() \ No newline at end of file diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 86a6ceec25..25ff6f8942 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -1,4 +1,4 @@ -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import FastAPI, HTTPException, Request, status @@ -301,3 +301,99 @@ def test_customer_endpoints_error_schema_consistency(mock_prisma_client, mock_us for key in ["message", "type", "code"]: assert isinstance(error1[key], str), f"error1[{key}] should be a string" assert isinstance(error2[key], str), f"error2[{key}] should be a string" + + +@pytest.mark.asyncio +async def test_get_customer_daily_activity_admin_param_passing(monkeypatch): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints import customer_endpoints + from litellm.proxy.management_endpoints.customer_endpoints import ( + get_customer_daily_activity, + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock( + return_value=[] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse") + get_daily_activity_mock = AsyncMock(return_value=mocked_response) + monkeypatch.setattr( + customer_endpoints, "get_daily_activity", get_daily_activity_mock + ) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1") + result = await get_customer_daily_activity( + end_user_ids="end-user-1,end-user-2", + start_date="2024-01-01", + end_date="2024-01-31", + model="gpt-4", + api_key="test-key", + page=2, + page_size=5, + exclude_end_user_ids="end-user-3", + user_api_key_dict=auth, + ) + + get_daily_activity_mock.assert_awaited_once() + kwargs = get_daily_activity_mock.call_args.kwargs + assert kwargs["table_name"] == "litellm_dailyenduserspend" + assert kwargs["entity_id_field"] == "end_user_id" + assert kwargs["entity_id"] == ["end-user-1", "end-user-2"] + assert kwargs["exclude_entity_ids"] == ["end-user-3"] + assert kwargs["start_date"] == "2024-01-01" + assert kwargs["end_date"] == "2024-01-31" + assert kwargs["model"] == "gpt-4" + assert kwargs["api_key"] == "test-key" + assert kwargs["page"] == 2 + assert kwargs["page_size"] == 5 + + assert result is mocked_response + + +@pytest.mark.asyncio +async def test_get_customer_daily_activity_with_end_user_aliases(monkeypatch): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints import customer_endpoints + from litellm.proxy.management_endpoints.customer_endpoints import ( + get_customer_daily_activity, + ) + + mock_prisma_client = AsyncMock() + mock_end_user1 = MagicMock() + mock_end_user1.user_id = "end-user-1" + mock_end_user1.alias = "Customer One" + mock_end_user2 = MagicMock() + mock_end_user2.user_id = "end-user-2" + mock_end_user2.alias = "Customer Two" + + mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock( + return_value=[mock_end_user1, mock_end_user2] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse") + get_daily_activity_mock = AsyncMock(return_value=mocked_response) + monkeypatch.setattr( + customer_endpoints, "get_daily_activity", get_daily_activity_mock + ) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1") + await get_customer_daily_activity( + end_user_ids="end-user-1,end-user-2", + start_date="2024-01-01", + end_date="2024-01-31", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_end_user_ids=None, + user_api_key_dict=auth, + ) + + kwargs = get_daily_activity_mock.call_args.kwargs + assert kwargs["entity_metadata_field"] == { + "end-user-1": {"alias": "Customer One"}, + "end-user-2": {"alias": "Customer Two"}, + }