From 6680af4a35e4772bc60a90dd72ee0b347ec9d4f7 Mon Sep 17 00:00:00 2001
From: bartektomczyk <bartlomiejtomczyk@gmail.com>
Date: Thu, 11 Jan 2024 14:34:49 +0100
Subject: [PATCH] fix: move to the last step when editing data

---
 src/hooks/useSubStep/index.ts                              | 7 ++++++-
 .../BeneficialOwnerInfo/BeneficialOwnerInfo.tsx            | 6 ++++++
 .../ReimbursementAccount/BusinessInfo/BusinessInfo.tsx     | 7 ++++++-
 .../CompleteVerification/CompleteVerification.tsx          | 7 ++++++-
 .../ReimbursementAccount/PersonalInfo/PersonalInfo.tsx     | 7 ++++++-
 5 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/hooks/useSubStep/index.ts b/src/hooks/useSubStep/index.ts
index 0efd875998e2..f7d1ed273959 100644
--- a/src/hooks/useSubStep/index.ts
+++ b/src/hooks/useSubStep/index.ts
@@ -42,5 +42,10 @@ export default function useSubStep<T>({bodyContent, onFinished, startFrom = 0}:
         setScreenIndex(0);
     }, []);
 
-    return {componentToRender: bodyContent[screenIndex], isEditing: isEditing.current, screenIndex, prevScreen, nextScreen, moveTo, resetScreenIndex};
+    const goToTheLastStep = useCallback(() => {
+        isEditing.current = false;
+        setScreenIndex(bodyContent.length - 1);
+    }, [bodyContent]);
+
+    return {componentToRender: bodyContent[screenIndex], isEditing: isEditing.current, screenIndex, prevScreen, nextScreen, moveTo, resetScreenIndex, goToTheLastStep};
 }
diff --git a/src/pages/ReimbursementAccount/BeneficialOwnerInfo/BeneficialOwnerInfo.tsx b/src/pages/ReimbursementAccount/BeneficialOwnerInfo/BeneficialOwnerInfo.tsx
index e5145ef3ad55..8037820a55fb 100644
--- a/src/pages/ReimbursementAccount/BeneficialOwnerInfo/BeneficialOwnerInfo.tsx
+++ b/src/pages/ReimbursementAccount/BeneficialOwnerInfo/BeneficialOwnerInfo.tsx
@@ -122,6 +122,7 @@ function BeneficialOwnerInfo({reimbursementAccount, reimbursementAccountDraft, o
         prevScreen,
         moveTo,
         resetScreenIndex,
+        goToTheLastStep,
     } = useSubStep<{beneficialOwnerBeingModifiedID: string; setBeneficialOwnerBeingModifiedID?: (id: string) => void}>({
         bodyContent: BODY_CONTENT,
         startFrom: 0,
@@ -192,6 +193,11 @@ function BeneficialOwnerInfo({reimbursementAccount, reimbursementAccountDraft, o
     };
 
     const handleBackButtonPress = () => {
+        if (isEditing) {
+            goToTheLastStep();
+            return;
+        }
+
         // User goes back to previous step
         if (currentUBOSubstep === SUBSTEP.IS_USER_UBO) {
             onBackButtonPress();
diff --git a/src/pages/ReimbursementAccount/BusinessInfo/BusinessInfo.tsx b/src/pages/ReimbursementAccount/BusinessInfo/BusinessInfo.tsx
index 3bf0b7180664..27941ab25612 100644
--- a/src/pages/ReimbursementAccount/BusinessInfo/BusinessInfo.tsx
+++ b/src/pages/ReimbursementAccount/BusinessInfo/BusinessInfo.tsx
@@ -88,9 +88,14 @@ function BusinessInfo({reimbursementAccount, reimbursementAccountDraft, policyID
 
     const startFrom = useMemo(() => getInitialSubstepForBusinessInfo(values), [values]);
 
-    const {componentToRender: SubStep, isEditing, screenIndex, nextScreen, prevScreen, moveTo} = useSubStep({bodyContent, startFrom, onFinished: submit});
+    const {componentToRender: SubStep, isEditing, screenIndex, nextScreen, prevScreen, moveTo, goToTheLastStep} = useSubStep({bodyContent, startFrom, onFinished: submit});
 
     const handleBackButtonPress = () => {
+        if (isEditing) {
+            goToTheLastStep();
+            return;
+        }
+
         if (screenIndex === 0) {
             onBackButtonPress();
         } else {
diff --git a/src/pages/ReimbursementAccount/CompleteVerification/CompleteVerification.tsx b/src/pages/ReimbursementAccount/CompleteVerification/CompleteVerification.tsx
index 6ba89c0ec3e4..e8fb9b487fa8 100644
--- a/src/pages/ReimbursementAccount/CompleteVerification/CompleteVerification.tsx
+++ b/src/pages/ReimbursementAccount/CompleteVerification/CompleteVerification.tsx
@@ -58,9 +58,14 @@ function CompleteVerification({reimbursementAccount, reimbursementAccountDraft,
         });
     }, [reimbursementAccount, reimbursementAccountDraft, values]);
 
-    const {componentToRender: SubStep, isEditing, screenIndex, nextScreen, prevScreen, moveTo} = useSubStep({bodyContent: BODY_CONTENT, startFrom: 0, onFinished: submit});
+    const {componentToRender: SubStep, isEditing, screenIndex, nextScreen, prevScreen, moveTo, goToTheLastStep} = useSubStep({bodyContent: BODY_CONTENT, startFrom: 0, onFinished: submit});
 
     const handleBackButtonPress = () => {
+        if (isEditing) {
+            goToTheLastStep();
+            return;
+        }
+
         if (screenIndex === 0) {
             onBackButtonPress();
         } else {
diff --git a/src/pages/ReimbursementAccount/PersonalInfo/PersonalInfo.tsx b/src/pages/ReimbursementAccount/PersonalInfo/PersonalInfo.tsx
index 0c67773dc096..038917978865 100644
--- a/src/pages/ReimbursementAccount/PersonalInfo/PersonalInfo.tsx
+++ b/src/pages/ReimbursementAccount/PersonalInfo/PersonalInfo.tsx
@@ -55,9 +55,14 @@ function PersonalInfo({reimbursementAccount, reimbursementAccountDraft, onBackBu
     }, [values]);
     const startFrom = useMemo(() => getInitialSubstepForPersonalInfo(values), [values]);
 
-    const {componentToRender: SubStep, isEditing, screenIndex, nextScreen, prevScreen, moveTo} = useSubStep({bodyContent, startFrom, onFinished: submit});
+    const {componentToRender: SubStep, isEditing, screenIndex, nextScreen, prevScreen, moveTo, goToTheLastStep} = useSubStep({bodyContent, startFrom, onFinished: submit});
 
     const handleBackButtonPress = () => {
+        if (isEditing) {
+            goToTheLastStep();
+            return;
+        }
+
         if (screenIndex === 0) {
             onBackButtonPress();
         } else {