address feedback

packages/assets-controllers/src/CurrencyRateController.ts

@@ -11,7 +11,6 @@ import type {
   NetworkClientId,
   NetworkControllerGetNetworkClientByIdAction,
 } from '@metamask/network-controller';
-import type { IPollingController } from '@metamask/polling-controller';
 import { StaticIntervalPollingController } from '@metamask/polling-controller';
 import { Mutex } from 'async-mutex';
 
@@ -83,14 +82,11 @@ const defaultState = {
  * Controller that passively polls on a set interval for an exchange rate from the current network
  * asset to the user's preferred currency.
  */
-export class CurrencyRateController
-  extends StaticIntervalPollingController<
-    typeof name,
-    CurrencyRateState,
-    CurrencyRateMessenger
-  >
-  implements IPollingController
-{
+export class CurrencyRateController extends StaticIntervalPollingController<
+  typeof name,
+  CurrencyRateState,
+  CurrencyRateMessenger
+> {
   private readonly mutex = new Mutex();
 
   private readonly fetchExchangeRate;

packages/assets-controllers/src/NftDetectionController.ts

@@ -11,7 +11,6 @@ import type {
   NetworkState,
   NetworkClient,
 } from '@metamask/network-controller';
-import type { IPollingController } from '@metamask/polling-controller';
 import { StaticIntervalPollingControllerV1 } from '@metamask/polling-controller';
 import type { PreferencesState } from '@metamask/preferences-controller';
 import type { Hex } from '@metamask/utils';
@@ -148,10 +147,10 @@ export interface NftDetectionConfig extends BaseConfig {
 /**
  * Controller that passively polls on a set interval for NFT auto detection
  */
-export class NftDetectionController
-  extends StaticIntervalPollingControllerV1<NftDetectionConfig, BaseState>
-  implements IPollingController
-{
+export class NftDetectionController extends StaticIntervalPollingControllerV1<
+  NftDetectionConfig,
+  BaseState
+> {
   private intervalId?: ReturnType<typeof setTimeout>;
 
   private getOwnerNftApi({

packages/assets-controllers/src/TokenDetectionController.ts

@@ -8,7 +8,6 @@ import type {
   NetworkController,
   NetworkState,
 } from '@metamask/network-controller';
-import type { IPollingController } from '@metamask/polling-controller';
 import { StaticIntervalPollingControllerV1 } from '@metamask/polling-controller';
 import type { PreferencesState } from '@metamask/preferences-controller';
 import type { Hex } from '@metamask/utils';
@@ -45,10 +44,10 @@ export interface TokenDetectionConfig extends BaseConfig {
 /**
  * Controller that passively polls on a set interval for Tokens auto detection
  */
-export class TokenDetectionController
-  extends StaticIntervalPollingControllerV1<TokenDetectionConfig, BaseState>
-  implements IPollingController
-{
+export class TokenDetectionController extends StaticIntervalPollingControllerV1<
+  TokenDetectionConfig,
+  BaseState
+> {
   private intervalId?: ReturnType<typeof setTimeout>;
 
   /**

packages/assets-controllers/src/TokenListController.ts

@@ -10,7 +10,6 @@ import type {
   NetworkState,
   NetworkControllerGetNetworkClientByIdAction,
 } from '@metamask/network-controller';
-import type { IPollingController } from '@metamask/polling-controller';
 import { StaticIntervalPollingController } from '@metamask/polling-controller';
 import type { Hex } from '@metamask/utils';
 import { Mutex } from 'async-mutex';
@@ -92,14 +91,11 @@ const defaultState: TokenListState = {
 /**
  * Controller that passively polls on a set interval for the list of tokens from metaswaps api
  */
-export class TokenListController
-  extends StaticIntervalPollingController<
-    typeof name,
-    TokenListState,
-    TokenListMessenger
-  >
-  implements IPollingController
-{
+export class TokenListController extends StaticIntervalPollingController<
+  typeof name,
+  TokenListState,
+  TokenListMessenger
+> {
   private readonly mutex = new Mutex();
 
   private intervalId?: ReturnType<typeof setTimeout>;

packages/assets-controllers/src/TokenRatesController.ts

@@ -10,7 +10,6 @@ import type {
   NetworkController,
   NetworkState,
 } from '@metamask/network-controller';
-import type { IPollingController } from '@metamask/polling-controller';
 import { StaticIntervalPollingControllerV1 } from '@metamask/polling-controller';
 import type { PreferencesState } from '@metamask/preferences-controller';
 import type { Hex } from '@metamask/utils';
@@ -137,10 +136,10 @@ async function getCurrencyConversionRate({
  * Controller that passively polls on a set interval for token-to-fiat exchange rates
  * for tokens stored in the TokensController
  */
-export class TokenRatesController
-  extends StaticIntervalPollingControllerV1<TokenRatesConfig, TokenRatesState>
-  implements IPollingController
-{
+export class TokenRatesController extends StaticIntervalPollingControllerV1<
+  TokenRatesConfig,
+  TokenRatesState
+> {
   private handle?: ReturnType<typeof setTimeout>;
 
   #pollState = PollState.Inactive;

packages/gas-fee-controller/src/GasFeeController.ts

@@ -18,7 +18,6 @@ import type {
   NetworkState,
   ProviderProxy,
 } from '@metamask/network-controller';
-import type { IPollingController } from '@metamask/polling-controller';
 import { StaticIntervalPollingController } from '@metamask/polling-controller';
 import type { Hex } from '@metamask/utils';
 import { v1 as random } from 'uuid';
@@ -254,14 +253,11 @@ const defaultState: GasFeeState = {
 /**
  * Controller that retrieves gas fee estimate data and polls for updated data on a set interval
  */
-export class GasFeeController
-  extends StaticIntervalPollingController<
-    typeof name,
-    GasFeeState,
-    GasFeeMessenger
-  >
-  implements IPollingController
-{
+export class GasFeeController extends StaticIntervalPollingController<
+  typeof name,
+  GasFeeState,
+  GasFeeMessenger
+> {
   private intervalId?: ReturnType<typeof setTimeout>;
 
   private readonly intervalDelay;

packages/polling-controller/src/AbstractPollingController.ts

@@ -75,7 +75,7 @@ export function AbstractPollingControllerBaseMixin<TBase extends Constructor>(
       const pollToken = random();
       const key = getKey(networkClientId, options);
       const pollingTokenSet =
-        this.#pollingTokenSets.get(key) || new Set<string>();
+        this.#pollingTokenSets.get(key) ?? new Set<string>();
       pollingTokenSet.add(pollToken);
       this.#pollingTokenSets.set(key, pollingTokenSet);
 
@@ -110,15 +110,16 @@ export function AbstractPollingControllerBaseMixin<TBase extends Constructor>(
       }
 
       if (keyToDelete) {
-        // TODO figure out why this is necessary
-        const nonNullKey = keyToDelete;
-        this._stopPollingByPollingTokenSetId(nonNullKey);
-        this.#pollingTokenSets.delete(nonNullKey);
-        this.#callbacks.get(nonNullKey)?.forEach((callback) => {
-          // for some reason this typescript can't tell that keyToDelete is not null here
-          callback(nonNullKey);
-        });
-        this.#callbacks.get(nonNullKey)?.clear();
+        this._stopPollingByPollingTokenSetId(keyToDelete);
+        this.#pollingTokenSets.delete(keyToDelete);
+        const callbacks = this.#callbacks.get(keyToDelete);
+        if (callbacks) {
+          for (const callback of callbacks) {
+            // eslint-disable-next-line n/callback-return
+            callback(keyToDelete);
+          }
+          callbacks.clear();
+        }
       }
     }
 
@@ -128,7 +129,7 @@ export function AbstractPollingControllerBaseMixin<TBase extends Constructor>(
       options: Json = {},
     ) {
       const key = getKey(networkClientId, options);
-      const callbacks = this.#callbacks.get(key) || new Set<typeof callback>();
+      const callbacks = this.#callbacks.get(key) ?? new Set<typeof callback>();
       callbacks.add(callback);
       this.#callbacks.set(key, callbacks);
     }

packages/polling-controller/src/BlockTrackerPollingController.test.ts

@@ -254,4 +254,15 @@ describe('BlockTrackerPollingController', () => {
       controller.stopAllPolling();
     });
   });
+
+  describe('onPollingCompleteByNetworkClientId', () => {
+    it('should publish "pollingComplete" callback function set by "onPollingCompleteByNetworkClientId" when polling stops', async () => {
+      const pollingComplete: any = jest.fn();
+      controller.onPollingCompleteByNetworkClientId('mainnet', pollingComplete);
+      const pollingToken = controller.startPollingByNetworkClientId('mainnet');
+      controller.stopPollingByPollingToken(pollingToken);
+      expect(pollingComplete).toHaveBeenCalledTimes(1);
+      expect(pollingComplete).toHaveBeenCalledWith('mainnet:{}');
+    });
+  });
 });

packages/polling-controller/src/StaticIntervalPollingController.test.ts

@@ -4,7 +4,7 @@ import { useFakeTimers } from 'sinon';
 import { advanceTime } from '../../../tests/helpers';
 import { StaticIntervalPollingController } from './StaticIntervalPollingController';
 
-const TICK_TIME = 1000;
+const TICK_TIME = 5;
 
 const createExecutePollMock = () => {
   const executePollMock = jest.fn().mockImplementation(async () => {
@@ -19,8 +19,6 @@ class ChildBlockTrackerPollingController extends StaticIntervalPollingController
   any
 > {
   _executePoll = createExecutePollMock();
-
-  _intervalLength = TICK_TIME;
 }
 
 describe('StaticIntervalPollingController', () => {
@@ -35,29 +33,24 @@ describe('StaticIntervalPollingController', () => {
       name: 'PollingController',
       state: { foo: 'bar' },
     });
+    controller.setIntervalLength(TICK_TIME);
     clock = useFakeTimers();
   });
   afterEach(() => {
     clock.restore();
   });
 
   describe('startPollingByNetworkClientId', () => {
-    it('should start polling if not polling', async () => {
+    it('should start polling if not already polling', async () => {
       controller.startPollingByNetworkClientId('mainnet');
       await advanceTime({ clock, duration: 0 });
       expect(controller._executePoll).toHaveBeenCalledTimes(1);
       await advanceTime({ clock, duration: TICK_TIME });
       expect(controller._executePoll).toHaveBeenCalledTimes(2);
       controller.stopAllPolling();
     });
-    it('should call _executePoll immediately and on interval if polling', async () => {
-      controller.startPollingByNetworkClientId('mainnet');
-      await advanceTime({ clock, duration: 0 });
-      expect(controller._executePoll).toHaveBeenCalledTimes(1);
-      await advanceTime({ clock, duration: TICK_TIME * 2 });
-      expect(controller._executePoll).toHaveBeenCalledTimes(3);
-    });
-    it('should call _executePoll immediately once and continue calling _executePoll on interval when start is called again with the same networkClientId', async () => {
+
+    it('should call _executePoll immediately once and continue calling _executePoll on interval when called again with the same networkClientId', async () => {
       controller.startPollingByNetworkClientId('mainnet');
       await advanceTime({ clock, duration: 0 });
 
@@ -184,20 +177,11 @@ describe('StaticIntervalPollingController', () => {
     it('should error if no pollingToken is passed', () => {
       controller.startPollingByNetworkClientId('mainnet');
       expect(() => {
-        controller.stopPollingByPollingToken(undefined as unknown as any);
+        controller.stopPollingByPollingToken();
       }).toThrow('pollingToken required');
       controller.stopAllPolling();
     });
 
-    it('should publish "pollingComplete" when stop is called', async () => {
-      const pollingComplete: any = jest.fn();
-      controller.onPollingCompleteByNetworkClientId('mainnet', pollingComplete);
-      const pollingToken = controller.startPollingByNetworkClientId('mainnet');
-      controller.stopPollingByPollingToken(pollingToken);
-      expect(pollingComplete).toHaveBeenCalledTimes(1);
-      expect(pollingComplete).toHaveBeenCalledWith('mainnet:{}');
-    });
-
     it('should start and stop polling sessions for different networkClientIds with the same options', async () => {
       controller.setIntervalLength(TICK_TIME);
       const pollToken1 = controller.startPollingByNetworkClientId('mainnet', {
@@ -239,4 +223,15 @@ describe('StaticIntervalPollingController', () => {
       ]);
     });
   });
+
+  describe('onPollingCompleteByNetworkClientId', () => {
+    it('should publish "pollingComplete" callback function set by "onPollingCompleteByNetworkClientId" when polling stops', async () => {
+      const pollingComplete: any = jest.fn();
+      controller.onPollingCompleteByNetworkClientId('mainnet', pollingComplete);
+      const pollingToken = controller.startPollingByNetworkClientId('mainnet');
+      controller.stopPollingByPollingToken(pollingToken);
+      expect(pollingComplete).toHaveBeenCalledTimes(1);
+      expect(pollingComplete).toHaveBeenCalledWith('mainnet:{}');
+    });
+  });
 });

packages/polling-controller/src/StaticIntervalPollingController.ts

@@ -29,21 +29,21 @@ function StaticIntervalPollingControllerMixin<TBase extends Constructor>(
   {
     readonly #intervalIds: Record<PollingTokenSetId, NodeJS.Timeout> = {};
 
-    intervalLength: number | undefined = 1000;
+    #intervalLength: number | undefined = 1000;
 
     setIntervalLength(intervalLength: number) {
-      this.intervalLength = intervalLength;
+      this.#intervalLength = intervalLength;
     }
 
     getIntervalLength() {
-      return this.intervalLength;
+      return this.#intervalLength;
     }
 
     _startPollingByNetworkClientId(
       networkClientId: NetworkClientId,
       options: Json,
     ) {
-      if (!this.intervalLength) {
+      if (!this.#intervalLength) {
         throw new Error('intervalLength must be defined and greater than 0');
       }
 
@@ -60,7 +60,7 @@ function StaticIntervalPollingControllerMixin<TBase extends Constructor>(
           }
           this._startPollingByNetworkClientId(networkClientId, options);
         },
-        existingInterval ? this.intervalLength : 0,
+        existingInterval ? this.#intervalLength : 0,
       );
     }