Added stopAll and changed interface of stop to only take pollingToken

MetaMask · Sep 15, 2023 · 900c3d5 · 900c3d5
1 parent ee5d6f0
commit 900c3d5
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 49 deletions.
diff --git a/packages/gas-fee-controller/src/GasFeeControllerPolling.test.ts b/packages/gas-fee-controller/src/GasFeeControllerPolling.test.ts
@@ -22,7 +22,7 @@ describe('GasFeeControllerPolling', () => {
       });
       controller.start('mainnet');
       jest.advanceTimersByTime(1500);
-      controller.stop();
+      controller.stopAll();
       expect(executePollMock).toHaveBeenCalledTimes(1);
     });
   });
@@ -42,13 +42,10 @@ describe('GasFeeControllerPolling', () => {
       });
       const pollingToken = controller.start('mainnet');
       jest.advanceTimersByTime(1500);
-      controller.stop({
-        pollingToken,
-        networkClientId: 'mainnet',
-      });
+      controller.stop(pollingToken);
       jest.advanceTimersByTime(1500);
       expect(executePollMock).toHaveBeenCalledTimes(1);
-      controller.stop();
+      controller.stopAll();
     });
     it('should not stop polling if multiple polling tokens exist', () => {
       const executePollMock = jest.fn();
@@ -66,13 +63,10 @@ describe('GasFeeControllerPolling', () => {
       const pollingToken1 = controller.start('mainnet');
       controller.start('mainnet');
       jest.advanceTimersByTime(1200);
-      controller.stop({
-        pollingToken: pollingToken1,
-        networkClientId: 'mainnet',
-      });
+      controller.stop(pollingToken1);
       jest.advanceTimersByTime(1200);
       expect(executePollMock).toHaveBeenCalledTimes(2);
-      controller.stop();
+      controller.stopAll();
     });
     it('should error if no poll token is passed', () => {
       const executePollMock = jest.fn();
@@ -87,12 +81,30 @@ describe('GasFeeControllerPolling', () => {
         name: 'GasFeeControllerPolling',
         state: { foo: 'bar' },
       });
-      const pollingToken = controller.start('mainnet');
+      controller.start('mainnet');
+      expect(() => {
+        controller.stop(undefined as unknown as any);
+      }).toThrow('pollingToken required');
+      controller.stopAll();
+    });
+    it('should error if no poll token is found', () => {
+      const executePollMock = jest.fn();
+      class MyGasFeeController extends GasFeeControllerPolling<any, any, any> {
+        executePoll = executePollMock;
+      }
+      const mockMessenger = new ControllerMessenger<any, any>();
+
+      const controller = new MyGasFeeController({
+        messenger: mockMessenger,
+        metadata: {},
+        name: 'GasFeeControllerPolling',
+        state: { foo: 'bar' },
+      });
+      controller.start('mainnet');
       expect(() => {
-        controller.stop({
-          pollingToken,
-        });
-      }).toThrow('networkClientId is required when pollingToken is passed');
+        controller.stop('potato');
+      }).toThrow('pollingToken not found');
+      controller.stopAll();
     });
   });
   describe('poll', () => {
@@ -134,7 +146,7 @@ describe('GasFeeControllerPolling', () => {
       controller.start('mainnet');
       jest.advanceTimersByTime(2500);
       expect(executePollMock).toHaveBeenCalledTimes(2);
-      controller.stop();
+      controller.stopAll();
     });
     it('should publish polligComplete when stop is called', async () => {
       const executePollMock = jest.fn();
@@ -158,7 +170,7 @@ describe('GasFeeControllerPolling', () => {
         state: { foo: 'bar' },
       });
       const pollingToken = controller.start('mainnet');
-      controller.stop({ pollingToken, networkClientId: 'mainnet' });
+      controller.stop(pollingToken);
       expect(pollingComplete).toHaveBeenCalledTimes(1);
     });
   });
@@ -180,7 +192,7 @@ describe('GasFeeControllerPolling', () => {
       controller.start('rinkeby');
       jest.advanceTimersByTime(2200);
       expect(executePollMock).toHaveBeenCalledTimes(4);
-      controller.stop();
+      controller.stopAll();
     });
   });
 });
diff --git a/packages/gas-fee-controller/src/GasFeeControllerPolling.ts b/packages/gas-fee-controller/src/GasFeeControllerPolling.ts
@@ -46,38 +46,38 @@ export default abstract class GasFeeControllerPolling<
     return innerPollToken;
   }
 
-  stop({
-    pollingToken,
-    networkClientId,
-  }: {
-    pollingToken?: string;
-    networkClientId?: NetworkClientId;
-  } = {}) {
-    if (pollingToken && !networkClientId) {
-      throw new Error(
-        'networkClientId is required when pollingToken is passed',
-      );
-    }
-    if (!pollingToken || !networkClientId) {
-      this.networkClientIdTokensMap.forEach((tokens, _networkClientId) => {
-        tokens.forEach((token) => {
-          this.stop({
-            pollingToken: token,
-            networkClientId: _networkClientId,
-          });
-        });
+  stopAll() {
+    this.networkClientIdTokensMap.forEach((tokens, _networkClientId) => {
+      tokens.forEach((token) => {
+        this.stop(token);
       });
-      return;
+    });
+  }
+
+  stop(pollingToken: string) {
+    if (!pollingToken) {
+      throw new Error('pollingToken required');
     }
-    this.networkClientIdTokensMap.get(networkClientId)?.delete(pollingToken);
-    if (this.networkClientIdTokensMap.get(networkClientId)?.size === 0) {
-      clearInterval(this.intervalIds[networkClientId]);
-      delete this.intervalIds[networkClientId];
-      this.networkClientIdTokensMap.delete(networkClientId);
-      this.messagingSystem.publish(
-        `${this.name}:pollingComplete`,
-        networkClientId,
-      );
+    let found = false;
+    this.networkClientIdTokensMap.forEach((tokens, networkClientId) => {
+      if (tokens.has(pollingToken)) {
+        found = true;
+        this.networkClientIdTokensMap
+          .get(networkClientId)
+          ?.delete(pollingToken);
+        if (this.networkClientIdTokensMap.get(networkClientId)?.size === 0) {
+          clearInterval(this.intervalIds[networkClientId]);
+          delete this.intervalIds[networkClientId];
+          this.networkClientIdTokensMap.delete(networkClientId);
+          this.messagingSystem.publish(
+            `${this.name}:pollingComplete`,
+            networkClientId,
+          );
+        }
+      }
+    });
+    if (!found) {
+      throw new Error('pollingToken not found');
     }
   }