From 202980424b78354ea8347290ca9c8c9b59fc32a0 Mon Sep 17 00:00:00 2001
From: Carlos Martins <luacmartins@gmail.com>
Date: Tue, 17 Dec 2024 15:02:40 -0700
Subject: [PATCH] Revert "Revert "Replace navigation blur listeners with
 `useFocusEffect`""

---
 src/components/Button/index.tsx               |  5 ++
 src/components/ConfirmationPage.tsx           |  1 +
 src/hooks/useWaitForNavigation.ts             | 27 +++----
 src/pages/workspace/upgrade/UpgradeIntro.tsx  |  1 +
 .../upgrade/WorkspaceUpgradePage.tsx          | 29 ++++---
 tests/ui/WorkspaceUpgradeTest.tsx             | 76 +++++++++++++++++++
 6 files changed, 109 insertions(+), 30 deletions(-)
 create mode 100644 tests/ui/WorkspaceUpgradeTest.tsx

diff --git a/src/components/Button/index.tsx b/src/components/Button/index.tsx
index 07edd148778d..84767c6347e7 100644
--- a/src/components/Button/index.tsx
+++ b/src/components/Button/index.tsx
@@ -122,6 +122,9 @@ type ButtonProps = Partial<ChildrenProps> & {
     /** Id to use for this button */
     id?: string;
 
+    /** Used to locate this button in ui tests */
+    testID?: string;
+
     /** Accessibility label for the component */
     accessibilityLabel?: string;
 
@@ -237,6 +240,7 @@ function Button(
         shouldShowRightIcon = false,
 
         id = '',
+        testID = undefined,
         accessibilityLabel = '',
         isSplitButton = false,
         link = false,
@@ -405,6 +409,7 @@ function Button(
                 ]}
                 disabledStyle={disabledStyle}
                 id={id}
+                testID={testID}
                 accessibilityLabel={accessibilityLabel}
                 role={CONST.ROLE.BUTTON}
                 hoverDimmingValue={1}
diff --git a/src/components/ConfirmationPage.tsx b/src/components/ConfirmationPage.tsx
index a95cf9bf87d2..7b55f2317d46 100644
--- a/src/components/ConfirmationPage.tsx
+++ b/src/components/ConfirmationPage.tsx
@@ -82,6 +82,7 @@ function ConfirmationPage({
                         success
                         large
                         text={buttonText}
+                        testID="confirmation-button"
                         style={styles.mt6}
                         pressOnEnter
                         onPress={onButtonPress}
diff --git a/src/hooks/useWaitForNavigation.ts b/src/hooks/useWaitForNavigation.ts
index 73c0eb2bb14c..05981ec3322b 100644
--- a/src/hooks/useWaitForNavigation.ts
+++ b/src/hooks/useWaitForNavigation.ts
@@ -1,5 +1,5 @@
-import {useNavigation} from '@react-navigation/native';
-import {useEffect, useRef} from 'react';
+import {useFocusEffect} from '@react-navigation/native';
+import {useCallback, useRef} from 'react';
 
 type UseWaitForNavigation = (navigate: () => void) => () => Promise<void>;
 
@@ -8,21 +8,18 @@ type UseWaitForNavigation = (navigate: () => void) => () => Promise<void>;
  * Only use when navigating by react-navigation
  */
 export default function useWaitForNavigation(): UseWaitForNavigation {
-    const navigation = useNavigation();
     const resolvePromises = useRef<Array<() => void>>([]);
 
-    useEffect(() => {
-        const unsubscribeBlur = navigation.addListener('blur', () => {
-            resolvePromises.current.forEach((resolve) => {
-                resolve();
-            });
-            resolvePromises.current = [];
-        });
-
-        return () => {
-            unsubscribeBlur();
-        };
-    }, [navigation]);
+    useFocusEffect(
+        useCallback(() => {
+            return () => {
+                resolvePromises.current.forEach((resolve) => {
+                    resolve();
+                });
+                resolvePromises.current = [];
+            };
+        }, []),
+    );
 
     return (navigate: () => void) => () => {
         navigate();
diff --git a/src/pages/workspace/upgrade/UpgradeIntro.tsx b/src/pages/workspace/upgrade/UpgradeIntro.tsx
index d45e27905c28..029f8e78271f 100644
--- a/src/pages/workspace/upgrade/UpgradeIntro.tsx
+++ b/src/pages/workspace/upgrade/UpgradeIntro.tsx
@@ -75,6 +75,7 @@ function UpgradeIntro({feature, onUpgrade, buttonDisabled, loading, isCategorizi
                 <Button
                     isLoading={loading}
                     text={translate('common.upgrade')}
+                    testID="upgrade-button"
                     success
                     onPress={onUpgrade}
                     isDisabled={buttonDisabled}
diff --git a/src/pages/workspace/upgrade/WorkspaceUpgradePage.tsx b/src/pages/workspace/upgrade/WorkspaceUpgradePage.tsx
index 3bccc9c90eb0..5b9abd31b2d2 100644
--- a/src/pages/workspace/upgrade/WorkspaceUpgradePage.tsx
+++ b/src/pages/workspace/upgrade/WorkspaceUpgradePage.tsx
@@ -1,5 +1,5 @@
-import {useNavigation} from '@react-navigation/native';
-import React, {useCallback, useEffect} from 'react';
+import {useFocusEffect} from '@react-navigation/native';
+import React, {useCallback, useMemo} from 'react';
 import {useOnyx} from 'react-native-onyx';
 import HeaderWithBackButton from '@components/HeaderWithBackButton';
 import ScreenWrapper from '@components/ScreenWrapper';
@@ -37,20 +37,19 @@ function getFeatureNameAlias(featureName: string) {
 }
 
 function WorkspaceUpgradePage({route}: WorkspaceUpgradePageProps) {
-    const navigation = useNavigation();
     const styles = useThemeStyles();
     const policyID = route.params.policyID;
 
     const featureNameAlias = getFeatureNameAlias(route.params.featureName);
 
-    const feature = Object.values(CONST.UPGRADE_FEATURE_INTRO_MAPPING).find((f) => f.alias === featureNameAlias);
+    const feature = useMemo(() => Object.values(CONST.UPGRADE_FEATURE_INTRO_MAPPING).find((f) => f.alias === featureNameAlias), [featureNameAlias]);
     const {translate} = useLocalize();
     const [policy] = useOnyx(`policy_${policyID}`);
     const qboConfig = policy?.connections?.quickbooksOnline?.config;
     const {isOffline} = useNetwork();
 
     const canPerformUpgrade = !!feature && !!policy && PolicyUtils.isPolicyAdmin(policy);
-    const isUpgraded = React.useMemo(() => PolicyUtils.isControlPolicy(policy), [policy]);
+    const isUpgraded = useMemo(() => PolicyUtils.isControlPolicy(policy), [policy]);
 
     const perDiemCustomUnit = PolicyUtils.getPerDiemCustomUnit(policy);
     const categoryId = route.params?.categoryId;
@@ -155,16 +154,16 @@ function WorkspaceUpgradePage({route}: WorkspaceUpgradePageProps) {
         route.params.featureName,
     ]);
 
-    useEffect(() => {
-        const unsubscribeListener = navigation.addListener('blur', () => {
-            if (!isUpgraded || !canPerformUpgrade) {
-                return;
-            }
-            confirmUpgrade();
-        });
-
-        return unsubscribeListener;
-    }, [isUpgraded, canPerformUpgrade, confirmUpgrade, navigation]);
+    useFocusEffect(
+        useCallback(() => {
+            return () => {
+                if (!isUpgraded || !canPerformUpgrade) {
+                    return;
+                }
+                confirmUpgrade();
+            };
+        }, [isUpgraded, canPerformUpgrade, confirmUpgrade]),
+    );
 
     if (!canPerformUpgrade) {
         return <NotFoundPage />;
diff --git a/tests/ui/WorkspaceUpgradeTest.tsx b/tests/ui/WorkspaceUpgradeTest.tsx
new file mode 100644
index 000000000000..456ca3f737c8
--- /dev/null
+++ b/tests/ui/WorkspaceUpgradeTest.tsx
@@ -0,0 +1,76 @@
+/* eslint-disable @typescript-eslint/naming-convention */
+import {NavigationContainer} from '@react-navigation/native';
+import {act, fireEvent, render, screen} from '@testing-library/react-native';
+import React from 'react';
+import Onyx from 'react-native-onyx';
+import {WRITE_COMMANDS} from '@libs/API/types';
+import * as SequentialQueue from '@libs/Network/SequentialQueue';
+import createResponsiveStackNavigator from '@navigation/AppNavigator/createResponsiveStackNavigator';
+import type {SettingsNavigatorParamList} from '@navigation/types';
+import WorkspaceUpgradePage from '@pages/workspace/upgrade/WorkspaceUpgradePage';
+import ONYXKEYS from '@src/ONYXKEYS';
+import SCREENS from '@src/SCREENS';
+import type {Policy} from '@src/types/onyx';
+import * as LHNTestUtils from '../utils/LHNTestUtils';
+import * as TestHelper from '../utils/TestHelper';
+import waitForBatchedUpdates from '../utils/waitForBatchedUpdates';
+import waitForBatchedUpdatesWithAct from '../utils/waitForBatchedUpdatesWithAct';
+
+TestHelper.setupGlobalFetchMock();
+
+const RootStack = createResponsiveStackNavigator<SettingsNavigatorParamList>();
+
+const renderPage = (initialRouteName: typeof SCREENS.WORKSPACE.UPGRADE, initialParams: SettingsNavigatorParamList[typeof SCREENS.WORKSPACE.UPGRADE]) => {
+    return render(
+        <NavigationContainer>
+            <RootStack.Navigator initialRouteName={initialRouteName}>
+                <RootStack.Screen
+                    name={SCREENS.WORKSPACE.UPGRADE}
+                    component={WorkspaceUpgradePage}
+                    initialParams={initialParams}
+                />
+            </RootStack.Navigator>
+        </NavigationContainer>,
+    );
+};
+
+describe('WorkspaceUpgrade', () => {
+    beforeAll(() => {
+        Onyx.init({
+            keys: ONYXKEYS,
+        });
+    });
+
+    afterEach(async () => {
+        await SequentialQueue.waitForIdle();
+        await act(async () => {
+            await Onyx.clear();
+        });
+
+        jest.clearAllMocks();
+    });
+
+    it('should enable policy rules', async () => {
+        const policy: Policy = LHNTestUtils.getFakePolicy();
+
+        // Given that a policy is initialized in Onyx
+        await Onyx.merge(`${ONYXKEYS.COLLECTION.POLICY}${policy.id}`, policy);
+
+        // And WorkspaceUpgradePage for rules is opened
+        const {unmount} = renderPage(SCREENS.WORKSPACE.UPGRADE, {policyID: policy.id, featureName: 'rules'});
+
+        // When the policy is upgraded by clicking on the Upgrade button
+        fireEvent.press(screen.getByTestId('upgrade-button'));
+        await waitForBatchedUpdatesWithAct();
+
+        // Then "Upgrade to Corporate" API request should be made
+        TestHelper.expectAPICommandToHaveBeenCalled(WRITE_COMMANDS.UPGRADE_TO_CORPORATE, 1);
+
+        // When WorkspaceUpgradePage is unmounted
+        unmount();
+        await waitForBatchedUpdates();
+
+        // Then "Set policy rules enabled" API request should be made
+        TestHelper.expectAPICommandToHaveBeenCalled(WRITE_COMMANDS.SET_POLICY_RULES_ENABLED, 1);
+    });
+});