From 553c52357c63cf15f16f37e14c3504480e2fcff0 Mon Sep 17 00:00:00 2001
From: Oleg Chendighelean <oleg.chendighelean@trili.tech>
Date: Fri, 6 Sep 2024 14:33:04 +0100
Subject: [PATCH] Update useMultiForm

---
 .../components/BackButton/ModalBackButton.tsx |   2 +-
 .../DeriveMnemonicAccount.tsx                 |  37 +++---
 .../DynamicDisclosure.test.tsx                | 115 +++++++++++++++---
 .../DynamicDisclosure/DynamicDisclosure.tsx   |  16 ++-
 .../DynamicDisclosure/useMultiForm.test.tsx   |   4 +-
 .../src/DynamicDisclosure/useMultiForm.tsx    |   1 +
 6 files changed, 128 insertions(+), 47 deletions(-)

diff --git a/apps/web/src/components/BackButton/ModalBackButton.tsx b/apps/web/src/components/BackButton/ModalBackButton.tsx
index d8b92328bd..2d064fa216 100644
--- a/apps/web/src/components/BackButton/ModalBackButton.tsx
+++ b/apps/web/src/components/BackButton/ModalBackButton.tsx
@@ -5,5 +5,5 @@ import { BaseBackButton } from "./BaseBackButton";
 export const ModalBackButton = () => {
   const { hasPrevious, goBack } = useDynamicModalContext();
 
-  return hasPrevious ? <BaseBackButton onClick={goBack} /> : null;
+  return hasPrevious ? <BaseBackButton onClick={() => goBack()} /> : null;
 };
diff --git a/apps/web/src/components/DeriveMnemonicAccount/DeriveMnemonicAccount.tsx b/apps/web/src/components/DeriveMnemonicAccount/DeriveMnemonicAccount.tsx
index f506d81c67..df3394a8fc 100644
--- a/apps/web/src/components/DeriveMnemonicAccount/DeriveMnemonicAccount.tsx
+++ b/apps/web/src/components/DeriveMnemonicAccount/DeriveMnemonicAccount.tsx
@@ -2,7 +2,6 @@ import { useToast } from "@chakra-ui/react";
 import { useDynamicModalContext } from "@umami/components";
 import { DEFAULT_ACCOUNT_LABEL, type MnemonicAccount } from "@umami/core";
 import { useAsyncActionHandler, useDeriveMnemonicAccount } from "@umami/state";
-import { useCallback } from "react";
 
 import { MasterPasswordModal } from "../MasterPasswordModal";
 import { NameAccountModal } from "../NameAccountModal";
