Adding tests

ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.test.tsx

@@ -151,6 +151,117 @@ describe("BaseSSOSettingsForm", () => {
       expect(screen.getByText("Default Role")).toBeInTheDocument();
     });
   });
+
+  it("should show team mappings checkbox for okta provider", async () => {
+    const TestWrapper = () => {
+      const [form] = Form.useForm();
+      const handleSubmit = vi.fn();
+
+      return <BaseSSOSettingsForm form={form} onFormSubmit={handleSubmit} />;
+    };
+
+    renderWithProviders(<TestWrapper />);
+
+    const providerSelect = screen.getByLabelText("SSO Provider");
+    await act(async () => {
+      fireEvent.mouseDown(providerSelect);
+    });
+
+    await waitFor(() => {
+      const oktaOption = screen.getByText(/okta/i);
+      fireEvent.click(oktaOption);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Use Team Mappings")).toBeInTheDocument();
+    });
+  });
+
+  it("should show team mappings checkbox for generic provider", async () => {
+    const TestWrapper = () => {
+      const [form] = Form.useForm();
+      const handleSubmit = vi.fn();
+
+      return <BaseSSOSettingsForm form={form} onFormSubmit={handleSubmit} />;
+    };
+
+    renderWithProviders(<TestWrapper />);
+
+    const providerSelect = screen.getByLabelText("SSO Provider");
+    await act(async () => {
+      fireEvent.mouseDown(providerSelect);
+    });
+
+    await waitFor(() => {
+      const genericOption = screen.getByText(/generic sso/i);
+      fireEvent.click(genericOption);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Use Team Mappings")).toBeInTheDocument();
+    });
+  });
+
+  it("should show team IDs JWT field when use_team_mappings is checked for okta provider", async () => {
+    const TestWrapper = () => {
+      const [form] = Form.useForm();
+      const handleSubmit = vi.fn();
+
+      return <BaseSSOSettingsForm form={form} onFormSubmit={handleSubmit} />;
+    };
+
+    renderWithProviders(<TestWrapper />);
+
+    const providerSelect = screen.getByLabelText("SSO Provider");
+    await act(async () => {
+      fireEvent.mouseDown(providerSelect);
+    });
+
+    await waitFor(() => {
+      const oktaOption = screen.getByText(/okta/i);
+      fireEvent.click(oktaOption);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Use Team Mappings")).toBeInTheDocument();
+    });
+
+    const checkbox = screen.getByLabelText("Use Team Mappings");
+    await act(async () => {
+      fireEvent.click(checkbox);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Team IDs JWT Field")).toBeInTheDocument();
+    });
+  });
+
+  it("should not show team mappings checkbox for google provider", async () => {
+    const TestWrapper = () => {
+      const [form] = Form.useForm();
+      const handleSubmit = vi.fn();
+
+      return <BaseSSOSettingsForm form={form} onFormSubmit={handleSubmit} />;
+    };
+
+    renderWithProviders(<TestWrapper />);
+
+    const providerSelect = screen.getByLabelText("SSO Provider");
+    await act(async () => {
+      fireEvent.mouseDown(providerSelect);
+    });
+
+    await waitFor(() => {
+      const googleOption = screen.getByText(/google sso/i);
+      fireEvent.click(googleOption);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Google Client ID")).toBeInTheDocument();
+    });
+
+    expect(screen.queryByText("Use Team Mappings")).not.toBeInTheDocument();
+  });
 });
 
 describe("renderProviderFields", () => {

ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx

@@ -105,6 +105,14 @@ const createRoleMappingsSSOData = (overrides: Record<string, any> = {}) =>
     ...overrides,
   });
 
+const createTeamMappingsSSOData = (overrides: Record<string, any> = {}) =>
+  createGenericSSOData({
+    team_mappings: {
+      team_ids_jwt_field: overrides.team_ids_jwt_field || "teams",
+    },
+    ...overrides,
+  });
+
 // Mock utilities
 const createMockHooks = (): {
   useSSOSettings: SSOSettingsHookReturn;
@@ -577,6 +585,104 @@ describe("EditSSOSettingsModal", () => {
         });
       });
     });
