diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.test.ts new file mode 100644 index 00000000000..a392a940f98 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.test.ts @@ -0,0 +1,253 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useCurrentUser } from "./useCurrentUser"; +import { userInfoCall } from "@/components/networking"; +import type { UserInfo } from "@/components/view_users/types"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + userInfoCall: vi.fn(), +})); + +// Mock the queryKeysFactory - we'll mock the specific return value +vi.mock("../common/queryKeysFactory", () => ({ + createQueryKeys: vi.fn((resource: string) => ({ + all: [resource], + lists: () => [resource, "list"], + list: (params?: any) => [resource, "list", { params }], + details: () => [resource, "detail"], + detail: (uid: string) => [resource, "detail", uid], + })), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock data - response from userInfoCall should have user_info property +const mockUserInfoResponse = { + user_info: { + user_id: "test-user-id", + user_email: "test@example.com", + user_alias: "Test User", + user_role: "Admin", + spend: 150.75, + max_budget: 1000.0, + key_count: 5, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + sso_user_id: null, + budget_duration: "monthly", + } as UserInfo, +}; + +describe("useCurrentUser", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return user info data when query is successful", async () => { + // Mock successful API call + (userInfoCall as any).mockResolvedValue(mockUserInfoResponse); + + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockUserInfoResponse.user_info); + expect(result.current.error).toBeNull(); + expect(userInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", false, null, null); + expect(userInfoCall).toHaveBeenCalledTimes(1); + }); + + it("should handle error when userInfoCall fails", async () => { + const errorMessage = "Failed to fetch user info"; + const testError = new Error(errorMessage); + + // Mock failed API call + (userInfoCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(userInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", false, null, null); + expect(userInfoCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: "test-user-id", + userRole: "Admin", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(userInfoCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userId is missing", async () => { + // Mock missing userId + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: null, + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(userInfoCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is missing", async () => { + // Mock missing userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: null, + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(userInfoCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when all auth values are missing", async () => { + // Mock all auth values missing + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: null, + userRole: null, + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(userInfoCall).not.toHaveBeenCalled(); + }); + + it("should execute query when all auth values are present", async () => { + // Mock successful API call + (userInfoCall as any).mockResolvedValue(mockUserInfoResponse); + + // Ensure all auth values are present (already set in beforeEach) + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(userInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", false, null, null); + expect(userInfoCall).toHaveBeenCalledTimes(1); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (userInfoCall as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.ts new file mode 100644 index 00000000000..f4028ada0dc --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.ts @@ -0,0 +1,19 @@ +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { UserInfo, userInfoCall } from "@/components/networking"; +import { useQuery, UseQueryResult } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +const userKeys = createQueryKeys("users"); + +export const useCurrentUser = (): UseQueryResult => { + const { accessToken, userId, userRole } = useAuthorized(); + return useQuery({ + queryKey: userKeys.detail(userId!), + queryFn: async () => { + const data = await userInfoCall(accessToken!, userId!, userRole!, false, null, null); + console.log(`userInfo: ${JSON.stringify(data)}`); + return data.user_info; + }, + enabled: Boolean(accessToken && userId && userRole), + }); +}; diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx index d72515b8e4f..09415e87443 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx @@ -1,8 +1,9 @@ import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents"; import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { act, fireEvent, screen, waitFor } from "@testing-library/react"; import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../tests/test-utils"; import type { Organization } from "../../networking"; import * as networking from "../../networking"; import NewUsagePage from "./UsagePageView"; @@ -332,7 +333,7 @@ describe("NewUsage", () => { }); it("should render and fetch usage data on mount", async () => { - render(); + renderWithProviders(); // Wait for data to be fetched await waitFor(() => { @@ -349,7 +350,7 @@ describe("NewUsage", () => { }); it("should display usage metrics and charts", async () => { - render(); + renderWithProviders(); await waitFor(() => { expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); @@ -367,7 +368,7 @@ describe("NewUsage", () => { }); it("should switch between usage views correctly", async () => { - render(); + renderWithProviders(); await waitFor(() => { expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); @@ -401,7 +402,7 @@ describe("NewUsage", () => { }); it("should show organization usage banner and view for admins", async () => { - render(); + renderWithProviders(); await waitFor(() => { expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); @@ -426,7 +427,7 @@ describe("NewUsage", () => { error: null, } as any); - render(); + renderWithProviders(); await waitFor(() => { expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); @@ -450,7 +451,7 @@ describe("NewUsage", () => { error: null, } as any); - render(); + renderWithProviders(); await waitFor(() => { expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx index a7107814113..a9c86c064ef 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx @@ -52,6 +52,7 @@ import { valueFormatterSpend } from "../utils/value_formatters"; import EntityUsage, { EntityList } from "./EntityUsage/EntityUsage"; import TopKeyView from "./EntityUsage/TopKeyView"; import { UsageOption, UsageViewSelect } from "./UsageViewSelect/UsageViewSelect"; +import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; interface UsagePageProps { teams: Team[]; @@ -82,6 +83,9 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const [allTags, setAllTags] = useState([]); const { data: customers = [] } = useCustomers(); const { data: agentsResponse } = useAgents(); + const { data: currentUser } = useCurrentUser(); + console.log(`currentUser: ${JSON.stringify(currentUser)}`); + console.log(`currentUser max budget: ${currentUser?.max_budget}`); const [modelViewType, setModelViewType] = useState<"groups" | "individual">("groups"); const [isCloudZeroModalOpen, setIsCloudZeroModalOpen] = useState(false); const [isGlobalExportModalOpen, setIsGlobalExportModalOpen] = useState(false); @@ -477,7 +481,11 @@ const UsagePage: React.FC = ({ teams, organizations }) => { )} - + diff --git a/ui/litellm-dashboard/src/components/view_user_spend.tsx b/ui/litellm-dashboard/src/components/view_user_spend.tsx index 51e87f8bc0d..c05850361a3 100644 --- a/ui/litellm-dashboard/src/components/view_user_spend.tsx +++ b/ui/litellm-dashboard/src/components/view_user_spend.tsx @@ -17,7 +17,6 @@ interface ViewUserSpendProps { } const ViewUserSpend: React.FC = ({ userSpend, userMaxBudget, selectedTeam }) => { const { accessToken, userRole, userId: userID } = useAuthorized(); - console.log(`userSpend: ${userSpend}`); let [spend, setSpend] = useState(userSpend !== null ? userSpend : 0.0); const [maxBudget, setMaxBudget] = useState( selectedTeam ? Number(formatNumberWithCommas(selectedTeam.max_budget, 4)) : null, @@ -61,6 +60,8 @@ const ViewUserSpend: React.FC = ({ userSpend, userMaxBudget, setMaxBudget(selectedTeam.max_budget); } } + } else { + setMaxBudget(userMaxBudget); } }, [selectedTeam, userMaxBudget]); const [userModels, setUserModels] = useState([]);