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 e5d805f04c7..97d1d4a2f38 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 @@ -986,39 +986,45 @@ async def test_ui_view_spend_logs_with_status(client, monkeypatch): start_date, end_date = _default_date_range() - # Test success status - response = client.get( - "/spend/logs/ui", - params={ - "status_filter": "success", - "start_date": start_date, - "end_date": end_date, - }, - headers={"Authorization": "Bearer sk-test"}, + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN ) + try: + # Test success status + response = client.get( + "/spend/logs/ui", + params={ + "status_filter": "success", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) - assert response.status_code == 200 - data = response.json() - assert data["total"] == 1 - assert len(data["data"]) == 1 - assert data["data"][0]["status"] == "success" + assert response.status_code == 200 + data = response.json() + assert data["total"] == 1 + assert len(data["data"]) == 1 + assert data["data"][0]["status"] == "success" - # Test failure status - response = client.get( - "/spend/logs/ui", - params={ - "status_filter": "failure", - "start_date": start_date, - "end_date": end_date, - }, - headers={"Authorization": "Bearer sk-test"}, - ) + # Test failure status + response = client.get( + "/spend/logs/ui", + params={ + "status_filter": "failure", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) - assert response.status_code == 200 - data = response.json() - assert data["total"] == 1 - assert len(data["data"]) == 1 - assert data["data"][0]["status"] == "failure" + assert response.status_code == 200 + data = response.json() + assert data["total"] == 1 + assert len(data["data"]) == 1 + assert data["data"][0]["status"] == "failure" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) @pytest.mark.asyncio @@ -1060,25 +1066,31 @@ async def test_ui_view_spend_logs_with_model(client, monkeypatch): start_date, end_date = _default_date_range() - # Make the request with model filter - response = client.get( - "/spend/logs/ui", - params={ - "model": "gpt-3.5-turbo", - "start_date": start_date, - "end_date": end_date, - }, - headers={"Authorization": "Bearer sk-test"}, + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN ) + try: + # Make the request with model filter + response = client.get( + "/spend/logs/ui", + params={ + "model": "gpt-3.5-turbo", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) - # Assert response - assert response.status_code == 200 - data = response.json() + # 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]["model"] == "gpt-3.5-turbo" + # Verify the filtered data + assert data["total"] == 1 + assert len(data["data"]) == 1 + assert data["data"][0]["model"] == "gpt-3.5-turbo" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) @pytest.mark.asyncio @@ -1123,21 +1135,27 @@ async def test_ui_view_spend_logs_with_model_id(client, monkeypatch): start_date, end_date = _default_date_range() - response = client.get( - "/spend/logs/ui", - params={ - "model_id": "deployment-id-1", - "start_date": start_date, - "end_date": end_date, - }, - headers={"Authorization": "Bearer sk-test"}, + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN ) + try: + response = client.get( + "/spend/logs/ui", + params={ + "model_id": "deployment-id-1", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) - assert response.status_code == 200 - data = response.json() - assert data["total"] == 1 - assert len(data["data"]) == 1 - assert data["data"][0]["model_id"] == "deployment-id-1" + assert response.status_code == 200 + data = response.json() + assert data["total"] == 1 + assert len(data["data"]) == 1 + assert data["data"][0]["model_id"] == "deployment-id-1" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) @pytest.mark.asyncio @@ -1682,69 +1700,75 @@ async def test_view_spend_logs_summarize_parameter(client, monkeypatch): ) end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d") - # Test 1: summarize=false should return individual log entries - response = client.get( - "/spend/logs", - params={ - "start_date": start_date, - "end_date": end_date, - "summarize": "false", - }, - headers={"Authorization": "Bearer sk-test"}, + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN ) + try: + # Test 1: summarize=false should return individual log entries + response = client.get( + "/spend/logs", + params={ + "start_date": start_date, + "end_date": end_date, + "summarize": "false", + }, + headers={"Authorization": "Bearer sk-test"}, + ) - assert response.status_code == 200 - data = response.json() + assert response.status_code == 200 + data = response.json() - # Should return the raw log entries - assert isinstance(data, list) - assert len(data) == 2 - assert data[0]["id"] == "log1" - assert data[1]["id"] == "log2" - assert data[0]["request_id"] == "req1" - assert data[1]["request_id"] == "req2" + # Should return the raw log entries + assert isinstance(data, list) + assert len(data) == 2 + assert data[0]["id"] == "log1" + assert data[1]["id"] == "log2" + assert data[0]["request_id"] == "req1" + assert data[1]["request_id"] == "req2" - # Test 2: summarize=true should return grouped data - response = client.get( - "/spend/logs", - params={ - "start_date": start_date, - "end_date": end_date, - "summarize": "true", - }, - headers={"Authorization": "Bearer sk-test"}, - ) + # Test 2: summarize=true should return grouped data + response = client.get( + "/spend/logs", + params={ + "start_date": start_date, + "end_date": end_date, + "summarize": "true", + }, + headers={"Authorization": "Bearer sk-test"}, + ) - assert response.status_code == 200 - data = response.json() + assert response.status_code == 200 + data = response.json() - # Should return grouped/summarized data - assert isinstance(data, list) - # The structure should be different - grouped by date with aggregated spend - assert "startTime" in data[0] - assert "spend" in data[0] - assert "users" in data[0] - assert "models" in data[0] + # Should return grouped/summarized data + assert isinstance(data, list) + # The structure should be different - grouped by date with aggregated spend + assert "startTime" in data[0] + assert "spend" in data[0] + assert "users" in data[0] + assert "models" in data[0] - # Test 3: default behavior (no summarize parameter) should maintain backward compatibility - response = client.get( - "/spend/logs", - params={ - "start_date": start_date, - "end_date": end_date, - }, - headers={"Authorization": "Bearer sk-test"}, - ) + # Test 3: default behavior (no summarize parameter) should maintain backward compatibility + response = client.get( + "/spend/logs", + params={ + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) - assert response.status_code == 200 - data = response.json() + assert response.status_code == 200 + data = response.json() - # Should return grouped/summarized data (same as summarize=true) - assert isinstance(data, list) - assert "startTime" in data[0] - assert "spend" in data[0] - assert "users" in data[0] - assert "models" in data[0] + # Should return grouped/summarized data (same as summarize=true) + assert isinstance(data, list) + assert "startTime" in data[0] + assert "spend" in data[0] + assert "users" in data[0] + assert "models" in data[0] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) @pytest.mark.asyncio