+  });
+
+  describe("Team Mappings", () => {
+    it("processes team mappings when team_mappings exists", async () => {
+      const ssoData = createTeamMappingsSSOData();
+
+      setupMocks({
+        useSSOSettings: { data: ssoData, isLoading: false, error: null },
+      });
+
+      renderComponent();
+
+      await waitFor(() => {
+        expect(mockForm.setFieldsValue).toHaveBeenCalledWith({
+          sso_provider: SSO_PROVIDERS.GENERIC,
+          ...ssoData.values,
+          use_team_mappings: true,
+          team_ids_jwt_field: "teams",
+        });
+      });
+    });
+
+    it("handles team mappings with custom JWT field name", async () => {
+      const ssoData = createTeamMappingsSSOData({
+        team_ids_jwt_field: "custom_teams_field",
+      });
+
+      setupMocks({
+        useSSOSettings: { data: ssoData, isLoading: false, error: null },
+      });
+
+      renderComponent();
+
+      await waitFor(() => {
+        expect(mockForm.setFieldsValue).toHaveBeenCalledWith({
+          sso_provider: SSO_PROVIDERS.GENERIC,
+          ...ssoData.values,
+          use_team_mappings: true,
+          team_ids_jwt_field: "custom_teams_field",
+        });
+      });
+    });
+
+    it("handles team mappings and role mappings together", async () => {
+      const ssoData = createGenericSSOData({
+        role_mappings: {
+          group_claim: "groups",
+          default_role: "internal_user",
+          roles: {
+            proxy_admin: ["admin-group"],
+            proxy_admin_viewer: [],
+            internal_user: [],
+            internal_user_viewer: [],
+          },
+        },
+        team_mappings: {
+          team_ids_jwt_field: "teams",
+        },
+      });
+
+      setupMocks({
+        useSSOSettings: { data: ssoData, isLoading: false, error: null },
+      });
+
+      renderComponent();
+
+      await waitFor(() => {
+        expect(mockForm.setFieldsValue).toHaveBeenCalledWith({
+          sso_provider: SSO_PROVIDERS.GENERIC,
+          ...ssoData.values,
+          use_role_mappings: true,
+          group_claim: "groups",
+          default_role: "internal_user",
+          proxy_admin_teams: "admin-group",
+          admin_viewer_teams: "",
+          internal_user_teams: "",
+          internal_viewer_teams: "",
+          use_team_mappings: true,
+          team_ids_jwt_field: "teams",
+        });
+      });
+    });
+
+    it("does not set team mapping fields when team_mappings is not present", async () => {
+      const ssoData = createGenericSSOData();
+
+      setupMocks({
+        useSSOSettings: { data: ssoData, isLoading: false, error: null },
+      });
+
+      renderComponent();
+
+      await waitFor(() => {
+        const callArgs = mockForm.setFieldsValue.mock.calls[0][0];
+        expect(callArgs.use_team_mappings).toBeUndefined();
+        expect(callArgs.team_ids_jwt_field).toBeUndefined();
+      });
+    });
 
     it("handles provider detection with partial SSO data", async () => {
       const ssoData = createSSOData({

ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts

@@ -12,6 +12,8 @@ describe("processSSOSettingsPayload", () => {
         default_role: "proxy_admin",
         group_claim: "groups",
         use_role_mappings: false,
+        use_team_mappings: false,
+        team_ids_jwt_field: "teams",
         other_field: "value",
         another_field: 123,
       };
@@ -23,6 +25,7 @@ describe("processSSOSettingsPayload", () => {
         another_field: 123,
       });
       expect(result.role_mappings).toBeUndefined();
+      expect(result.team_mappings).toBeUndefined();
     });
 
     it("should return all fields except role mapping fields when use_role_mappings is not present", () => {
@@ -33,6 +36,8 @@ describe("processSSOSettingsPayload", () => {
         internal_viewer_teams: "viewer1",
         default_role: "proxy_admin",
         group_claim: "groups",
+        use_team_mappings: false,
+        team_ids_jwt_field: "teams",
         other_field: "value",
       };
 
@@ -42,6 +47,7 @@ describe("processSSOSettingsPayload", () => {
         other_field: "value",
       });
       expect(result.role_mappings).toBeUndefined();
+      expect(result.team_mappings).toBeUndefined();
     });
   });
 
@@ -253,6 +259,143 @@ describe("processSSOSettingsPayload", () => {
     });
   });
 
+  describe("without team mappings", () => {
+    it("should return all fields except team mapping fields when use_team_mappings is false", () => {
+      const formValues = {
+        use_team_mappings: false,
+        team_ids_jwt_field: "teams",
+        sso_provider: "okta",
+        other_field: "value",
+      };
+
+      const result = processSSOSettingsPayload(formValues);
+
+      expect(result).toEqual({
+        sso_provider: "okta",
+        other_field: "value",
+      });
+      expect(result.team_mappings).toBeUndefined();
+    });
+
+    it("should return all fields except team mapping fields when use_team_mappings is not present", () => {
+      const formValues = {
+        team_ids_jwt_field: "teams",
+        sso_provider: "generic",
+        other_field: "value",
+      };
+
+      const result = processSSOSettingsPayload(formValues);
+
+      expect(result).toEqual({
+        sso_provider: "generic",
+        other_field: "value",
+      });
+      expect(result.team_mappings).toBeUndefined();
+    });
+
+    it("should not include team mappings for unsupported providers even when use_team_mappings is true", () => {
+      const formValues = {
+        use_team_mappings: true,
+        team_ids_jwt_field: "teams",
+        sso_provider: "google",
+        other_field: "value",
+      };
+
+      const result = processSSOSettingsPayload(formValues);
+
+      expect(result).toEqual({
+        sso_provider: "google",
+        other_field: "value",
+      });
+      expect(result.team_mappings).toBeUndefined();
+    });
+
+    it("should not include team mappings for microsoft provider even when use_team_mappings is true", () => {
+      const formValues = {
+        use_team_mappings: true,
+        team_ids_jwt_field: "teams",
+        sso_provider: "microsoft",
+        other_field: "value",
+      };
+
+      const result = processSSOSettingsPayload(formValues);
+
+      expect(result).toEqual({
+        sso_provider: "microsoft",
+        other_field: "value",
+      });
+      expect(result.team_mappings).toBeUndefined();
+    });
+  });
+
+  describe("with team mappings enabled", () => {
+    it("should create team mappings for okta provider when use_team_mappings is true", () => {
+      const formValues = {
+        use_team_mappings: true,
+        team_ids_jwt_field: "teams",
+        sso_provider: "okta",
+        other_field: "value",
+      };
+
+      const result = processSSOSettingsPayload(formValues);
+
+      expect(result.other_field).toBe("value");
+      expect(result.team_mappings).toEqual({
+        team_ids_jwt_field: "teams",
+      });
+    });
+
+    it("should create team mappings for generic provider when use_team_mappings is true", () => {
+      const formValues = {
+        use_team_mappings: true,
+        team_ids_jwt_field: "custom_teams",
+        sso_provider: "generic",
+        other_field: "value",
+      };
+
+      const result = processSSOSettingsPayload(formValues);
+
+      expect(result.other_field).toBe("value");
+      expect(result.team_mappings).toEqual({
+        team_ids_jwt_field: "custom_teams",
+      });
+    });
+
+    it("should exclude team mapping fields from payload when team mappings are included", () => {
+      const formValues = {
+        use_team_mappings: true,
+        team_ids_jwt_field: "teams",
+        sso_provider: "okta",
+        other_field: "value",
+      };
+
+      const result = processSSOSettingsPayload(formValues);
+
+      expect(result.use_team_mappings).toBeUndefined();
+      expect(result.team_ids_jwt_field).toBeUndefined();
+    });
+
+    it("should handle team mappings and role mappings together", () => {
+      const formValues = {
+        use_team_mappings: true,
+        team_ids_jwt_field: "teams",
+        use_role_mappings: true,
+        group_claim: "groups",
+        default_role: "internal_user",
+        sso_provider: "okta",
+        other_field: "value",
+      };
+
+      const result = processSSOSettingsPayload(formValues);
+
+      expect(result.team_mappings).toEqual({
+        team_ids_jwt_field: "teams",
+      });
+      expect(result.role_mappings).toBeDefined();
+      expect(result.role_mappings.group_claim).toBe("groups");
+    });
+  });
+
   describe("edge cases", () => {
     it("should handle empty form values", () => {
       const result = processSSOSettingsPayload({});
@@ -263,6 +406,7 @@ describe("processSSOSettingsPayload", () => {
     it("should preserve other fields in the payload", () => {
       const formValues = {
         use_role_mappings: false,
+        use_team_mappings: false,
         sso_provider: "google",
         client_id: "123",
         client_secret: "secret",