refactor: migrate tokenRates controller to modular controller init architecture

salimtb · salimtb · commit d975a69c18ac · 2025-03-24T22:18:59.000+01:00
diff --git a/app/scripts/controller-init/assets/index.ts b/app/scripts/controller-init/assets/index.ts
@@ -0,0 +1 @@
+export { TokenRatesControllerInit } from './token-rates-controller-init';
diff --git a/app/scripts/controller-init/assets/token-rates-controller-init.test.ts b/app/scripts/controller-init/assets/token-rates-controller-init.test.ts
@@ -0,0 +1,47 @@
+import {
+  TokenRatesController,
+  TokenRatesControllerMessenger,
+} from '@metamask/assets-controllers';
+import { Messenger } from '@metamask/base-controller';
+import { buildControllerInitRequestMock } from '../test/utils';
+import { ControllerInitRequest } from '../types';
+import { getTokenRatesControllerMessenger } from '../messengers/assets';
+import { TokenRatesControllerInit } from './token-rates-controller-init';
+
+jest.mock('@metamask/assets-controllers');
+
+function buildInitRequestMock(): jest.Mocked<
+  ControllerInitRequest<TokenRatesControllerMessenger>
+> {
+  const baseControllerMessenger = new Messenger();
+
+  return {
+    ...buildControllerInitRequestMock(),
+    controllerMessenger: getTokenRatesControllerMessenger(
+      baseControllerMessenger,
+    ),
+    initMessenger: undefined,
+  };
+}
+
+describe('TokenRatesControllerInit', () => {
+  const tokenRatesControllerClassMock = jest.mocked(TokenRatesController);
+
+  beforeEach(() => {
+    jest.resetAllMocks();
+  });
+
+  it('returns controller instance', () => {
+    const requestMock = buildInitRequestMock();
+    expect(TokenRatesControllerInit(requestMock).controller).toBeInstanceOf(
+      TokenRatesController,
+    );
+  });
+
+  it('initializes with correct messenger and state', () => {
+    const requestMock = buildInitRequestMock();
+    TokenRatesControllerInit(requestMock);
+
+    expect(tokenRatesControllerClassMock).toHaveBeenCalled();
+  });
+});
diff --git a/app/scripts/controller-init/assets/token-rates-controller-init.ts b/app/scripts/controller-init/assets/token-rates-controller-init.ts
@@ -0,0 +1,30 @@
+import {
+  CodefiTokenPricesServiceV2,
+  TokenRatesController,
+} from '@metamask/assets-controllers';
+import { ControllerInitFunction } from '../types';
+import { TokenRatesControllerMessenger } from '../messengers/assets';
+
+/**
+ * Initialize the Token Rates controller.
+ *
+ * @param request - The request object.
+ * @param request.controllerMessenger - The messenger to use for the controller.
+ * @param request.persistedState - The persisted state of the extension.
+ * @returns The initialized controller.
+ */
+export const TokenRatesControllerInit: ControllerInitFunction<
+  TokenRatesController,
+  TokenRatesControllerMessenger
+> = ({ controllerMessenger, persistedState }) => {
+  const controller = new TokenRatesController({
+    messenger: controllerMessenger,
+    state: persistedState.TokenRatesController,
+    tokenPricesService: new CodefiTokenPricesServiceV2(),
+    disabled: !persistedState.PreferencesController?.useCurrencyRateCheck,
+  });
+
+  return {
+    controller,
+  };
+};
diff --git a/app/scripts/controller-init/controller-list.ts b/app/scripts/controller-init/controller-list.ts
@@ -15,6 +15,7 @@ import {
   MultichainAssetsController,
   MultichainAssetsRatesController,
   MultichainBalancesController,
+  TokenRatesController,
 } from '@metamask/assets-controllers';
 import { MultichainNetworkController } from '@metamask/multichain-network-controller';
 import { MultichainTransactionsController } from '@metamask/multichain-transactions-controller';
@@ -69,7 +70,8 @@ export type Controller =
       name: 'TransactionUpdateController';
       state: Record<string, unknown>;
     })
