diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index c62d41cca..96dfeed0f 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -42,6 +42,8 @@ from litellm.proxy.proxy_server import ( info_key_fn, update_key_fn, generate_key_fn, + spend_user_fn, + spend_key_fn, ) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm._logging import verbose_proxy_logger @@ -830,3 +832,39 @@ def test_call_with_key_over_budget_stream(prisma_client): error_detail = e.message assert "Authentication Error, ExceededTokenBudget:" in error_detail print(vars(e)) + + +@pytest.mark.asyncio() +async def test_view_spend_per_user(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() + try: + user_by_spend = await spend_user_fn() + assert type(user_by_spend) == list + assert len(user_by_spend) > 0 + first_user = user_by_spend[0] + + print("\nfirst_user=", first_user) + assert first_user.spend > 0 + except Exception as e: + print("Got Exception", e) + pytest.fail(f"Got exception {e}") + + +@pytest.mark.asyncio() +async def test_view_spend_per_key(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() + try: + key_by_spend = await spend_key_fn() + assert type(key_by_spend) == list + assert len(key_by_spend) > 0 + first_key = key_by_spend[0] + + print("\nfirst_key=", first_key) + assert first_key.spend > 0 + except Exception as e: + print("Got Exception", e) + pytest.fail(f"Got exception {e}")