diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx index c2837cf412e..627a6e35c47 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx @@ -3,7 +3,6 @@ import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { KeyResponse } from "../../../key_team_helpers/key_list"; -import * as transformKeyInfo from "../../../key_team_helpers/transform_key_info"; import * as networking from "../../../networking"; import TopKeyView from "./TopKeyView"; @@ -13,11 +12,7 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ })); vi.mock("../../../networking", () => ({ - keyInfoV1Call: vi.fn(), -})); - -vi.mock("../../../key_team_helpers/transform_key_info", () => ({ - transformKeyInfo: vi.fn(), + keyListCall: vi.fn(), })); vi.mock("../../../templates/key_info_view", () => ({ @@ -31,8 +26,7 @@ vi.mock("../../../templates/key_info_view", () => ({ describe("TopKeyView", () => { const mockUseAuthorized = vi.mocked(useAuthorized); - const mockKeyInfoV1Call = vi.mocked(networking.keyInfoV1Call); - const mockTransformKeyInfo = vi.mocked(transformKeyInfo.transformKeyInfo); + const mockKeyListCall = vi.mocked(networking.keyListCall); const mockAuth = { token: "mock-token", @@ -58,8 +52,7 @@ describe("TopKeyView", () => { beforeEach(() => { mockUseAuthorized.mockReturnValue(mockAuth); mockSetTopKeysLimit.mockClear(); - mockKeyInfoV1Call.mockClear(); - mockTransformKeyInfo.mockClear(); + mockKeyListCall.mockClear(); }); it("should render", () => { @@ -129,10 +122,8 @@ describe("TopKeyView", () => { }); it("renders cyan bars with truncated aliases in chart view and opens the key info modal on bar click", async () => { - const mockKeyInfo = { key: "info" }; const mockTransformedData = { transformed: "data" } as unknown as KeyResponse; - mockKeyInfoV1Call.mockResolvedValue(mockKeyInfo); - mockTransformKeyInfo.mockReturnValue(mockTransformedData); + mockKeyListCall.mockResolvedValue({ keys: [mockTransformedData] }); const user = userEvent.setup(); const { container } = render( @@ -158,7 +149,19 @@ describe("TopKeyView", () => { fireEvent.click(bars[0]); await waitFor(() => { - expect(mockKeyInfoV1Call).toHaveBeenCalledWith("test-token", "key-123"); + expect(mockKeyListCall).toHaveBeenCalledWith( + "test-token", + null, + null, + null, + null, + "key-123", + 1, + 1, + null, + null, + "user", + ); }); await waitFor(() => { @@ -391,10 +394,8 @@ describe("TopKeyView", () => { }); it("should open modal when key ID is clicked", async () => { - const mockKeyInfo = { key: "info" }; const mockTransformedData = { transformed: "data" } as unknown as KeyResponse; - mockKeyInfoV1Call.mockResolvedValue(mockKeyInfo); - mockTransformKeyInfo.mockReturnValue(mockTransformedData); + mockKeyListCall.mockResolvedValue({ keys: [mockTransformedData] }); const user = userEvent.setup(); render( @@ -416,7 +417,19 @@ describe("TopKeyView", () => { } await waitFor(() => { - expect(mockKeyInfoV1Call).toHaveBeenCalledWith("test-token", "key-123"); + expect(mockKeyListCall).toHaveBeenCalledWith( + "test-token", + null, + null, + null, + null, + "key-123", + 1, + 1, + null, + null, + "user", + ); }); await waitFor(() => { @@ -425,10 +438,8 @@ describe("TopKeyView", () => { }); it("should close modal when close button is clicked", async () => { - const mockKeyInfo = { key: "info" }; const mockTransformedData = { transformed: "data" } as unknown as KeyResponse; - mockKeyInfoV1Call.mockResolvedValue(mockKeyInfo); - mockTransformKeyInfo.mockReturnValue(mockTransformedData); + mockKeyListCall.mockResolvedValue({ keys: [mockTransformedData] }); const user = userEvent.setup(); render( @@ -462,10 +473,8 @@ describe("TopKeyView", () => { }); it("should close modal when escape key is pressed", async () => { - const mockKeyInfo = { key: "info" }; const mockTransformedData = { transformed: "data" } as unknown as KeyResponse; - mockKeyInfoV1Call.mockResolvedValue(mockKeyInfo); - mockTransformKeyInfo.mockReturnValue(mockTransformedData); + mockKeyListCall.mockResolvedValue({ keys: [mockTransformedData] }); const user = userEvent.setup(); render( @@ -498,10 +507,8 @@ describe("TopKeyView", () => { }); it("should close modal when clicking outside modal", async () => { - const mockKeyInfo = { key: "info" }; const mockTransformedData = { transformed: "data" } as unknown as KeyResponse; - mockKeyInfoV1Call.mockResolvedValue(mockKeyInfo); - mockTransformKeyInfo.mockReturnValue(mockTransformedData); + mockKeyListCall.mockResolvedValue({ keys: [mockTransformedData] }); const user = userEvent.setup(); const { container } = render( @@ -562,7 +569,7 @@ describe("TopKeyView", () => { } await waitFor(() => { - expect(mockKeyInfoV1Call).not.toHaveBeenCalled(); + expect(mockKeyListCall).not.toHaveBeenCalled(); }); expect(screen.queryByTestId("key-info-view")).not.toBeInTheDocument(); @@ -570,7 +577,7 @@ describe("TopKeyView", () => { it("should handle error when fetching key info", async () => { const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); - mockKeyInfoV1Call.mockRejectedValue(new Error("Network error")); + mockKeyListCall.mockRejectedValue(new Error("Network error")); const user = userEvent.setup(); render( @@ -592,7 +599,7 @@ describe("TopKeyView", () => { } await waitFor(() => { - expect(mockKeyInfoV1Call).toHaveBeenCalled(); + expect(mockKeyListCall).toHaveBeenCalled(); }); await waitFor(() => { diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx index df1a51d8e38..5888295c4c4 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx @@ -8,8 +8,8 @@ import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import { SimpleTooltip } from "@/components/ui/tooltip"; import React, { useState } from "react"; import { formatNumberWithCommas } from "../../../../utils/dataUtils"; -import { transformKeyInfo } from "../../../key_team_helpers/transform_key_info"; -import { keyInfoV1Call } from "../../../networking"; +import type { KeyResponse } from "../../../key_team_helpers/key_list"; +import { keyListCall } from "../../../networking"; import KeyInfoView from "../../../templates/key_info_view"; import { TagUsage } from "../../types"; @@ -47,10 +47,14 @@ const TopKeyView: React.FC = ({ topKeys, teams, showTags = fals if (!accessToken) return; try { - const keyInfo = await keyInfoV1Call(accessToken, item.api_key); - const transformedKeyData = transformKeyInfo(keyInfo); + const keyList = await keyListCall(accessToken, null, null, null, null, item.api_key, 1, 1, null, null, "user"); + const expandedKey = keyList.keys[0] as KeyResponse | undefined; + if (!expandedKey) return; - setKeyData(transformedKeyData); + setKeyData({ + ...expandedKey, + user_email: expandedKey.user?.user_email ?? expandedKey.user_email, + }); setSelectedKey(item.api_key); setIsModalOpen(true); // Open modal when key is clicked } catch (error) { diff --git a/ui/litellm-dashboard/tests/top_key_view.test.tsx b/ui/litellm-dashboard/tests/top_key_view.test.tsx index 51662b8f453..a49736559bd 100644 --- a/ui/litellm-dashboard/tests/top_key_view.test.tsx +++ b/ui/litellm-dashboard/tests/top_key_view.test.tsx @@ -1,17 +1,19 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; -import { renderWithProviders, screen, fireEvent } from "./test-utils"; +import { renderWithProviders, screen, fireEvent, waitFor } from "./test-utils"; import TopKeyView from "../src/components/UsagePage/components/EntityUsage/TopKeyView"; import { TagUsage } from "../src/components/UsagePage/types"; import useAuthorized from "../src/app/(dashboard)/hooks/useAuthorized"; +import { keyListCall } from "../src/components/networking"; // Mock the networking module vi.mock("../src/components/networking", () => ({ - keyInfoV1Call: vi.fn(), + keyListCall: vi.fn(), })); -// Mock the transform function -vi.mock("../src/components/key_team_helpers/transform_key_info", () => ({ - transformKeyInfo: vi.fn((data) => data), +vi.mock("../src/components/templates/key_info_view", () => ({ + default: ({ keyData }: { keyData: { user_id?: string; user_email?: string; user?: { user_alias?: string } } }) => ( +
{`Owner: ${keyData.user?.user_alias}|${keyData.user_email}|${keyData.user_id}`}
+ ), })); vi.mock("../src/app/(dashboard)/hooks/useAuthorized", () => ({ @@ -20,6 +22,7 @@ vi.mock("../src/app/(dashboard)/hooks/useAuthorized", () => ({ describe("TopKeyView", () => { const mockUseAuthorized = vi.mocked(useAuthorized); + const mockKeyListCall = vi.mocked(keyListCall); const mockProps = { topKeys: [], accessToken: "test-token", @@ -305,5 +308,50 @@ describe("TopKeyView", () => { // Check that spend is formatted correctly expect(screen.getByText("$25.50")).toBeInTheDocument(); }); + + it("should expand the owner when opening key details", async () => { + mockKeyListCall.mockResolvedValue({ + keys: [ + { + token: "key-1", + user_id: "user-1", + user_email: null, + user: { + user_id: "user-1", + user_alias: "Alice Example", + user_email: "alice@example.com", + }, + }, + ], + } as Awaited>); + + renderWithProviders( + , + ); + + fireEvent.click(screen.getByText("key-1")); + + await waitFor(() => + expect(mockKeyListCall).toHaveBeenCalledWith( + "test-token", + null, + null, + null, + null, + "key-1", + 1, + 1, + null, + null, + "user", + ), + ); + expect(await screen.findByText("Owner: Alice Example|alice@example.com|user-1")).toBeInTheDocument(); + }); }); });