-  | UserStorageController;
+  | UserStorageController
+  | TokenRatesController;
 
 /**
  * Flat state object for all controllers supporting or required by modular initialization.
@@ -100,4 +102,5 @@ export type ControllerFlatState = AccountsController['state'] &
   SnapInterfaceController['state'] &
   TransactionController['state'] &
   SwapsController['state'] &
-  UserStorageController['state'];
+  UserStorageController['state'] &
+  TokenRatesController['state'];
diff --git a/app/scripts/controller-init/messengers/assets/index.ts b/app/scripts/controller-init/messengers/assets/index.ts
@@ -0,0 +1,3 @@
+export { getTokenRatesControllerMessenger } from './token-rates-controller-messenger';
+
+export type { TokenRatesControllerMessenger } from './token-rates-controller-messenger';
diff --git a/app/scripts/controller-init/messengers/assets/token-rates-controller-messenger.test.ts b/app/scripts/controller-init/messengers/assets/token-rates-controller-messenger.test.ts
@@ -0,0 +1,12 @@
+import { Messenger, RestrictedMessenger } from '@metamask/base-controller';
+import { getTokenRatesControllerMessenger } from './token-rates-controller-messenger';
+
+describe('getTokenRatesControllerMessenger', () => {
+  it('returns a restricted messenger', () => {
+    const messenger = new Messenger<never, never>();
+    const tokenRatesControllerMessenger =
+      getTokenRatesControllerMessenger(messenger);
+
+    expect(tokenRatesControllerMessenger).toBeInstanceOf(RestrictedMessenger);
+  });
+});
diff --git a/app/scripts/controller-init/messengers/assets/token-rates-controller-messenger.ts b/app/scripts/controller-init/messengers/assets/token-rates-controller-messenger.ts
@@ -0,0 +1,61 @@
+import { Messenger } from '@metamask/base-controller';
+import {
+  NetworkControllerGetStateAction,
+  NetworkControllerGetNetworkClientByIdAction,
+  NetworkControllerStateChangeEvent,
+} from '@metamask/network-controller';
+import {
+  AccountsControllerGetSelectedAccountAction,
+  AccountsControllerGetAccountAction,
+  AccountsControllerSelectedEvmAccountChangeEvent,
+} from '@metamask/accounts-controller';
+import { PreferencesControllerStateChangeEvent } from '@metamask/preferences-controller';
+import {
+  TokensControllerGetStateAction,
+  TokensControllerStateChangeEvent,
+} from '@metamask/assets-controllers';
+
+type Actions =
+  | TokensControllerGetStateAction
+  | NetworkControllerGetNetworkClientByIdAction
+  | NetworkControllerGetStateAction
+  | AccountsControllerGetAccountAction
+  | AccountsControllerGetSelectedAccountAction;
+
+type Events =
+  | NetworkControllerStateChangeEvent
+  | AccountsControllerSelectedEvmAccountChangeEvent
+  | PreferencesControllerStateChangeEvent
+  | TokensControllerStateChangeEvent;
+
+export type TokenRatesControllerMessenger = ReturnType<
+  typeof getTokenRatesControllerMessenger
+>;
+
+/**
+ * Get a restricted messenger for the Token Rates controller. This is scoped to the
+ * actions and events that the Token Rates controller is allowed to handle.
+ *
+ * @param messenger - The controller messenger to restrict.
+ * @returns The restricted controller messenger.
+ */
+export function getTokenRatesControllerMessenger(
+  messenger: Messenger<Actions, Events>,
+) {
+  return messenger.getRestricted({
+    name: 'TokenRatesController',
+    allowedActions: [
+      'TokensController:getState',
+      'NetworkController:getNetworkClientById',
+      'NetworkController:getState',
+      'AccountsController:getAccount',
+      'AccountsController:getSelectedAccount',
+    ],
+    allowedEvents: [
+      'NetworkController:stateChange',
+      'AccountsController:selectedEvmAccountChange',
+      'PreferencesController:stateChange',
+      'TokensController:stateChange',
+    ],
+  });
+}
diff --git a/app/scripts/controller-init/messengers/index.ts b/app/scripts/controller-init/messengers/index.ts
@@ -29,6 +29,7 @@ import {
   getAuthenticationControllerMessenger,
   getUserStorageControllerMessenger,
 } from './identity';
