diff --git a/tests/proxy_admin_ui_tests/test_usage_endpoints.py b/tests/proxy_admin_ui_tests/test_usage_endpoints.py index cb5a9f7ad..b77d74fbf 100644 --- a/tests/proxy_admin_ui_tests/test_usage_endpoints.py +++ b/tests/proxy_admin_ui_tests/test_usage_endpoints.py @@ -12,3 +12,162 @@ For all tests - test the following: - Response is valid - Response for Admin User is different from response from Internal User """ + +import os +import sys +import traceback +import uuid +from datetime import datetime + +from dotenv import load_dotenv +from fastapi import Request +from fastapi.routing import APIRoute + +load_dotenv() +import io +import os +import time + +# this file is to test litellm/proxy + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import asyncio +import logging + +import pytest + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.proxy.management_endpoints.internal_user_endpoints import ( + new_user, + user_info, + user_update, +) +from litellm.proxy.management_endpoints.key_management_endpoints import ( + delete_key_fn, + generate_key_fn, + generate_key_helper_fn, + info_key_fn, + regenerate_key_fn, + update_key_fn, +) +from litellm.proxy.management_endpoints.team_endpoints import ( + new_team, + team_info, + update_team, +) +from litellm.proxy.proxy_server import ( + LitellmUserRoles, + audio_transcriptions, + chat_completion, + completion, + embeddings, + image_generation, + model_list, + moderations, + new_end_user, + user_api_key_auth, +) +from litellm.proxy.spend_tracking.spend_management_endpoints import ( + global_spend, + global_spend_logs, + spend_key_fn, + spend_user_fn, + view_spend_logs, +) +from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token, update_spend + +verbose_proxy_logger.setLevel(level=logging.DEBUG) + +from starlette.datastructures import URL + +from litellm.caching import DualCache +from litellm.proxy._types import ( + DynamoDBArgs, + GenerateKeyRequest, + KeyRequest, + LiteLLM_UpperboundKeyGenerateParams, + NewCustomerRequest, + NewTeamRequest, + NewUserRequest, + ProxyErrorTypes, + ProxyException, + UpdateKeyRequest, + UpdateTeamRequest, + UpdateUserRequest, + UserAPIKeyAuth, +) +from litellm.proxy.utils import DBClient + +proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + + +@pytest.fixture +def prisma_client(): + from litellm.proxy.proxy_cli import append_query_params + + ### add connection pool + pool timeout args + params = {"connection_limit": 100, "pool_timeout": 60} + database_url = os.getenv("DATABASE_URL") + modified_url = append_query_params(database_url, params) + os.environ["DATABASE_URL"] = modified_url + + # Assuming DBClient is a class that needs to be instantiated + prisma_client = PrismaClient( + database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj + ) + + # Reset litellm.proxy.proxy_server.prisma_client to None + litellm.proxy.proxy_server.custom_db_client = None + litellm.proxy.proxy_server.litellm_proxy_budget_name = ( + f"litellm-proxy-budget-{time.time()}" + ) + litellm.proxy.proxy_server.user_custom_key_generate = None + + return prisma_client + + +@pytest.mark.asyncio() +async def test_view_daily_spend_ui(prisma_client): + print("prisma client=", prisma_client) + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + + await litellm.proxy.proxy_server.prisma_client.connect() + from litellm.proxy.proxy_server import user_api_key_cache + + spend_logs_for_admin = await global_spend_logs( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-1234", + user_role=LitellmUserRoles.PROXY_ADMIN, + ), + api_key=None, + ) + + print("spend_logs_for_admin=", spend_logs_for_admin) + + spend_logs_for_internal_user = await global_spend_logs( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-1234", user_role=LitellmUserRoles.INTERNAL_USER, user_id="1234" + ), + api_key=None, + ) + + print("spend_logs_for_internal_user=", spend_logs_for_internal_user) + + # Calculate total spend for admin + admin_total_spend = sum(log.get("spend", 0) for log in spend_logs_for_admin) + + # Calculate total spend for internal user (0 in this case, but we'll keep it generic) + internal_user_total_spend = sum( + log.get("spend", 0) for log in spend_logs_for_internal_user + ) + + print("total_spend_for_admin=", admin_total_spend) + print("total_spend_for_internal_user=", internal_user_total_spend) + + assert ( + admin_total_spend > internal_user_total_spend + ), "Admin should have more spend than internal user"