adding tests

Author: adonesky1

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

@@ -1,4 +1,5 @@
 import { ControllerMessenger } from '@metamask/base-controller';
+import EventEmitter from 'events';
 import { useFakeTimers } from 'sinon';
 
 import { advanceTime } from '../../../tests/helpers';
@@ -277,30 +278,30 @@ describe('PollingController', () => {
       controller.startPollingByNetworkClientId('mainnet');
       await advanceTime({ clock, duration: 0 });
 
-      controller.startPollingByNetworkClientId('rinkeby');
+      controller.startPollingByNetworkClientId('goerli');
       await advanceTime({ clock, duration: 0 });
 
       expect(controller._executePoll.mock.calls).toMatchObject([
         ['mainnet', {}],
-        ['rinkeby', {}],
+        ['goerli', {}],
       ]);
       await advanceTime({ clock, duration: TICK_TIME });
 
       expect(controller._executePoll.mock.calls).toMatchObject([
         ['mainnet', {}],
-        ['rinkeby', {}],
+        ['goerli', {}],
         ['mainnet', {}],
-        ['rinkeby', {}],
+        ['goerli', {}],
       ]);
       await advanceTime({ clock, duration: TICK_TIME });
 
       expect(controller._executePoll.mock.calls).toMatchObject([
         ['mainnet', {}],
-        ['rinkeby', {}],
+        ['goerli', {}],
         ['mainnet', {}],
-        ['rinkeby', {}],
+        ['goerli', {}],
         ['mainnet', {}],
-        ['rinkeby', {}],
+        ['goerli', {}],
       ]);
       controller.stopAllPolling();
     });
@@ -383,4 +384,145 @@ describe('PollingController', () => {
       expect(c.stopPollingByPollingToken).toBeDefined();
     });
   });
+  describe('startPollingByNetworkClientId after setPollOnNewBlocks', () => {
+    class TestBlockTracker extends EventEmitter {
+      private latestBlockNumber: number;
+
+      constructor({ interval } = { interval: 1000 }) {
+        super();
+        this.latestBlockNumber = 0;
+        this.start(interval);
+      }
+
+      private start(interval: number) {
+        setInterval(() => {
+          this.latestBlockNumber += 1;
+          this.emit('latest', this.latestBlockNumber);
+        }, interval);
+      }
+    }
+
+    it('should start polling for the specified networkClientId', async () => {
+      class MyGasFeeController extends PollingController<any, any, any> {
+        _executePoll = createExecutePollMock();
+      }
+      const mockMessenger = new ControllerMessenger<any, any>();
+
+      const controller = new MyGasFeeController({
+        messenger: mockMessenger,
+        metadata: {},
+        name: 'PollingController',
+        state: { foo: 'bar' },
+      });
+
+      const getNetworkClientById = jest.fn().mockReturnValue({
+        blockTracker: new TestBlockTracker({ interval: 5 }),
+      });
+
+      controller.setPollOnNewBlocks(getNetworkClientById);
+
+      controller.startPollingByNetworkClientId('mainnet');
+
+      expect(getNetworkClientById).toHaveBeenCalledWith('mainnet');
+
+      await advanceTime({ clock, duration: 5 });
+
+      expect(controller._executePoll).toHaveBeenCalledTimes(1);
+
+      await advanceTime({ clock, duration: 1 });
+
+      expect(controller._executePoll).toHaveBeenCalledTimes(1);
+
+      await advanceTime({ clock, duration: 4 });
+
+      expect(controller._executePoll.mock.calls).toMatchObject([
+        expect.arrayContaining(['mainnet', {}]),
+        expect.arrayContaining(['mainnet', {}]),
+      ]);
+
+      // Stop all polling
+      controller.stopAllPolling();
+    });
+
+    it('should start polling with different intervals for each networkClientId', async () => {
+      class MyGasFeeController extends PollingController<any, any, any> {
+        _executePoll = createExecutePollMock();
+      }
+      const mockMessenger = new ControllerMessenger<any, any>();
+
+      const controller = new MyGasFeeController({
+        messenger: mockMessenger,
+        metadata: {},
+        name: 'PollingController',
+        state: { foo: 'bar' },
+      });
+
+      const getNetworkClientById = jest
+        .fn()
+        .mockImplementation((networkClientId) => {
+          switch (networkClientId) {
+            case 'mainnet':
+              return {
+                blockTracker: new TestBlockTracker({ interval: 5 }),
+              };
+            case 'goerli':
+              return {
+                blockTracker: new TestBlockTracker({ interval: 10 }),
+              };
+            case 'sepolia':
+              return {
+                blockTracker: new TestBlockTracker({ interval: 15 }),
+              };
+            default:
+              throw new Error(`Unknown networkClientId: ${networkClientId}`);
+          }
+        });
+
+      controller.setPollOnNewBlocks(getNetworkClientById);
+
+      controller.startPollingByNetworkClientId('mainnet');
+      await advanceTime({ clock, duration: 5 });
+
+      expect(controller._executePoll).toHaveBeenCalledTimes(1);
+      expect(controller._executePoll).toHaveBeenCalledWith('mainnet', {}, 1);
+
+      // Start polling for goerli, 10ms interval
+      controller.startPollingByNetworkClientId('goerli');
+      await advanceTime({ clock, duration: 5 });
+
+      expect(controller._executePoll.mock.calls).toMatchObject([
+        ['mainnet', {}, 1],
+        ['mainnet', {}, 2],
+      ]);
+
+      await advanceTime({ clock, duration: 5 });
+
+      expect(controller._executePoll.mock.calls).toMatchObject([
+        ['mainnet', {}, 1],
+        ['mainnet', {}, 2],
+        ['mainnet', {}, 3],
+        ['goerli', {}, 1],
+      ]);
+
+      // Start polling for sepolia, 15ms interval
+      controller.startPollingByNetworkClientId('sepolia');
+
+      await advanceTime({ clock, duration: 15 });
+
+      expect(controller._executePoll.mock.calls).toMatchObject([
+        ['mainnet', {}, 1],
+        ['mainnet', {}, 2],
+        ['mainnet', {}, 3],
+        ['goerli', {}, 1],
+        ['mainnet', {}, 4],
+        ['mainnet', {}, 5],
+        ['goerli', {}, 2],
+        ['mainnet', {}, 6],
+        ['sepolia', {}, 1],
+      ]);
+
+      // Stop all polling
+      controller.stopAllPolling();
+    });
+  });
 });