+import { getTokenRatesControllerMessenger } from './assets';
 
 export const CONTROLLER_MESSENGERS = {
   AuthenticationController: {
@@ -94,4 +95,8 @@ export const CONTROLLER_MESSENGERS = {
     getMessenger: getUserStorageControllerMessenger,
     getInitMessenger: noop,
   },
+  TokenRatesController: {
+    getMessenger: getTokenRatesControllerMessenger,
+    getInitMessenger: noop,
+  },
 } as const;
diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js
@@ -7,9 +7,7 @@ import {
   NftDetectionController,
   TokenDetectionController,
   TokenListController,
-  TokenRatesController,
   TokensController,
-  CodefiTokenPricesServiceV2,
   RatesController,
   fetchMultiExchangeRate,
   TokenBalancesController,
@@ -381,6 +379,7 @@ import {
   ///: END:ONLY_INCLUDE_IF
   MultichainNetworkControllerInit,
 } from './controller-init/multichain';
+import { TokenRatesControllerInit } from './controller-init/assets';
 import { TransactionControllerInit } from './controller-init/confirmations/transaction-controller-init';
 import { PPOMControllerInit } from './controller-init/confirmations/ppom-controller-init';
 import { initControllers } from './controller-init/utils';
@@ -1041,31 +1040,6 @@ export default class MetamaskController extends EventEmitter {
       fetchMultiExchangeRate,
     });
 
-    const tokenRatesMessenger = this.controllerMessenger.getRestricted({
-      name: 'TokenRatesController',
-      allowedActions: [
-        'TokensController:getState',
-        'NetworkController:getNetworkClientById',
-        'NetworkController:getState',
-        'AccountsController:getAccount',
-        'AccountsController:getSelectedAccount',
-      ],
-      allowedEvents: [
-        'NetworkController:stateChange',
-        'AccountsController:selectedEvmAccountChange',
-        'PreferencesController:stateChange',
-        'TokensController:stateChange',
-      ],
-    });
-
-    // token exchange rate tracker
-    this.tokenRatesController = new TokenRatesController({
-      state: initState.TokenRatesController,
-      messenger: tokenRatesMessenger,
-      tokenPricesService: new CodefiTokenPricesServiceV2(),
-      disabled: !this.preferencesController.state.useCurrencyRateCheck,
-    });
-
     this.controllerMessenger.subscribe(
       'PreferencesController:stateChange',
       previousValueComparator((prevState, currState) => {
@@ -2015,6 +1989,7 @@ export default class MetamaskController extends EventEmitter {
       MultichainTransactionsController: MultichainTransactionsControllerInit,
       ///: END:ONLY_INCLUDE_IF
       MultichainNetworkController: MultichainNetworkControllerInit,
+      TokenRatesController: TokenRatesControllerInit,
       AuthenticationController: AuthenticationControllerInit,
       UserStorageController: UserStorageControllerInit,
     };
@@ -2054,6 +2029,7 @@ export default class MetamaskController extends EventEmitter {
     this.multichainAssetsRatesController =
       controllersByName.MultichainAssetsRatesController;
     ///: END:ONLY_INCLUDE_IF
+    this.tokenRatesController = controllersByName.TokenRatesController;
     this.multichainNetworkController =
       controllersByName.MultichainNetworkController;
     this.authenticationController = controllersByName.AuthenticationController;
@@ -2260,6 +2236,7 @@ export default class MetamaskController extends EventEmitter {
         MultichainTransactionsController: this.multichainTransactionsController,
         MultichainAssetsRatesController: this.multichainAssetsRatesController,
         ///: END:ONLY_INCLUDE_IF
+        TokenRatesController: this.tokenRatesController,
         MultichainNetworkController: this.multichainNetworkController,
         NetworkController: this.networkController,
         KeyringController: this.keyringController,

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+export { TokenRatesControllerInit } from './token-rates-controller-init';`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+export { getTokenRatesControllerMessenger } from './token-rates-controller-messenger';`
	`2`	`+`
	`3`	`+export type { TokenRatesControllerMessenger } from './token-rates-controller-messenger';`