Changed onPollingComplete to be by networkClientId

MetaMask · Sep 27, 2023 · 82255b8 · 82255b8
1 parent ddf7c70
commit 82255b8
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 7 deletions.
diff --git a/packages/polling-controller/src/PollingController.test.ts b/packages/polling-controller/src/PollingController.test.ts
@@ -175,14 +175,13 @@ describe('PollingController', () => {
         PollingCompleteType<typeof name>
       >();
 
-
       const controller = new MyGasFeeController({
         messenger: mockMessenger,
         metadata: {},
         name,
         state: { foo: 'bar' },
       });
-      controller.onPollingComplete(pollingComplete);
+      controller.onPollingComplete('mainnet', pollingComplete);
       const pollingToken = controller.start('mainnet');
       controller.stop(pollingToken);
       expect(pollingComplete).toHaveBeenCalledTimes(1);

diff --git a/packages/polling-controller/src/PollingController.ts b/packages/polling-controller/src/PollingController.ts
@@ -28,7 +28,10 @@ function PollingControllerMixin<TBase extends Constructor>(Base: TBase) {
 
     readonly #intervalIds: Record<NetworkClientId, NodeJS.Timeout> = {};
 
-    #callbacks: Set<(networkClientId: NetworkClientId) => void> = new Set();
+    #callbacks: Map<
+      NetworkClientId,
+      Set<(networkClientId: NetworkClientId) => void>
+    > = new Map();
 
     #intervalLength = 1000;
 
@@ -91,10 +94,10 @@ function PollingControllerMixin<TBase extends Constructor>(Base: TBase) {
             clearTimeout(this.#intervalIds[networkClientId]);
             delete this.#intervalIds[networkClientId];
             this.#networkClientIdTokensMap.delete(networkClientId);
-            this.#callbacks.forEach((callback) => {
+            this.#callbacks.get(networkClientId)?.forEach((callback) => {
               callback(networkClientId);
             });
-            this.#callbacks.clear();
+            this.#callbacks.get(networkClientId)?.clear();
           }
         }
       });
@@ -125,8 +128,17 @@ function PollingControllerMixin<TBase extends Constructor>(Base: TBase) {
       }, this.#intervalLength);
     }
 
-    onPollingComplete(callback: (networkClientId: NetworkClientId) => void) {
-      this.#callbacks.add(callback);
+    onPollingComplete(
+      networkClientId: NetworkClientId,
+      callback: (networkClientId: NetworkClientId) => void,
+    ) {
+      if (this.#callbacks.has(networkClientId)) {
+        this.#callbacks.get(networkClientId)?.add(callback);
+      } else {
+        const set = new Set<typeof callback>();
+        set.add(callback);
+        this.#callbacks.set(networkClientId, set);
+      }
     }
   }
   return PollingControllerBase;