mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
feat(ui): migrate useCustomers to the typed fetchClient
Converts useCustomers from allEndUsersCall to fetchClient.GET("/customer/list"); the response is typed as LiteLLM_EndUserTable[] from the schema, so the hand-written Customer/CustomersResponse types are deleted. They were also inaccurate (allowed_model_region was string but is "eu"|"us", and a budget_id the table has no field for). No cast; the schema type flows to the one consumer. First caller on the new pattern.
This commit is contained in:
ryan-crabbe-berri
parent
80e5f22e26
commit
069c3ca910
2 changed files with 48 additions and 301 deletions
|
|
@ -1,334 +1,104 @@
|
|||
import { allEndUsersCall } from "@/components/networking";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import { renderHook, waitFor } from "@testing-library/react";
|
||||
import React, { ReactNode } from "react";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import type { Customer, CustomersResponse } from "./useCustomers";
|
||||
import { useCustomers } from "./useCustomers";
|
||||
import { useCustomers, type EndUser } from "./useCustomers";
|
||||
|
||||
// Mock the networking function
|
||||
vi.mock("@/components/networking", () => ({
|
||||
allEndUsersCall: vi.fn(),
|
||||
const mockGet = vi.fn();
|
||||
vi.mock("@/lib/http/api", () => ({
|
||||
fetchClient: { GET: (...args: unknown[]) => mockGet(...args) },
|
||||
}));
|
||||
|
||||
// Mock useAuthorized hook - we can override this in individual tests
|
||||
const mockUseAuthorized = vi.fn();
|
||||
vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
|
||||
default: () => mockUseAuthorized(),
|
||||
}));
|
||||
|
||||
// Import actual roles instead of mocking them
|
||||
|
||||
// Mock data
|
||||
const mockCustomers: Customer[] = [
|
||||
{
|
||||
user_id: "customer-1",
|
||||
alias: "Test Customer 1",
|
||||
spend: 150.5,
|
||||
blocked: false,
|
||||
allowed_model_region: "us-east-1",
|
||||
default_model: "gpt-3.5-turbo",
|
||||
budget_id: "budget-1",
|
||||
litellm_budget_table: {
|
||||
budget_id: "budget-1",
|
||||
max_budget: 1000,
|
||||
soft_budget: 800,
|
||||
max_parallel_requests: 10,
|
||||
tpm_limit: 1000,
|
||||
rpm_limit: 100,
|
||||
model_max_budget: { "gpt-4": 500 },
|
||||
budget_duration: "monthly",
|
||||
budget_reset_at: "2024-02-01T00:00:00Z",
|
||||
created_at: "2024-01-01T00:00:00Z",
|
||||
created_by: "admin-1",
|
||||
updated_at: "2024-01-01T00:00:00Z",
|
||||
updated_by: "admin-1",
|
||||
},
|
||||
},
|
||||
{
|
||||
user_id: "customer-2",
|
||||
alias: null,
|
||||
spend: 0,
|
||||
blocked: true,
|
||||
allowed_model_region: null,
|
||||
default_model: null,
|
||||
budget_id: null,
|
||||
litellm_budget_table: null,
|
||||
},
|
||||
const mockCustomers: EndUser[] = [
|
||||
{ user_id: "customer-1", alias: "Test Customer 1", spend: 150.5, blocked: false },
|
||||
{ user_id: "customer-2", alias: null, spend: 0, blocked: true },
|
||||
];
|
||||
|
||||
const mockCustomersResponse: CustomersResponse = mockCustomers;
|
||||
const authorized = {
|
||||
accessToken: "test-access-token",
|
||||
userRole: "Admin",
|
||||
userId: "test-user-id",
|
||||
token: "test-token",
|
||||
userEmail: "test@example.com",
|
||||
premiumUser: false,
|
||||
disabledPersonalKeyCreation: null,
|
||||
showSSOBanner: false,
|
||||
};
|
||||
|
||||
describe("useCustomers", () => {
|
||||
let queryClient: QueryClient;
|
||||
|
||||
beforeEach(() => {
|
||||
queryClient = new QueryClient({
|
||||
defaultOptions: {
|
||||
queries: {
|
||||
retry: false,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// Reset all mocks
|
||||
queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } });
|
||||
vi.clearAllMocks();
|
||||
|
||||
// Set default mock for useAuthorized (enabled state)
|
||||
mockUseAuthorized.mockReturnValue({
|
||||
accessToken: "test-access-token",
|
||||
userRole: "Admin",
|
||||
userId: "test-user-id",
|
||||
token: "test-token",
|
||||
userEmail: "test@example.com",
|
||||
premiumUser: false,
|
||||
disabledPersonalKeyCreation: null,
|
||||
showSSOBanner: false,
|
||||
});
|
||||
mockUseAuthorized.mockReturnValue(authorized);
|
||||
});
|
||||
|
||||
const wrapper = ({ children }: { children: ReactNode }) =>
|
||||
React.createElement(QueryClientProvider, { client: queryClient }, children);
|
||||
|
||||
it("should return customers data when query is successful", async () => {
|
||||
// Mock successful API call
|
||||
(allEndUsersCall as any).mockResolvedValue(mockCustomersResponse);
|
||||
it("fetches /customer/list and returns the typed list on success", async () => {
|
||||
mockGet.mockResolvedValue({ data: mockCustomers });
|
||||
|
||||
const { result } = renderHook(() => useCustomers(), { 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(mockCustomersResponse);
|
||||
expect(result.current.error).toBeNull();
|
||||
expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token");
|
||||
expect(allEndUsersCall).toHaveBeenCalledTimes(1);
|
||||
expect(result.current.data).toEqual(mockCustomers);
|
||||
expect(mockGet).toHaveBeenCalledWith("/customer/list");
|
||||
expect(mockGet).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("should handle error when allEndUsersCall fails", async () => {
|
||||
const errorMessage = "Failed to fetch customers";
|
||||
const testError = new Error(errorMessage);
|
||||
|
||||
// Mock failed API call
|
||||
(allEndUsersCall as any).mockRejectedValue(testError);
|
||||
it("surfaces an error when the request rejects", async () => {
|
||||
const testError = new Error("Failed to fetch customers");
|
||||
mockGet.mockRejectedValue(testError);
|
||||
|
||||
const { result } = renderHook(() => useCustomers(), { 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(allEndUsersCall).toHaveBeenCalledWith("test-access-token");
|
||||
expect(allEndUsersCall).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("should not execute query when accessToken is missing", async () => {
|
||||
// Mock missing accessToken
|
||||
mockUseAuthorized.mockReturnValue({
|
||||
accessToken: null,
|
||||
userRole: "Admin",
|
||||
userId: "test-user-id",
|
||||
token: null,
|
||||
userEmail: "test@example.com",
|
||||
premiumUser: false,
|
||||
disabledPersonalKeyCreation: null,
|
||||
showSSOBanner: false,
|
||||
});
|
||||
it("falls back to an empty list when the response has no body", async () => {
|
||||
mockGet.mockResolvedValue({ data: undefined });
|
||||
|
||||
const { result } = renderHook(() => useCustomers(), { 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(allEndUsersCall).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("should not execute query when userRole is not an admin role", async () => {
|
||||
// Mock non-admin userRole
|
||||
mockUseAuthorized.mockReturnValue({
|
||||
accessToken: "test-access-token",
|
||||
userRole: "member", // Not in all_admin_roles
|
||||
userId: "test-user-id",
|
||||
token: "test-token",
|
||||
userEmail: "test@example.com",
|
||||
premiumUser: false,
|
||||
disabledPersonalKeyCreation: null,
|
||||
showSSOBanner: false,
|
||||
});
|
||||
|
||||
const { result } = renderHook(() => useCustomers(), { 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(allEndUsersCall).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("should not execute query when userRole is null", async () => {
|
||||
// Mock null userRole
|
||||
mockUseAuthorized.mockReturnValue({
|
||||
accessToken: "test-access-token",
|
||||
userRole: null,
|
||||
userId: "test-user-id",
|
||||
token: "test-token",
|
||||
userEmail: "test@example.com",
|
||||
premiumUser: false,
|
||||
disabledPersonalKeyCreation: null,
|
||||
showSSOBanner: false,
|
||||
});
|
||||
|
||||
const { result } = renderHook(() => useCustomers(), { 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(allEndUsersCall).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("should not execute query when userRole is empty string", async () => {
|
||||
// Mock empty string userRole
|
||||
mockUseAuthorized.mockReturnValue({
|
||||
accessToken: "test-access-token",
|
||||
userRole: "",
|
||||
userId: "test-user-id",
|
||||
token: "test-token",
|
||||
userEmail: "test@example.com",
|
||||
premiumUser: false,
|
||||
disabledPersonalKeyCreation: null,
|
||||
showSSOBanner: false,
|
||||
});
|
||||
|
||||
const { result } = renderHook(() => useCustomers(), { 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(allEndUsersCall).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("should not execute query when both accessToken and userRole are missing", async () => {
|
||||
// Mock both auth values missing
|
||||
mockUseAuthorized.mockReturnValue({
|
||||
accessToken: null,
|
||||
userRole: null,
|
||||
userId: "test-user-id",
|
||||
token: null,
|
||||
userEmail: "test@example.com",
|
||||
premiumUser: false,
|
||||
disabledPersonalKeyCreation: null,
|
||||
showSSOBanner: false,
|
||||
});
|
||||
|
||||
const { result } = renderHook(() => useCustomers(), { 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(allEndUsersCall).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("should execute query when accessToken is present and userRole is Admin", async () => {
|
||||
// Mock successful API call
|
||||
(allEndUsersCall as any).mockResolvedValue(mockCustomersResponse);
|
||||
|
||||
// Ensure auth values are set (already done in beforeEach)
|
||||
const { result } = renderHook(() => useCustomers(), { wrapper });
|
||||
|
||||
// Wait for query to execute
|
||||
await waitFor(() => {
|
||||
expect(result.current.isLoading).toBe(false);
|
||||
});
|
||||
|
||||
expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token");
|
||||
expect(allEndUsersCall).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("should execute query when accessToken is present and userRole is proxy_admin", async () => {
|
||||
// Mock successful API call
|
||||
(allEndUsersCall as any).mockResolvedValue(mockCustomersResponse);
|
||||
|
||||
// Mock proxy_admin role
|
||||
mockUseAuthorized.mockReturnValue({
|
||||
accessToken: "test-access-token",
|
||||
userRole: "proxy_admin",
|
||||
userId: "test-user-id",
|
||||
token: "test-token",
|
||||
userEmail: "test@example.com",
|
||||
premiumUser: false,
|
||||
disabledPersonalKeyCreation: null,
|
||||
showSSOBanner: false,
|
||||
});
|
||||
|
||||
const { result } = renderHook(() => useCustomers(), { wrapper });
|
||||
|
||||
// Wait for query to execute
|
||||
await waitFor(() => {
|
||||
expect(result.current.isLoading).toBe(false);
|
||||
});
|
||||
|
||||
expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token");
|
||||
expect(allEndUsersCall).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("should return empty customers array when API returns empty data", async () => {
|
||||
// Mock API returning empty customers array
|
||||
(allEndUsersCall as any).mockResolvedValue([]);
|
||||
|
||||
const { result } = renderHook(() => useCustomers(), { wrapper });
|
||||
|
||||
// Wait for success
|
||||
await waitFor(() => {
|
||||
expect(result.current.isLoading).toBe(false);
|
||||
expect(result.current.isSuccess).toBe(true);
|
||||
});
|
||||
|
||||
expect(result.current.data).toEqual([]);
|
||||
expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token");
|
||||
});
|
||||
|
||||
it("should handle network timeout error", async () => {
|
||||
const timeoutError = new Error("Network timeout");
|
||||
|
||||
// Mock network timeout
|
||||
(allEndUsersCall as any).mockRejectedValue(timeoutError);
|
||||
it("does not fetch when the access token is missing", () => {
|
||||
mockUseAuthorized.mockReturnValue({ ...authorized, accessToken: null, token: null });
|
||||
|
||||
const { result } = renderHook(() => useCustomers(), { wrapper });
|
||||
|
||||
// Wait for error
|
||||
await waitFor(() => {
|
||||
expect(result.current.isError).toBe(true);
|
||||
});
|
||||
expect(result.current.isFetched).toBe(false);
|
||||
expect(mockGet).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
expect(result.current.error).toEqual(timeoutError);
|
||||
expect(result.current.data).toBeUndefined();
|
||||
it("does not fetch when the user is not an admin", () => {
|
||||
mockUseAuthorized.mockReturnValue({ ...authorized, userRole: "member" });
|
||||
|
||||
const { result } = renderHook(() => useCustomers(), { wrapper });
|
||||
|
||||
expect(result.current.isFetched).toBe(false);
|
||||
expect(mockGet).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,42 +1,19 @@
|
|||
import { allEndUsersCall } from "@/components/networking";
|
||||
import { useQuery } from "@tanstack/react-query";
|
||||
import { createQueryKeys } from "../common/queryKeysFactory";
|
||||
import { fetchClient } from "@/lib/http/api";
|
||||
import { all_admin_roles } from "@/utils/roles";
|
||||
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
|
||||
import type { components } from "@/lib/http/schema";
|
||||
|
||||
export type EndUser = components["schemas"]["LiteLLM_EndUserTable"];
|
||||
|
||||
const customersKeys = createQueryKeys("customers");
|
||||
|
||||
export interface Customer {
|
||||
user_id: string;
|
||||
alias?: string | null;
|
||||
spend: number;
|
||||
blocked: boolean;
|
||||
allowed_model_region?: string | null;
|
||||
default_model?: string | null;
|
||||
budget_id?: string | null;
|
||||
litellm_budget_table?: {
|
||||
budget_id: string;
|
||||
max_budget?: number | null;
|
||||
soft_budget?: number | null;
|
||||
max_parallel_requests?: number | null;
|
||||
tpm_limit?: number | null;
|
||||
rpm_limit?: number | null;
|
||||
model_max_budget?: Record<string, unknown> | null;
|
||||
budget_duration?: string | null;
|
||||
budget_reset_at?: string | null;
|
||||
created_at: string;
|
||||
created_by: string;
|
||||
updated_at: string;
|
||||
updated_by: string;
|
||||
} | null;
|
||||
}
|
||||
|
||||
export type CustomersResponse = Customer[];
|
||||
|
||||
export const useCustomers = () => {
|
||||
const { accessToken, userRole } = useAuthorized();
|
||||
return useQuery<CustomersResponse>({
|
||||
return useQuery({
|
||||
queryKey: customersKeys.list({}),
|
||||
queryFn: async () => await allEndUsersCall(accessToken!),
|
||||
queryFn: async () => (await fetchClient.GET("/customer/list")).data ?? [],
|
||||
enabled: Boolean(accessToken) && all_admin_roles.includes(userRole!),
|
||||
});
|
||||
};
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue