Merge pull request #21980 from BerriAI/litellm_ui_testing_coverage_00

[Refactor] UI - Onboarding: Extract testable view components

ui/litellm-dashboard/src/app/(dashboard)/hooks/onboarding/useOnboarding.test.ts

@@ -0,0 +1,136 @@
+import { getOnboardingCredentials, claimOnboardingToken } from "@/components/networking";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import { renderHook, waitFor, act } from "@testing-library/react";
+import React, { ReactNode } from "react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { useOnboardingCredentials, useClaimOnboardingToken } from "./useOnboarding";
+
+vi.mock("@/components/networking", () => ({
+  getOnboardingCredentials: vi.fn(),
+  claimOnboardingToken: vi.fn(),
+}));
+
+const mockUseUIConfig = vi.fn();
+vi.mock("@/app/(dashboard)/hooks/uiConfig/useUIConfig", () => ({
+  useUIConfig: () => mockUseUIConfig(),
+}));
+
+const mockCredentialsResponse = { token: "mock.jwt.token", login_url: "http://example.com/login" };
+
+describe("useOnboardingCredentials", () => {
+  let queryClient: QueryClient;
+
+  beforeEach(() => {
+    queryClient = new QueryClient({
+      defaultOptions: {
+        queries: { retry: false },
+        mutations: { retry: false },
+      },
+    });
+    vi.clearAllMocks();
+    mockUseUIConfig.mockReturnValue({ isLoading: false });
+  });
+
+  const wrapper = ({ children }: { children: ReactNode }) =>
+    React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+  it("fetches credentials when inviteId is provided and UIConfig is loaded", async () => {
+    (getOnboardingCredentials as any).mockResolvedValue(mockCredentialsResponse);
+
+    const { result } = renderHook(() => useOnboardingCredentials("invite-123"), { wrapper });
+
+    expect(result.current.isLoading).toBe(true);
+
+    await waitFor(() => expect(result.current.isSuccess).toBe(true));
+
+    expect(result.current.data).toEqual(mockCredentialsResponse);
+    expect(getOnboardingCredentials).toHaveBeenCalledWith("invite-123");
+    expect(getOnboardingCredentials).toHaveBeenCalledTimes(1);
+  });
+
+  it("does not fetch when inviteId is null", async () => {
+    const { result } = renderHook(() => useOnboardingCredentials(null), { wrapper });
+
+    expect(result.current.isLoading).toBe(false);
+    expect(result.current.isFetched).toBe(false);
+    expect(getOnboardingCredentials).not.toHaveBeenCalled();
+  });
+
+  it("does not fetch while UIConfig is loading", async () => {
+    mockUseUIConfig.mockReturnValue({ isLoading: true });
+
+    const { result } = renderHook(() => useOnboardingCredentials("invite-123"), { wrapper });
+
+    expect(result.current.isLoading).toBe(false);
+    expect(result.current.isFetched).toBe(false);
+    expect(getOnboardingCredentials).not.toHaveBeenCalled();
+  });
+
+  it("exposes error state when fetch fails", async () => {
+    const error = new Error("Invalid invite");
+    (getOnboardingCredentials as any).mockRejectedValue(error);
+
+    const { result } = renderHook(() => useOnboardingCredentials("bad-invite"), { wrapper });
+
+    await waitFor(() => expect(result.current.isError).toBe(true));
+
+    expect(result.current.error).toEqual(error);
+  });
+});
+
+describe("useClaimOnboardingToken", () => {
+  let queryClient: QueryClient;
+
+  beforeEach(() => {
+    queryClient = new QueryClient({
+      defaultOptions: {
+        queries: { retry: false },
+        mutations: { retry: false },
+      },
+    });
+    vi.clearAllMocks();
+    mockUseUIConfig.mockReturnValue({ isLoading: false });
+  });
+
+  const wrapper = ({ children }: { children: ReactNode }) =>
+    React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+  it("calls claimOnboardingToken with correct params", async () => {
+    (claimOnboardingToken as any).mockResolvedValue({ success: true });
+
+    const { result } = renderHook(() => useClaimOnboardingToken(), { wrapper });
+
+    act(() => {
+      result.current.mutate({
+        accessToken: "acc-token",
+        inviteId: "invite-123",
+        userId: "user-456",
+        password: "secret",
+      });
+    });
+
+    await waitFor(() => expect(result.current.isSuccess).toBe(true));
+
+    expect(claimOnboardingToken).toHaveBeenCalledWith("acc-token", "invite-123", "user-456", "secret");
+  });
+
+  it("exposes error state when mutation fails", async () => {
+    const error = new Error("Claim failed");
+    (claimOnboardingToken as any).mockRejectedValue(error);
+
+    const { result } = renderHook(() => useClaimOnboardingToken(), { wrapper });
+
+    act(() => {
+      result.current.mutate({
+        accessToken: "acc-token",
+        inviteId: "invite-123",
+        userId: "user-456",
+        password: "secret",
+      });
+    });
+
+    await waitFor(() => expect(result.current.isError).toBe(true));
+
+    expect(result.current.error).toEqual(error);
+  });
+});

ui/litellm-dashboard/src/app/(dashboard)/hooks/onboarding/useOnboarding.ts

@@ -0,0 +1,37 @@
+import { claimOnboardingToken, getOnboardingCredentials } from "@/components/networking";
+import { useMutation, useQuery } from "@tanstack/react-query";
+import { useUIConfig } from "../uiConfig/useUIConfig";
+import { createQueryKeys } from "../common/queryKeysFactory";
+
+const onboardingKeys = createQueryKeys("onboarding");
+
+export interface OnboardingCredentials {
+  token: string;
+  login_url: string;
+}
+
+export const useOnboardingCredentials = (inviteId: string | null) => {
+  const { isLoading: isUIConfigLoading } = useUIConfig();
+  return useQuery<OnboardingCredentials>({
+    queryKey: onboardingKeys.detail(inviteId ?? ""),
+    queryFn: async () => {
+      if (!inviteId) throw new Error("inviteId is required");
+      return getOnboardingCredentials(inviteId);
+    },
+    enabled: Boolean(inviteId) && !isUIConfigLoading,
+  });
+};
+
+export interface ClaimTokenParams {
+  accessToken: string;
+  inviteId: string;
+  userId: string;
+  password: string;
+}
+
+export const useClaimOnboardingToken = () => {
+  return useMutation({
+    mutationFn: async ({ accessToken, inviteId, userId, password }: ClaimTokenParams) =>
+      await claimOnboardingToken(accessToken, inviteId, userId, password),
+  });
+};

ui/litellm-dashboard/src/app/onboarding/OnboardingErrorView.test.tsx

@@ -0,0 +1,25 @@
+import React from "react";
+import { render, screen } from "@testing-library/react";
+import { describe, it, expect } from "vitest";
+import { OnboardingErrorView } from "./OnboardingErrorView";
+
+describe("OnboardingErrorView", () => {
+  it("should show the failed to load invitation message", () => {
+    render(<OnboardingErrorView />);
+    expect(screen.getByText("Failed to load invitation")).toBeInTheDocument();
+  });
+
+  it("should show the expiry description", () => {
+    render(<OnboardingErrorView />);
+    expect(
+      screen.getByText("The invitation link may be invalid or expired.")
+    ).toBeInTheDocument();
+  });
+
+  it("should render a Back to Login link pointing to /ui/login", () => {
+    render(<OnboardingErrorView />);
+    // antd Button with href renders as an <a> element
+    const link = screen.getByRole("link", { name: "Back to Login" });
+    expect(link).toHaveAttribute("href", "/ui/login");
+  });
+});

ui/litellm-dashboard/src/app/onboarding/OnboardingErrorView.tsx

@@ -0,0 +1,18 @@
+import React from "react";
+import { Alert, Button } from "antd";
+
+export function OnboardingErrorView() {
+  return (
+    <div className="mx-auto w-full max-w-md mt-10">
+      <Alert
+        type="error"
+        message="Failed to load invitation"
+        description="The invitation link may be invalid or expired."
+        showIcon
+      />
+      <div className="mt-4">
+        <Button href="/ui/login">Back to Login</Button>
+      </div>
+    </div>
+  );
+}

ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx

@@ -0,0 +1,70 @@
+"use client";
+import React from "react";
+import { useSearchParams } from "next/navigation";
+import { jwtDecode } from "jwt-decode";
+import { useOnboardingCredentials, useClaimOnboardingToken } from "@/app/(dashboard)/hooks/onboarding/useOnboarding";
+import { getProxyBaseUrl } from "@/components/networking";
+import { OnboardingLoadingView } from "./OnboardingLoadingView";
+import { OnboardingErrorView } from "./OnboardingErrorView";
+import { OnboardingFormBody } from "./OnboardingFormBody";
+
+type OnboardingFormProps = {
+  variant: "signup" | "reset_password";
+};
+
+export function OnboardingForm({ variant }: OnboardingFormProps) {
+  const searchParams = useSearchParams()!;
+  const inviteId = searchParams.get("invitation_id");
+  const [claimError, setClaimError] = React.useState<string | null>(null);
+
+  const {
+    data: credentialsData,
+    isLoading: isCredentialsLoading,
+    isError: isCredentialsError,
+  } = useOnboardingCredentials(inviteId);
+
+  const { mutate: claimToken, isPending } = useClaimOnboardingToken();
+
+  const decoded = credentialsData?.token
+    ? (jwtDecode(credentialsData.token) as { [key: string]: any })
+    : null;
+  const userEmail: string = decoded?.user_email ?? "";
+  const userId: string | null = decoded?.user_id ?? null;
+  const accessToken: string | null = decoded?.key ?? null;
+  const jwtToken: string | null = credentialsData?.token ?? null;
+
+  const handleSubmit = (formValues: { password: string }) => {
+    if (!accessToken || !jwtToken || !userId || !inviteId) return;
+
+    setClaimError(null);
+
+    claimToken(
+      { accessToken, inviteId, userId, password: formValues.password },
+      {
+        onSuccess: () => {
+          document.cookie = `token=${jwtToken}; path=/; SameSite=Lax`;
+          const proxyBaseUrl = getProxyBaseUrl();
+          window.location.href = proxyBaseUrl
+            ? `${proxyBaseUrl}/ui/?login=success`
+            : "/ui/?login=success";
+        },
+        onError: (error: Error) => {
+          setClaimError(error.message || "Failed to submit. Please try again.");
+        },
+      }
+    );
+  };
+
+  if (isCredentialsLoading) return <OnboardingLoadingView />;
+  if (isCredentialsError) return <OnboardingErrorView />;
+
+  return (
+    <OnboardingFormBody
+      variant={variant}
+      userEmail={userEmail}
+      isPending={isPending}
+      claimError={claimError}
+      onSubmit={handleSubmit}
+    />
+  );
+}

ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.test.tsx

@@ -0,0 +1,89 @@
+import React from "react";
+import { render, screen, waitFor } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { OnboardingFormBody } from "./OnboardingFormBody";
+
+const defaultProps = {
+  variant: "signup" as const,
+  userEmail: "test@example.com",
+  isPending: false,
+  claimError: null,
+  onSubmit: vi.fn(),
+};
+
+describe("OnboardingFormBody", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("should show 'Sign Up' heading for signup variant", () => {
+    render(<OnboardingFormBody {...defaultProps} />);
+    expect(screen.getByRole("heading", { name: "Sign Up" })).toBeInTheDocument();
+  });
+
+  it("should show 'Reset Password' heading for reset_password variant", () => {
+    render(<OnboardingFormBody {...defaultProps} variant="reset_password" />);
+    expect(screen.getByRole("heading", { name: "Reset Password" })).toBeInTheDocument();
+  });
+
+  it("should show SSO alert for signup variant", () => {
+    render(<OnboardingFormBody {...defaultProps} />);
+    expect(screen.getByText("SSO")).toBeInTheDocument();
+  });
+
+  it("should hide SSO alert for reset_password variant", () => {
+    render(<OnboardingFormBody {...defaultProps} variant="reset_password" />);
+    expect(screen.queryByText("SSO")).not.toBeInTheDocument();
+  });
+
+  it("should pre-fill the email field with userEmail", async () => {
+    render(<OnboardingFormBody {...defaultProps} userEmail="user@example.com" />);
+    await waitFor(() => {
+      expect(screen.getByLabelText("Email Address")).toHaveValue("user@example.com");
+    });
+  });
+
+  it("should disable the email field", () => {
+    render(<OnboardingFormBody {...defaultProps} />);
+    expect(screen.getByLabelText("Email Address")).toBeDisabled();
+  });
+
+  it("should show claimError message when claimError is set", () => {
+    render(<OnboardingFormBody {...defaultProps} claimError="Something went wrong" />);
+    expect(screen.getByText("Something went wrong")).toBeInTheDocument();
+  });
+
+  it("should not show claimError message when claimError is null", () => {
+    render(<OnboardingFormBody {...defaultProps} claimError={null} />);
+    expect(screen.queryByText("Something went wrong")).not.toBeInTheDocument();
+  });
+
+  it("should show a loading indicator on the submit button when isPending is true", () => {
+    render(<OnboardingFormBody {...defaultProps} isPending={true} />);
+    // antd v5 renders a loading icon with aria-label="loading" inside the button
+    expect(screen.getByRole("img", { name: "loading" })).toBeInTheDocument();
+  });
+
+  it("should call onSubmit with the typed password on form submit", async () => {
+    const user = userEvent.setup();
+    const onSubmit = vi.fn();
+    render(<OnboardingFormBody {...defaultProps} onSubmit={onSubmit} />);
+
+    await user.type(screen.getByLabelText("Password"), "mypassword");
+    await user.click(screen.getByRole("button", { name: /sign up/i }));
+
+    await waitFor(() => {
+      expect(onSubmit).toHaveBeenCalledWith(
+        expect.objectContaining({ password: "mypassword" })
+      );
+    });
+  });
+
+  it("should show 'Reset Password' on the submit button for reset_password variant", () => {
+    render(<OnboardingFormBody {...defaultProps} variant="reset_password" />);
+    expect(
+      screen.getByRole("button", { name: /reset password/i })
+    ).toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/app/onboarding/OnboardingFormBody.tsx

@@ -0,0 +1,93 @@
+import React from "react";
+import { Alert, Button, Card, Form, Input, Typography } from "antd";
+
+type OnboardingFormBodyProps = {
+  variant: "signup" | "reset_password";
+  userEmail: string;
+  isPending: boolean;
+  claimError: string | null;
+  onSubmit: (values: { password: string }) => void;
+};
+
+export function OnboardingFormBody({
+  variant,
+  userEmail,
+  isPending,
+  claimError,
+  onSubmit,
+}: OnboardingFormBodyProps) {
+  const [form] = Form.useForm();
+
+  React.useEffect(() => {
+    if (userEmail) form.setFieldValue("user_email", userEmail);
+  }, [userEmail, form]);
+
+  return (
+    <div className="mx-auto w-full max-w-md mt-10">
+      <Card>
+        <Typography.Title level={5} className="text-center mb-5">
+          🚅 LiteLLM
+        </Typography.Title>
+        <Typography.Title level={3}>
+          {variant === "reset_password" ? "Reset Password" : "Sign Up"}
+        </Typography.Title>
+        <Typography.Text>
+          {variant === "reset_password"
+            ? "Reset your password to access Admin UI."
+            : "Claim your user account to login to Admin UI."}
+        </Typography.Text>
+
+        {variant === "signup" && (
+          <Alert
+            className="mt-4"
+            type="info"
+            message="SSO"
+            description={
+              <div className="flex justify-between items-center">
+                <span>SSO is under the Enterprise Tier.</span>
+                <Button
+                  type="primary"
+                  size="small"
+                  href="https://forms.gle/W3U4PZpJGFHWtHyA9"
+                  target="_blank"
+                >
+                  Get Free Trial
+                </Button>
+              </div>
+            }
+            showIcon
+          />
+        )}
+
+        <Form className="mt-10 mb-5" layout="vertical" form={form} onFinish={(values) => onSubmit({ password: values.password })}>
+          <Form.Item label="Email Address" name="user_email">
+            <Input type="email" disabled />
+          </Form.Item>
+
+          <Form.Item
+            label="Password"
+            name="password"
+            rules={[{ required: true, message: "password required to sign up" }]}
+            help={
+              variant === "reset_password"
+                ? "Enter your new password"
+                : "Create a password for your account"
+            }
+          >
+            <Input.Password />
+          </Form.Item>
+
+          {claimError && (
+            <Alert type="error" message={claimError} showIcon className="mb-4" />
+          )}
+
+          <div className="mt-10">
+            <Button htmlType="submit" loading={isPending}>
+              {variant === "reset_password" ? "Reset Password" : "Sign Up"}
+            </Button>
+          </div>
+        </Form>
+      </Card>
+    </div>
+  );
+}

ui/litellm-dashboard/src/app/onboarding/OnboardingLoadingView.test.tsx

@@ -0,0 +1,16 @@
+import React from "react";
+import { render } from "@testing-library/react";
+import { describe, it, expect } from "vitest";
+import { OnboardingLoadingView } from "./OnboardingLoadingView";
+
+describe("OnboardingLoadingView", () => {
+  it("should render a spinner container", () => {
+    const { container } = render(<OnboardingLoadingView />);
+    expect(container.firstChild).toBeInTheDocument();
+  });
+
+  it("should apply centering layout classes", () => {
+    const { container } = render(<OnboardingLoadingView />);
+    expect(container.firstChild).toHaveClass("flex", "justify-center");
+  });
+});

ui/litellm-dashboard/src/app/onboarding/OnboardingLoadingView.tsx

@@ -0,0 +1,11 @@
+import React from "react";
+import { Spin } from "antd";
+import { LoadingOutlined } from "@ant-design/icons";
+
+export function OnboardingLoadingView() {
+  return (
+    <div className="mx-auto w-full max-w-md mt-10 flex justify-center">
+      <Spin indicator={<LoadingOutlined spin />} size="large" />
+    </div>
+  );
+}

ui/litellm-dashboard/src/app/onboarding/page.tsx

@@ -1,150 +1,28 @@
 "use client";
-import React, { Suspense, useEffect, useState } from "react";
+import React, { Suspense } from "react";
 import { useSearchParams } from "next/navigation";
-import { Card, Title, Text, TextInput, Callout, Button, Grid, Col } from "@tremor/react";
-import { RiCheckboxCircleLine } from "@remixicon/react";
-import {
-  getOnboardingCredentials,
-  claimOnboardingToken,
-  getUiConfig,
-  getProxyBaseUrl,
-} from "@/components/networking";
-import { jwtDecode } from "jwt-decode";
-import { Form, Button as Button2 } from "antd";
-import { getCookie } from "@/utils/cookieUtils";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import { OnboardingForm } from "./OnboardingForm";
+
+const queryClient = new QueryClient();
 
 function OnboardingContent() {
-  const [form] = Form.useForm();
   const searchParams = useSearchParams()!;
-  const token = getCookie("token");
-  const inviteID = searchParams.get("invitation_id");
   const action = searchParams.get("action");
-  const [accessToken, setAccessToken] = useState<string | null>(null);
-  const [defaultUserEmail, setDefaultUserEmail] = useState<string>("");
-  const [userEmail, setUserEmail] = useState<string>("");
-  const [userID, setUserID] = useState<string | null>(null);
-  const [loginUrl, setLoginUrl] = useState<string>("");
-  const [jwtToken, setJwtToken] = useState<string>("");
-  const [getUiConfigLoading, setGetUiConfigLoading] = useState<boolean>(true);
-
-  useEffect(() => {
-    getUiConfig().then((data) => {
-      // get the information for constructing the proxy base url, and then set the token and auth loading
-      console.log("ui config in onboarding.tsx:", data);
-      setGetUiConfigLoading(false);
-    });
-  }, []);
-
-  useEffect(() => {
-    if (!inviteID || getUiConfigLoading) {
-      // wait for the ui config to be loaded
-      return;
-    }
-
-    getOnboardingCredentials(inviteID).then((data) => {
-      const login_url = data.login_url;
-      console.log("login_url:", login_url);
-      setLoginUrl(login_url);
-
-      const token = data.token;
-      const decoded = jwtDecode(token) as { [key: string]: any };
-      setJwtToken(token);
-
-      console.log("decoded:", decoded);
-      setAccessToken(decoded.key);
-
-      console.log("decoded user email:", decoded.user_email);
-      const user_email = decoded.user_email;
-      setUserEmail(user_email);
-
-      const user_id = decoded.user_id;
-      setUserID(user_id);
-    });
-  }, [inviteID, getUiConfigLoading]);
-
-  const handleSubmit = (formValues: Record<string, any>) => {
-    console.log("in handle submit. accessToken:", accessToken, "token:", jwtToken, "formValues:", formValues);
-    if (!accessToken || !jwtToken) {
-      return;
-    }
-
-    formValues.user_email = userEmail;
-
-    if (!userID || !inviteID) {
-      return;
-    }
-    claimOnboardingToken(accessToken, inviteID, userID, formValues.password).then((data) => {
-      // set cookie "token" to jwtToken
-      document.cookie = "token=" + jwtToken;
-      
-      const proxyBaseUrl = getProxyBaseUrl();
-      console.log("proxyBaseUrl:", proxyBaseUrl);
-      
-      // Construct the full redirect URL using the proxyBaseUrl which includes the server root path
-      let redirectUrl = proxyBaseUrl ? `${proxyBaseUrl}/ui/?login=success` : "/ui/?login=success";
-      console.log("redirecting to:", redirectUrl);
-
-      window.location.href = redirectUrl;
-    });
-
-    // redirect to login page
-  };
-  return (
-    <div className="mx-auto w-full max-w-md mt-10">
-      <Card>
-        <Title className="text-sm mb-5 text-center">🚅 LiteLLM</Title>
-        <Title className="text-xl">{action === "reset_password" ? "Reset Password" : "Sign up"}</Title>
-        <Text>
-          {action === "reset_password"
-            ? "Reset your password to access Admin UI."
-            : "Claim your user account to login to Admin UI."}
-        </Text>
-
-        {action !== "reset_password" && (
-          <Callout className="mt-4" title="SSO" icon={RiCheckboxCircleLine} color="sky">
-            <Grid numItems={2} className="flex justify-between items-center">
-              <Col>SSO is under the Enterprise Tier.</Col>
-
-              <Col>
-                <Button variant="primary" className="mb-2">
-                  <a href="https://forms.gle/W3U4PZpJGFHWtHyA9" target="_blank">
-                    Get Free Trial
-                  </a>
-                </Button>
-              </Col>
-            </Grid>
-          </Callout>
-        )}
-
-        <Form className="mt-10 mb-5 mx-auto" layout="vertical" onFinish={handleSubmit}>
-          <>
-            <Form.Item label="Email Address" name="user_email">
-              <TextInput type="email" disabled={true} value={userEmail} defaultValue={userEmail} className="max-w-md" />
-            </Form.Item>
-
-            <Form.Item
-              label="Password"
-              name="password"
-              rules={[{ required: true, message: "password required to sign up" }]}
-              help={action === "reset_password" ? "Enter your new password" : "Create a password for your account"}
-            >
-              <TextInput placeholder="" type="password" className="max-w-md" />
-            </Form.Item>
-          </>
-
-          <div className="mt-10">
-            <Button2 htmlType="submit">{action === "reset_password" ? "Reset Password" : "Sign Up"}</Button2>
-          </div>
-        </Form>
-      </Card>
-    </div>
-  );
+  const variant = action === "reset_password" ? "reset_password" : "signup";
+  return <OnboardingForm variant={variant} />;
 }
 
 export default function Onboarding() {
   return (
-    <Suspense fallback={<div className="flex items-center justify-center min-h-screen">Loading...</div>}>
-      <OnboardingContent />
-    </Suspense>
+    <QueryClientProvider client={queryClient}>
+      <Suspense
+        fallback={
+          <div className="flex items-center justify-center min-h-screen">Loading...</div>
+        }
+      >
+        <OnboardingContent />
+      </Suspense>
+    </QueryClientProvider>
   );
 }