fix test assertion

Author: kantord

src/features/workspace/components/__tests__/workspace-preferred-model.test.tsx

@@ -3,7 +3,7 @@ import { screen, waitFor } from "@testing-library/react";
 import { WorkspacePreferredModel } from "../workspace-preferred-model";
 import userEvent from "@testing-library/user-event";
 
-test("render model overrides", () => {
+test("render model overrides", async () => {
   render(
     <WorkspacePreferredModel
       isArchived={false}
@@ -19,7 +19,10 @@ test("render model overrides", () => {
   expect(
     screen.getByRole("button", { name: /select the model/i }),
   ).toBeVisible();
-  expect(screen.getByRole("button", { name: /save/i })).toBeVisible();
+
+  await waitFor(() => {
+    expect(screen.getByRole("button", { name: /save/i })).toBeVisible();
+  });
 });
 
 test("submit preferred model", async () => {

src/features/workspace/components/workspace-custom-instructions.tsx

@@ -123,7 +123,11 @@ function useCustomInstructionsValue({
   options: V1GetWorkspaceCustomInstructionsData;
   queryClient: QueryClient;
 }) {
-  const formState = useFormState({ prompt: initialValue });
+  const initialFormValues = useMemo(
+    () => ({ prompt: initialValue }),
+    [initialValue],
+  );
+  const formState = useFormState(initialFormValues);
   const { values, updateFormValues } = formState;
 
   // Subscribe to changes in the workspace system prompt value in the query cache

src/features/workspace/components/workspace-preferred-model.tsx

@@ -1,6 +1,5 @@
 import {
   Alert,
-  Button,
   Card,
   CardBody,
   CardFooter,
@@ -16,6 +15,7 @@ import { FormEvent } from "react";
 import { usePreferredModelWorkspace } from "../hooks/use-preferred-preferred-model";
 import { Select, SelectButton } from "@stacklok/ui-kit";
 import { useQueryListAllModelsForAllProviders } from "@/hooks/use-query-list-all-models-for-all-providers";
+import { FormButtons } from "@/components/FormButtons";
 
 function MissingProviderBanner() {
   return (
@@ -39,11 +39,10 @@ export function WorkspacePreferredModel({
   workspaceName: string;
   isArchived: boolean | undefined;
 }) {
-  const { preferredModel, setPreferredModel, isPending } =
-    usePreferredModelWorkspace(workspaceName);
+  const { formState, isPending } = usePreferredModelWorkspace(workspaceName);
   const { mutateAsync } = useMutationPreferredModelWorkspace();
   const { data: providerModels = [] } = useQueryListAllModelsForAllProviders();
-  const { model, provider_id } = preferredModel;
+  const { model, provider_id } = formState.values.preferredModel;
   const isModelsEmpty = !isPending && providerModels.length === 0;
 
   const handleSubmit = (event: FormEvent) => {
@@ -84,16 +83,18 @@ export function WorkspacePreferredModel({
                 isRequired
                 isDisabled={isModelsEmpty}
                 className="w-full"
-                selectedKey={preferredModel?.model}
+                selectedKey={formState.values.preferredModel?.model}
                 placeholder="Select the model"
                 onSelectionChange={(model) => {
                   const preferredModelProvider = providerModels.find(
                     (item) => item.name === model,
                   );
                   if (preferredModelProvider) {
-                    setPreferredModel({
-                      model: preferredModelProvider.name,
-                      provider_id: preferredModelProvider.provider_id,
+                    formState.updateFormValues({
+                      preferredModel: {
+                        model: preferredModelProvider.name,
+                        provider_id: preferredModelProvider.provider_id,
+                      },
                     });
                   }
                 }}
@@ -109,9 +110,11 @@ export function WorkspacePreferredModel({
           </div>
         </CardBody>
         <CardFooter className="justify-end">
-          <Button isDisabled={isArchived || isModelsEmpty} type="submit">
-            Save
-          </Button>
+          <FormButtons
+            isPending={isPending}
+            formState={formState}
+            canSubmit={!isArchived}
+          />
         </CardFooter>
       </Card>
     </Form>

src/features/workspace/hooks/use-preferred-preferred-model.ts

@@ -1,7 +1,8 @@
 import { MuxRule, V1GetWorkspaceMuxesData } from "@/api/generated";
 import { v1GetWorkspaceMuxesOptions } from "@/api/generated/@tanstack/react-query.gen";
+import { useFormState } from "@/hooks/useFormState";
 import { useQuery } from "@tanstack/react-query";
-import { useEffect, useMemo, useState } from "react";
+import { useMemo } from "react";
 
 type ModelRule = Omit<MuxRule, "matcher_type" | "matcher"> & {};
 
@@ -21,8 +22,6 @@ const usePreferredModel = (options: {
 };
 
 export const usePreferredModelWorkspace = (workspaceName: string) => {
-  const [preferredModel, setPreferredModel] =
-    useState<ModelRule>(DEFAULT_STATE);
   const options: V1GetWorkspaceMuxesData &
     Omit<V1GetWorkspaceMuxesData, "body"> = useMemo(
     () => ({
@@ -31,12 +30,10 @@ export const usePreferredModelWorkspace = (workspaceName: string) => {
     [workspaceName],
   );
   const { data, isPending } = usePreferredModel(options);
+  const providerModel = data?.[0];
+  const formState = useFormState<{ preferredModel: ModelRule }>({
+    preferredModel: providerModel ?? DEFAULT_STATE,
+  });
 
-  useEffect(() => {
-    const providerModel = data?.[0];
-
-    setPreferredModel(providerModel ?? DEFAULT_STATE);
-  }, [data, setPreferredModel]);
-
-  return { preferredModel, setPreferredModel, isPending };
+  return { isPending, formState };
 };

src/hooks/useFormState.ts

@@ -1,5 +1,5 @@
 import { isEqual } from "lodash";
-import { useState } from "react";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 
 export type FormState<T> = {
   values: T;
@@ -8,23 +8,50 @@ export type FormState<T> = {
   isDirty: boolean;
 };
 
+function useDeepMemo<T>(value: T): T {
+  const ref = useRef<T>(value);
+  if (!isEqual(ref.current, value)) {
+    ref.current = value;
+  }
+  return ref.current;
+}
+
 export function useFormState<Values extends Record<string, unknown>>(
   initialValues: Values,
 ): FormState<Values> {
+  const memoizedInitialValues = useDeepMemo(initialValues);
+
   // this could be replaced with some form library later
-  const [values, setValues] = useState<Values>(initialValues);
-  const updateFormValues = (newState: Partial<Values>) => {
+  const [values, setValues] = useState<Values>(memoizedInitialValues);
+  const [originalValues, setOriginalValues] = useState<Values>(values);
+
+  useEffect(() => {
+    // this logic supports the use case when the initialValues change
+    // due to an async request for instance
+    setOriginalValues(memoizedInitialValues);
+    setValues(memoizedInitialValues);
+  }, [memoizedInitialValues]);
+
+  const updateFormValues = useCallback((newState: Partial<Values>) => {
     setValues((prevState: Values) => ({
       ...prevState,
       ...newState,
     }));
-  };
+  }, []);
+
+  const resetForm = useCallback(() => {
+    setValues(originalValues);
+  }, [originalValues]);
 
-  const resetForm = () => {
-    setValues(initialValues);
-  };
+  const isDirty = useMemo(
+    () => !isEqual(values, originalValues),
+    [values, originalValues],
+  );
 
-  const isDirty = !isEqual(values, initialValues);
+  const formState = useMemo(
+    () => ({ values, updateFormValues, resetForm, isDirty }),
+    [values, updateFormValues, resetForm, isDirty],
+  );
 
-  return { values, updateFormValues, resetForm, isDirty };
+  return formState;
 }