WIP Multichainify TokenListController

MetaMask · Oct 3, 2023 · cd0de7a · cd0de7a
1 parent 014f632
commit cd0de7a
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 11 deletions.
diff --git a/packages/assets-controllers/src/TokenListController.test.ts b/packages/assets-controllers/src/TokenListController.test.ts
@@ -7,6 +7,7 @@ import {
   toHex,
 } from '@metamask/controller-utils';
 import type {
+  NetworkControllerGetNetworkClientByIdAction,
   NetworkControllerStateChangeEvent,
   NetworkState,
   ProviderConfig,
@@ -15,7 +16,10 @@ import { NetworkStatus } from '@metamask/network-controller';
 import nock from 'nock';
 import * as sinon from 'sinon';
 
-import { TOKEN_END_POINT_API } from './token-service';
+import {
+  TOKEN_END_POINT_API,
+  fetchTokenList as tokenServiceFetchTokenList,
+} from './token-service';
 import type {
   TokenListStateChange,
   GetTokenListState,
@@ -481,7 +485,7 @@ const expiredCacheExistingState: TokenListState = {
 };
 
 type MainControllerMessenger = ControllerMessenger<
-  GetTokenListState,
+  GetTokenListState | NetworkControllerGetNetworkClientByIdAction,
   TokenListStateChange | NetworkControllerStateChangeEvent
 >;
 
@@ -494,7 +498,7 @@ const getRestrictedMessenger = (
 ) => {
   const messenger = controllerMessenger.getRestricted({
     name,
-    allowedActions: [],
+    allowedActions: ['NetworkController:getNetworkClientById'],
     allowedEvents: [
       'TokenListController:stateChange',
       'NetworkController:stateChange',
@@ -1192,4 +1196,23 @@ describe('TokenListController', () => {
       );
     });
   });
+
+  describe('executePoll', () => {
+    it('should execute the poll', async () => {
+      const tokenServiceFetchTokenListSpy = sinon.spy(
+        tokenServiceFetchTokenList,
+      );
+
+      const controllerMessenger = getControllerMessenger();
+      const messenger = getRestrictedMessenger(controllerMessenger);
+      const controller = new TokenListController({
+        chainId: ChainId.mainnet,
+        preventPollingOnNetworkRestart: false,
+        messenger,
+        state: existingState,
+      });
+      await controller.executePoll('sepolia');
+      expect(tokenServiceFetchTokenListSpy.calledWith('0x1')).toBe(true);
+    });
+  });
 });
diff --git a/packages/assets-controllers/src/TokenListController.ts b/packages/assets-controllers/src/TokenListController.ts
@@ -1,10 +1,12 @@
 import type { RestrictedControllerMessenger } from '@metamask/base-controller';
-import { BaseControllerV2 } from '@metamask/base-controller';
 import { safelyExecute } from '@metamask/controller-utils';
 import type {
+  NetworkClientId,
   NetworkControllerStateChangeEvent,
   NetworkState,
+  NetworkControllerGetNetworkClientByIdAction,
 } from '@metamask/network-controller';
+import { PollingController } from '@metamask/polling-controller';
 import type { Hex } from '@metamask/utils';
 import { Mutex } from 'async-mutex';
 import type { Patch } from 'immer';
@@ -56,12 +58,11 @@ export type GetTokenListState = {
   type: `${typeof name}:getState`;
   handler: () => TokenListState;
 };
-
 type TokenListMessenger = RestrictedControllerMessenger<
   typeof name,
-  GetTokenListState,
+  GetTokenListState | NetworkControllerGetNetworkClientByIdAction,
   TokenListStateChange | NetworkControllerStateChangeEvent,
-  never,
+  NetworkControllerGetNetworkClientByIdAction['type'],
   TokenListStateChange['type'] | NetworkControllerStateChangeEvent['type']
 >;
 
@@ -80,7 +81,7 @@ const defaultState: TokenListState = {
 /**
  * Controller that passively polls on a set interval for the list of tokens from metaswaps api
  */
-export class TokenListController extends BaseControllerV2<
+export class TokenListController extends PollingController<
   typeof name,
   TokenListState,
   TokenListMessenger
@@ -232,8 +233,29 @@ export class TokenListController extends BaseControllerV2<
 
   /**
    * Fetching token list from the Token Service API.
+   *
+   * @param networkClientId - The ID of the network client triggering the fetch.
+   * @returns A promise that resolves when this operation completes.
    */
-  async fetchTokenList(): Promise<void> {
+  async executePoll(networkClientId: string): Promise<void> {
+    return this.fetchTokenList(networkClientId);
+  }
+
+  /**
+   * Fetching token list from the Token Service API.
+   *
+   * @param networkClientId - The ID of the network client triggering the fetch.
+   */
+  async fetchTokenList(networkClientId?: NetworkClientId): Promise<void> {
+    let networkClient;
+    if (networkClientId) {
+      networkClient = this.messagingSystem.call(
+        'NetworkController:getNetworkClientById',
+        networkClientId,
+      );
+    }
+
+    const chainId = networkClient?.configuration.chainId ?? this.chainId;
     const releaseLock = await this.mutex.acquire();
     try {
       const { tokensChainsCache } = this.state;
@@ -247,12 +269,12 @@ export class TokenListController extends BaseControllerV2<
       } else {
         // Fetch fresh token list
         const tokensFromAPI: TokenListToken[] = await safelyExecute(() =>
-          fetchTokenList(this.chainId, this.abortController.signal),
+          fetchTokenList(chainId, this.abortController.signal),
         );
 
         if (!tokensFromAPI) {
           // Fallback to expired cached tokens
-          tokenList = { ...(tokensChainsCache[this.chainId]?.data || {}) };
+          tokenList = { ...(tokensChainsCache[chainId]?.data || {}) };
 
           this.update(() => {
             return {