diff --git a/tests/proxy_admin_ui_tests/test_usage_endpoints.py b/tests/proxy_admin_ui_tests/test_usage_endpoints.py index b77d74fbf..9918015a1 100644 --- a/tests/proxy_admin_ui_tests/test_usage_endpoints.py +++ b/tests/proxy_admin_ui_tests/test_usage_endpoints.py @@ -73,6 +73,8 @@ from litellm.proxy.proxy_server import ( from litellm.proxy.spend_tracking.spend_management_endpoints import ( global_spend, global_spend_logs, + global_spend_models, + global_spend_keys, spend_key_fn, spend_user_fn, view_spend_logs, @@ -171,3 +173,149 @@ async def test_view_daily_spend_ui(prisma_client): assert ( admin_total_spend > internal_user_total_spend ), "Admin should have more spend than internal user" + + +@pytest.mark.asyncio +async def test_global_spend_models(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() + + # Test for admin user + models_spend_for_admin = await global_spend_models( + limit=10, + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-1234", + user_role=LitellmUserRoles.PROXY_ADMIN, + ), + ) + + print("models_spend_for_admin=", models_spend_for_admin) + + # Test for internal user + models_spend_for_internal_user = await global_spend_models( + limit=10, + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-1234", user_role=LitellmUserRoles.INTERNAL_USER, user_id="1234" + ), + ) + + print("models_spend_for_internal_user=", models_spend_for_internal_user) + + # Assertions + assert isinstance(models_spend_for_admin, list), "Admin response should be a list" + assert isinstance( + models_spend_for_internal_user, list + ), "Internal user response should be a list" + + # Check if the response has the expected shape for both admin and internal user + expected_keys = ["model", "total_spend"] + + if len(models_spend_for_admin) > 0: + assert all( + key in models_spend_for_admin[0] for key in expected_keys + ), f"Admin response should contain keys: {expected_keys}" + assert isinstance( + models_spend_for_admin[0]["model"], str + ), "Model should be a string" + assert isinstance( + models_spend_for_admin[0]["total_spend"], (int, float) + ), "Total spend should be a number" + + if len(models_spend_for_internal_user) > 0: + assert all( + key in models_spend_for_internal_user[0] for key in expected_keys + ), f"Internal user response should contain keys: {expected_keys}" + assert isinstance( + models_spend_for_internal_user[0]["model"], str + ), "Model should be a string" + assert isinstance( + models_spend_for_internal_user[0]["total_spend"], (int, float) + ), "Total spend should be a number" + + # Check if the lists are sorted by total_spend in descending order + if len(models_spend_for_admin) > 1: + assert all( + models_spend_for_admin[i]["total_spend"] + >= models_spend_for_admin[i + 1]["total_spend"] + for i in range(len(models_spend_for_admin) - 1) + ), "Admin response should be sorted by total_spend in descending order" + + if len(models_spend_for_internal_user) > 1: + assert all( + models_spend_for_internal_user[i]["total_spend"] + >= models_spend_for_internal_user[i + 1]["total_spend"] + for i in range(len(models_spend_for_internal_user) - 1) + ), "Internal user response should be sorted by total_spend in descending order" + + # Check if admin has access to more or equal models compared to internal user + assert len(models_spend_for_admin) >= len( + models_spend_for_internal_user + ), "Admin should have access to at least as many models as internal user" + + # Check if the response contains expected fields + if len(models_spend_for_admin) > 0: + assert all( + key in models_spend_for_admin[0] for key in ["model", "total_spend"] + ), "Admin response should contain model, total_spend, and total_tokens" + + if len(models_spend_for_internal_user) > 0: + assert all( + key in models_spend_for_internal_user[0] for key in ["model", "total_spend"] + ), "Internal user response should contain model, total_spend, and total_tokens" + + +@pytest.mark.asyncio +async def test_global_spend_keys(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() + + # Test for admin user + keys_spend_for_admin = await global_spend_keys( + limit=10, + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-1234", + user_role=LitellmUserRoles.PROXY_ADMIN, + ), + ) + + print("keys_spend_for_admin=", keys_spend_for_admin) + + # Test for internal user + keys_spend_for_internal_user = await global_spend_keys( + limit=10, + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-1234", user_role=LitellmUserRoles.INTERNAL_USER, user_id="1234" + ), + ) + + print("keys_spend_for_internal_user=", keys_spend_for_internal_user) + + # Assertions + assert isinstance(keys_spend_for_admin, list), "Admin response should be a list" + assert isinstance( + keys_spend_for_internal_user, list + ), "Internal user response should be a list" + + # Check if admin has access to more or equal keys compared to internal user + assert len(keys_spend_for_admin) >= len( + keys_spend_for_internal_user + ), "Admin should have access to at least as many keys as internal user" + + # Check if the response contains expected fields + if len(keys_spend_for_admin) > 0: + assert all( + key in keys_spend_for_admin[0] + for key in ["api_key", "total_spend", "key_alias", "key_name"] + ), "Admin response should contain api_key, total_spend, key_alias, and key_name" + + if len(keys_spend_for_internal_user) > 0: + assert all( + key in keys_spend_for_internal_user[0] + for key in ["api_key", "total_spend", "key_alias", "key_name"] + ), "Internal user response should contain api_key, total_spend, key_alias, and key_name"