diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 7acca1804f7..65b245b9e7e 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -15,15 +15,19 @@ sys.path.insert( from unittest.mock import MagicMock, patch import litellm -from litellm.proxy._types import SpendLogsPayload +import litellm.proxy.proxy_server as ps +from litellm.proxy._types import ( + LitellmUserRoles, + Member, + SpendLogsPayload, + UserAPIKeyAuth, +) from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger from litellm.proxy.proxy_server import app, prisma_client from litellm.proxy.spend_tracking import spend_management_endpoints from litellm.router import Router from litellm.types.utils import BudgetConfig -from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles, Member -from litellm.proxy.spend_tracking import spend_management_endpoints -import litellm.proxy.proxy_server as ps + @pytest.mark.asyncio async def test_is_admin_view_safe_true(monkeypatch): @@ -386,36 +390,50 @@ async def test_ui_view_spend_logs_with_team_id(client, monkeypatch): def __init__(self): self.db = MockDB() self.db.litellm_spendlogs = self.db - + # Apply the monkeypatch mock_prisma_client = MockPrismaClient() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - # Set up test dates - start_date = ( - datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7) - ).strftime("%Y-%m-%d %H:%M:%S") - end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") - - # Make the request with team_id filter - response = client.get( - "/spend/logs/ui", - params={ - "team_id": "team1", - "start_date": start_date, - "end_date": end_date, - }, - headers={"Authorization": "Bearer sk-test"}, + + # Mock _is_admin_view_safe to return True to bypass permission checks + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True ) - # Assert response - assert response.status_code == 200 - data = response.json() + # Override auth dependency to return PROXY_ADMIN + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) - # Verify the filtered data - assert data["total"] == 1 - assert len(data["data"]) == 1 - assert data["data"][0]["team_id"] == "team1" + try: + # Set up test dates + start_date = ( + datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7) + ).strftime("%Y-%m-%d %H:%M:%S") + end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + + # Make the request with team_id filter + response = client.get( + "/spend/logs/ui", + params={ + "team_id": "team1", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + + # Assert response + assert response.status_code == 200 + data = response.json() + + # Verify the filtered data + assert data["total"] == 1 + assert len(data["data"]) == 1 + assert data["data"][0]["team_id"] == "team1" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) @pytest.mark.asyncio