@@ -12,32 +11,28 @@ type DeriveMnemonicAccountProps = {
 };
 
 export const DeriveMnemonicAccount = ({ account }: DeriveMnemonicAccountProps) => {
-  const { onClose, allFormValues } = useDynamicModalContext();
+  const { goBack, allFormValues } = useDynamicModalContext();
 
   const { handleAsyncAction } = useAsyncActionHandler();
   const toast = useToast();
   const deriveMnemonicAccount = useDeriveMnemonicAccount();
 
-  const handlePasswordSubmit = useCallback(
-    (password: string) =>
-      handleAsyncAction(
-        async () => {
-          await deriveMnemonicAccount({
-            fingerPrint: account.seedFingerPrint,
-            password,
-            label: allFormValues.accountName?.trim() || DEFAULT_ACCOUNT_LABEL,
-          });
-          onClose();
+  const handlePasswordSubmit = (password: string) =>
+    handleAsyncAction(
+      async () => {
+        await deriveMnemonicAccount({
+          fingerPrint: account.seedFingerPrint,
+          password,
+          label: allFormValues.accountName?.trim() || DEFAULT_ACCOUNT_LABEL,
+        });
+        goBack(0);
 
-          toast({
-            description: `New account created! Successfully derived account from ${account.seedFingerPrint}`,
-          });
-        },
-        { title: "Failed to derive new account" }
-      ),
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-    [allFormValues, account]
-  );
+        toast({
+          description: `New account created! Successfully derived account from ${account.seedFingerPrint}`,
+        });
+      },
+      { title: "Failed to derive new account" }
+    );
 
   return (
     <NameAccountModal
diff --git a/packages/components/src/DynamicDisclosure/DynamicDisclosure.test.tsx b/packages/components/src/DynamicDisclosure/DynamicDisclosure.test.tsx
index 9e7de36feb..5ff17aefab 100644
--- a/packages/components/src/DynamicDisclosure/DynamicDisclosure.test.tsx
+++ b/packages/components/src/DynamicDisclosure/DynamicDisclosure.test.tsx
@@ -1,5 +1,5 @@
-import { useDynamicDrawer, useDynamicModal } from "./DynamicDisclosure";
-import { act, render, renderHook, screen } from "../testUtils";
+import { useDynamicDrawer, useDynamicModal, useDynamicModalContext } from "./DynamicDisclosure";
+import { act, render, renderHook, screen, waitFor } from "../testUtils";
 
 describe("DynamicDisclosure", () => {
   describe("useDynamicModal", () => {
@@ -43,15 +43,12 @@ describe("DynamicDisclosure", () => {
       expect(onClose).toHaveBeenCalled();
     });
 
-    it("handles multiple modals with goBack", async () => {
+    it("calls the onClose callback when the modal is closed", async () => {
+      const onClose = jest.fn();
       const view = renderHook(() => useDynamicModal());
-      await act(() => view.result.current.openWith(<div>test data 1</div>));
-      await act(() => view.result.current.openWith(<div>test data 2</div>));
-      act(() => view.result.current.goBack());
-      expect(view.result.current.isOpen).toBe(true);
-      render(view.result.current.content);
-      expect(screen.getByText("test data 1")).toBeVisible();
-      expect(screen.queryByText("test data 2")).not.toBeInTheDocument();
+      await act(() => view.result.current.openWith(<div>test data</div>, { onClose }));
+      act(() => view.result.current.onClose());
+      expect(onClose).toHaveBeenCalled();
     });
   });
 
@@ -95,16 +92,96 @@ describe("DynamicDisclosure", () => {
       act(() => view.result.current.onClose());
       expect(onClose).toHaveBeenCalled();
     });
+  });
 
-    it("handles multiple drawers with goBack", async () => {
-      const view = renderHook(() => useDynamicDrawer());
-      await act(() => view.result.current.openWith(<div>test data 1</div>));
-      await act(() => view.result.current.openWith(<div>test data 2</div>));
-      act(() => view.result.current.goBack());
-      expect(view.result.current.isOpen).toBe(true);
-      render(view.result.current.content);
-      expect(screen.getByText("test data 1")).toBeVisible();
-      expect(screen.queryByText("test data 2")).not.toBeInTheDocument();
+  describe("goBack functionality", () => {
+    it("should go back one step when goBack is called without parameters", () => {
+      const { result } = renderHook(() => useDynamicModalContext());
+
+      act(() => {
+        result.current.openWith(<div>First</div>);
+        result.current.openWith(<div>Second</div>);
+      });
+
+      expect(result.current.isOpen).toBe(true);
+      expect(result.current.hasPrevious).toBe(true);
+
+      act(() => {
+        result.current.goBack();
+      });
+
+      expect(result.current.isOpen).toBe(true);
+      expect(result.current.hasPrevious).toBe(false);
+    });
+
+    it("should go back to specific index when goBack is called with a valid index", async () => {
+      const { result } = renderHook(() => useDynamicModalContext());
+
+      act(() => {
+        result.current.openWith(<div>First</div>);
+      });
+
+      act(() => {
+        result.current.openWith(<div>Second</div>);
+      });
+
+      act(() => {
+        result.current.openWith(<div>Third</div>);
+      });
+
+      expect(result.current.hasPrevious).toBe(true);
+
+      act(() => {
+        result.current.goBack(0);
+      });
+
+      expect(result.current.hasPrevious).toBe(false);
+      expect(screen.getByText("First")).toBeVisible();
+    });
+
+    it("should update hasPrevious correctly", () => {
+      const { result } = renderHook(() => useDynamicModalContext());
+
+      expect(result.current.hasPrevious).toBe(false);
+
+      act(() => {
+        result.current.openWith(<div>First</div>);
+      });
+
+      expect(result.current.hasPrevious).toBe(false);
+
+      act(() => {
+        result.current.openWith(<div>Second</div>);
+      });
+
+      expect(result.current.hasPrevious).toBe(true);
+
+      act(() => {
+        result.current.goBack();
+      });
+
+      expect(result.current.hasPrevious).toBe(false);
+    });
+
+    it("should go back one step when goBack is called with out-of-bounds index", async () => {
+      const { result } = renderHook(() => useDynamicModalContext());
+
+      act(() => {
+        result.current.openWith(<div>First</div>);
+        result.current.openWith(<div>Second</div>);
+        result.current.openWith(<div>Third</div>);
+      });
+
+      expect(result.current.isOpen).toBe(true);
+      expect(result.current.hasPrevious).toBe(true);
+
+      act(() => {
+        result.current.goBack(5);
+      });
+
+      expect(result.current.isOpen).toBe(true);
+      expect(result.current.hasPrevious).toBe(true);
+      expect(screen.getByText("Second")).toBeVisible();
     });
   });
 });
diff --git a/packages/components/src/DynamicDisclosure/DynamicDisclosure.tsx b/packages/components/src/DynamicDisclosure/DynamicDisclosure.tsx
index c604b112ee..800aa1de6b 100644
--- a/packages/components/src/DynamicDisclosure/DynamicDisclosure.tsx
+++ b/packages/components/src/DynamicDisclosure/DynamicDisclosure.tsx
@@ -27,7 +27,7 @@ interface DynamicDisclosureContextType {
   ) => Promise<void>;
   onClose: () => void;
   isOpen: boolean;
-  goBack: () => void;
+  goBack: (index?: number) => void;
   hasPrevious: boolean;
   formValues: Record<string, any>;
   allFormValues: Record<string, any>;
@@ -100,9 +100,17 @@ export const useDynamicDisclosure = () => {
     return Promise.resolve();
   };
 
-  const goBack = () => {
-    setCurrentIndex(current => current - 1);
-    stackRef.current.pop();
+  const goBack = (index = -1) => {
+    if (index >= 0 && index < stackRef.current.length) {
+      // Go to specific index
+      const itemsToRemove = stackRef.current.length - index - 1;
+      stackRef.current.splice(index + 1, itemsToRemove);
+      setCurrentIndex(index);
+    } else {
+      // Default behavior: go back one step
+      setCurrentIndex(current => current - 1);
+      stackRef.current.pop();
+    }
   };
 
   const currentItem = stackRef.current[currentIndex] || null;
diff --git a/packages/components/src/DynamicDisclosure/useMultiForm.test.tsx b/packages/components/src/DynamicDisclosure/useMultiForm.test.tsx
index 2edf3667f8..f6b18a96c0 100644
--- a/packages/components/src/DynamicDisclosure/useMultiForm.test.tsx
+++ b/packages/components/src/DynamicDisclosure/useMultiForm.test.tsx
@@ -18,7 +18,7 @@ const Page3 = () => {
   return (
     <ModalContent>
       <ModalHeader>
-        <Button onClick={goBack}>Go back</Button>
+        <Button onClick={() => goBack()}>Go back</Button>
       </ModalHeader>
       <ModalBody>
         <Box data-testid="all-form-values">{JSON.stringify(allFormValues)}</Box>
@@ -34,7 +34,7 @@ const Page2 = () => {
   return (
     <ModalContent>
       <ModalHeader>
-        <Button onClick={goBack}>Go back</Button>
+        <Button onClick={() => goBack()}>Go back</Button>
       </ModalHeader>
       <ModalBody>
         <form data-testid="form2" onSubmit={form.handleSubmit(() => openWith(<Page3 />))}>
diff --git a/packages/components/src/DynamicDisclosure/useMultiForm.tsx b/packages/components/src/DynamicDisclosure/useMultiForm.tsx
index da166abed0..bad2b7c077 100644
--- a/packages/components/src/DynamicDisclosure/useMultiForm.tsx
+++ b/packages/components/src/DynamicDisclosure/useMultiForm.tsx
@@ -19,6 +19,7 @@ export const useMultiForm = <
 
   const formDefaultValues =
     usedIn === "drawer" ? drawerContext?.formValues : modalContext?.formValues;
+
   const form = useForm<TFieldValues, TContext, TTransformedValues>({
     ...props,
     defaultValues: { ...props?.defaultValues, ...formDefaultValues } as any,