Replace useSafeGasEstimatePolling with usePolling

app/scripts/metamask-controller.js

@@ -3301,6 +3301,12 @@ export default class MetamaskController extends EventEmitter {
       ),
 
       // GasFeeController
+      gasFeeStartPollingByNetworkClientId:
+        gasFeeController.startPollingByNetworkClientId.bind(gasFeeController),
+
+      gasFeeStopPollingByPollingToken:
+        gasFeeController.stopPollingByPollingToken.bind(gasFeeController),
+
       getGasFeeEstimatesAndStartPolling:
         gasFeeController.getGasFeeEstimatesAndStartPolling.bind(
           gasFeeController,

development/ts-migration-dashboard/files-to-convert.json

@@ -1130,7 +1130,6 @@
   "ui/hooks/useIncrementedGasFees.js",
   "ui/hooks/useOriginMetadata.js",
   "ui/hooks/usePrevious.js",
-  "ui/hooks/useSafeGasEstimatePolling.js",
   "ui/hooks/useSegmentContext.js",
   "ui/hooks/useShouldAnimateGasEstimations.js",
   "ui/hooks/useShouldShowSpeedUp.js",

ui/hooks/useGasFeeEstimates.js

@@ -6,7 +6,9 @@ import {
   getIsGasEstimatesLoading,
   getIsNetworkBusy,
 } from '../ducks/metamask/metamask';
-import { useSafeGasEstimatePolling } from './useSafeGasEstimatePolling';
+import usePolling from './usePolling';
+import { gasFeeStartPollingByNetworkClientId, gasFeeStopPollingByPollingToken } from '../store/actions';
+import { getSelectedNetworkClientId } from '../selectors';
 
 /**
  * @typedef {object} GasEstimates
@@ -27,12 +29,18 @@ import { useSafeGasEstimatePolling } from './useSafeGasEstimatePolling';
  *
  * @returns {GasEstimates} GasEstimates object
  */
-export function useGasFeeEstimates() {
+export function useGasFeeEstimates(networkClientId) {
   const gasEstimateType = useSelector(getGasEstimateType);
   const gasFeeEstimates = useSelector(getGasFeeEstimates, isEqual);
   const isGasEstimatesLoading = useSelector(getIsGasEstimatesLoading);
   const isNetworkBusy = useSelector(getIsNetworkBusy);
-  useSafeGasEstimatePolling();
+  const selectedNetworkClientId = useSelector(getSelectedNetworkClientId);
+
+  usePolling({
+    startPollingByNetworkClientId: gasFeeStartPollingByNetworkClientId,
+    stopPollingByPollingToken: gasFeeStopPollingByPollingToken,
+    networkClientId: networkClientId ?? selectedNetworkClientId,
+  });
 
   return {
     gasFeeEstimates,

ui/hooks/useGasFeeEstimates.test.js

@@ -7,19 +7,17 @@ import {
   getGasFeeEstimates,
   getIsGasEstimatesLoading,
 } from '../ducks/metamask/metamask';
-import { checkNetworkAndAccountSupports1559 } from '../selectors';
+import { checkNetworkAndAccountSupports1559, getSelectedNetworkClientId } from '../selectors';
 import {
-  disconnectGasFeeEstimatePoller,
-  getGasFeeEstimatesAndStartPolling,
+  gasFeeStopPollingByPollingToken,
+  gasFeeStartPollingByNetworkClientId,
 } from '../store/actions';
 
 import { useGasFeeEstimates } from './useGasFeeEstimates';
 
 jest.mock('../store/actions', () => ({
-  disconnectGasFeeEstimatePoller: jest.fn(),
-  getGasFeeEstimatesAndStartPolling: jest.fn(),
-  addPollingTokenToAppState: jest.fn(),
-  removePollingTokenFromAppState: jest.fn(),
+  gasFeeStopPollingByPollingToken: jest.fn(),
+  gasFeeStartPollingByNetworkClientId: jest.fn(),
 }));
 
 jest.mock('react-redux', () => {
@@ -51,6 +49,9 @@ const generateUseSelectorRouter =
         DEFAULT_OPTS.checkNetworkAndAccountSupports1559
       );
     }
+    if (selector === getSelectedNetworkClientId) {
+      return 'selectedNetworkClientId';
+    }
     if (selector === getGasEstimateType) {
       return opts.gasEstimateType ?? DEFAULT_OPTS.gasEstimateType;
     }
@@ -68,12 +69,12 @@ describe('useGasFeeEstimates', () => {
   beforeEach(() => {
     jest.clearAllMocks();
     tokens = [];
-    getGasFeeEstimatesAndStartPolling.mockImplementation(() => {
+    gasFeeStartPollingByNetworkClientId.mockImplementation(() => {
       const token = createRandomId();
       tokens.push(token);
       return Promise.resolve(token);
     });
-    disconnectGasFeeEstimatePoller.mockImplementation((token) => {
+    gasFeeStopPollingByPollingToken.mockImplementation((token) => {
       tokens = tokens.filter((tkn) => tkn !== token);
     });
   });
@@ -90,11 +91,25 @@ describe('useGasFeeEstimates', () => {
     expect(tokens).toHaveLength(1);
     const expectedToken = tokens[0];
     await cleanup();
-    expect(getGasFeeEstimatesAndStartPolling).toHaveBeenCalledTimes(1);
-    expect(disconnectGasFeeEstimatePoller).toHaveBeenCalledWith(expectedToken);
+    expect(gasFeeStartPollingByNetworkClientId).toHaveBeenCalledTimes(1);
+    expect(gasFeeStopPollingByPollingToken).toHaveBeenCalledWith(expectedToken);
     expect(tokens).toHaveLength(0);
   });
 
+  it('polls the selected networkClientId by default', () => {
+    useSelector.mockImplementation(generateUseSelectorRouter());
+    renderHook(() => useGasFeeEstimates());
+    expect(gasFeeStartPollingByNetworkClientId).toHaveBeenCalledTimes(1);
+    expect(gasFeeStartPollingByNetworkClientId).toHaveBeenCalledWith('selectedNetworkClientId', undefined)
+  });
+
+  it('polls the passed in networkClientId when provided', () => {
+    useSelector.mockImplementation(generateUseSelectorRouter());
+    renderHook(() => useGasFeeEstimates('networkClientId1'));
+    expect(gasFeeStartPollingByNetworkClientId).toHaveBeenCalledTimes(1);
+    expect(gasFeeStartPollingByNetworkClientId).toHaveBeenCalledWith('networkClientId1', undefined)
+  });
+
   it('works with LEGACY gas prices', () => {
     useSelector.mockImplementation(
       generateUseSelectorRouter({

ui/hooks/usePolling.test.js

@@ -0,0 +1,84 @@
+import React from 'react';
+import { renderHook, cleanup } from '@testing-library/react-hooks';
+import { Provider } from 'react-redux';
+import configureStore from '../store/store';
+import usePolling from './usePolling';
+
+describe('usePolling', () => {
+  // eslint-disable-next-line jest/no-done-callback
+  it('calls startPollingByNetworkClientId and callback option args with polling token when component instantiating the hook mounts', (done) => {
+    const mockStart = jest.fn().mockImplementation(() => {
+      return Promise.resolve('pollingToken');
+    });
+    const mockStop = jest.fn();
+    const networkClientId = 'mainnet';
+    const options = {};
+    const mockState = {
+      metamask: {},
+    };
+
+    const wrapper = ({ children }) => (
+      <>
+        <Provider store={configureStore(mockState)}>{children}</Provider>
+      </>
+    );
+
+    renderHook(
+      () => {
+        usePolling({
+          callback: (pollingToken) => {
+            expect(mockStart).toHaveBeenCalledWith(networkClientId, options);
+            expect(pollingToken).toBeDefined();
+            done();
+            return (_pollingToken) => {
+              // noop
+            };
+          },
+          startPollingByNetworkClientId: mockStart,
+          stopPollingByPollingToken: mockStop,
+          networkClientId,
+          options,
+        });
+      },
+      { wrapper },
+    );
+  });
+  // eslint-disable-next-line jest/no-done-callback
+  it('calls the cleanup function with the correct pollingToken when unmounted', (done) => {
+    const mockStart = jest.fn().mockImplementation(() => {
+      return Promise.resolve('pollingToken');
+    });
+    const mockStop = jest.fn();
+    const networkClientId = 'mainnet';
+    const options = {};
+    const mockState = {
+      metamask: {},
+    };
+
+    const wrapper = ({ children }) => (
+      <>
+        <Provider store={configureStore(mockState)}>{children}</Provider>
+      </>
+    );
+
+    renderHook(
+      () => {
+        usePolling({
+          callback: () => {
+            return (_pollingToken) => {
+              expect(mockStop).toHaveBeenCalledWith(_pollingToken);
+              expect(_pollingToken).toBeDefined();
+              done();
+            };
+          },
+          startPollingByNetworkClientId: mockStart,
+          stopPollingByPollingToken: mockStop,
+          networkClientId,
+          options,
+        });
+      },
+      { wrapper },
+    );
+    cleanup();
+  });
+});

ui/hooks/usePolling.ts

@@ -0,0 +1,56 @@
+import { useEffect, useRef } from 'react';
+
+type UsePollingOptions = {
+  callback?: (pollingToken: string) => (pollingToken: string) => void;
+  startPollingByNetworkClientId: (
+    networkClientId: string,
+    options: any,
+  ) => Promise<string>;
+  stopPollingByPollingToken: (pollingToken: string) => void;
+  networkClientId: string;
+  options?: any;
+};
+
+const usePolling = (usePollingOptions: UsePollingOptions) => {
+  const pollTokenRef = useRef<null | string>(null);
+  const cleanupRef = useRef<null | ((pollingToken: string) => void)>(null);
+  let isMounted = false;
+  useEffect(() => {
+    isMounted = true;
+
+    const cleanup = () => {
+      if (pollTokenRef.current) {
+        usePollingOptions.stopPollingByPollingToken(pollTokenRef.current);
+        cleanupRef.current?.(pollTokenRef.current);
+      }
+    };
+    // Start polling when the component mounts
+    usePollingOptions
+      .startPollingByNetworkClientId(
+        usePollingOptions.networkClientId,
+        usePollingOptions.options,
+      )
+      .then((pollToken) => {
+        pollTokenRef.current = pollToken;
+        cleanupRef.current = usePollingOptions.callback?.(pollToken) || null;
+        if (!isMounted) {
+          cleanup();
+        }
+      });
+
+    // Return a cleanup function to stop polling when the component unmounts
+    return () => {
+      isMounted = false;
+      cleanup();
+    };
+  }, [
+    usePollingOptions.networkClientId,
+    usePollingOptions.options &&
+      JSON.stringify(
+        usePollingOptions.options,
+        Object.keys(usePollingOptions.options).sort(),
+      ),
+  ]);
+};
+
+export default usePolling;