refactor: migrate AccountTreeController to @metamask/messenger

README.md

@@ -148,6 +148,7 @@ linkStyle default opacity:0.5
   transaction_controller(["@metamask/transaction-controller"]);
   user_operation_controller(["@metamask/user-operation-controller"]);
   account_tree_controller --> base_controller;
+  account_tree_controller --> messenger;
   account_tree_controller --> accounts_controller;
   account_tree_controller --> keyring_controller;
   account_tree_controller --> multichain_account_service;
@@ -216,6 +217,7 @@ linkStyle default opacity:0.5
   earn_controller --> transaction_controller;
   eip_5792_middleware --> transaction_controller;
   eip_5792_middleware --> keyring_controller;
+  eip_7702_internal_rpc_middleware --> controller_utils;
   eip1193_permission_middleware --> chain_agnostic_permission;
   eip1193_permission_middleware --> controller_utils;
   eip1193_permission_middleware --> json_rpc_engine;

packages/account-tree-controller/CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+
+- **BREAKING:** Use new `Messenger` from `@metamask/messenger` ([#6380](https://github.com/MetaMask/core/pull/6380))
+  - Previously, `AccountTreeController` accepted a `RestrictedMessenger` instance from `@metamask/base-controller`.
+
 ## [1.6.0]
 
 ### Changed

packages/account-tree-controller/package.json

@@ -48,6 +48,7 @@
   },
   "dependencies": {
     "@metamask/base-controller": "^8.4.2",
+    "@metamask/messenger": "^0.3.0",
     "@metamask/snaps-sdk": "^9.0.0",
     "@metamask/snaps-utils": "^11.0.0",
     "@metamask/superstruct": "^3.1.0",

packages/account-tree-controller/src/AccountTreeController.test.ts

@@ -9,7 +9,7 @@ import {
   type AccountGroupId,
 } from '@metamask/account-api';
 import type { AccountId } from '@metamask/accounts-controller';
-import { Messenger, deriveStateFromMetadata } from '@metamask/base-controller';
+import { deriveStateFromMetadata } from '@metamask/base-controller/next';
 import {
   BtcAccountType,
   EthAccountType,
@@ -36,14 +36,11 @@ import type { BackupAndSyncAnalyticsEventPayload } from './backup-and-sync/analy
 import { BackupAndSyncService } from './backup-and-sync/service';
 import { isAccountGroupNameUnique } from './group';
 import { getAccountWalletNameFromKeyringType } from './rules/keyring';
+import { type AccountTreeControllerState } from './types';
 import {
-  type AccountTreeControllerMessenger,
-  type AccountTreeControllerActions,
-  type AccountTreeControllerEvents,
-  type AccountTreeControllerState,
-  type AllowedActions,
-  type AllowedEvents,
-} from './types';
+  getAccountTreeControllerMessenger,
+  getRootMessenger,
+} from '../tests/mockMessenger';
 
 // Local mock of EMPTY_ACCOUNT to avoid circular dependency
 const EMPTY_ACCOUNT_MOCK: InternalAccount = {
@@ -233,53 +230,7 @@ const MOCK_HARDWARE_ACCOUNT_1: InternalAccount = {
   },
 };
 
-/**
- * Creates a new root messenger instance for testing.
- *
- * @returns A new Messenger instance.
- */
-function getRootMessenger() {
-  return new Messenger<
-    AccountTreeControllerActions | AllowedActions,
-    AccountTreeControllerEvents | AllowedEvents
-  >();
-}
-
-/**
- * Retrieves a restricted messenger for the AccountTreeController.
- *
- * @param messenger - The root messenger instance. Defaults to a new Messenger created by getRootMessenger().
- * @returns The restricted messenger for the AccountTreeController.
- */
-function getAccountTreeControllerMessenger(
-  messenger = getRootMessenger(),
-): AccountTreeControllerMessenger {
-  return messenger.getRestricted({
-    name: 'AccountTreeController',
-    allowedEvents: [
-      'AccountsController:accountAdded',
-      'AccountsController:accountRemoved',
-      'AccountsController:selectedAccountChange',
-      'UserStorageController:stateChange',
-      'MultichainAccountService:walletStatusChange',
-    ],
-    allowedActions: [
-      'AccountsController:listMultichainAccounts',
-      'AccountsController:getAccount',
-      'AccountsController:getSelectedMultichainAccount',
-      'AccountsController:setSelectedAccount',
-      'UserStorageController:getState',
-      'UserStorageController:performGetStorage',
-      'UserStorageController:performGetStorageAllFeatureEntries',
-      'UserStorageController:performSetStorage',
-      'UserStorageController:performBatchSetStorage',
-      'AuthenticationController:getSessionProfile',
-      'MultichainAccountService:createMultichainAccountGroup',
-      'KeyringController:getState',
-      'SnapController:get',
-    ],
-  });
-}
+const mockGetSelectedMultichainAccountActionHandler = jest.fn();
 
 /**
  * Sets up the AccountTreeController for testing.
@@ -317,10 +268,7 @@ function setup({
   },
 }: {
   state?: Partial<AccountTreeControllerState>;
-  messenger?: Messenger<
-    AccountTreeControllerActions | AllowedActions,
-    AccountTreeControllerEvents | AllowedEvents
-  >;
+  messenger?: ReturnType<typeof getRootMessenger>;
   accounts?: InternalAccount[];
   keyrings?: KeyringObject[];
   config?: {
@@ -338,9 +286,9 @@ function setup({
   };
 } = {}): {
   controller: AccountTreeController;
-  messenger: Messenger<
-    AccountTreeControllerActions | AllowedActions,
-    AccountTreeControllerEvents | AllowedEvents
+  messenger: ReturnType<typeof getRootMessenger>;
+  accountTreeControllerMessenger: ReturnType<
+    typeof getAccountTreeControllerMessenger
   >;
   spies: {
     consoleWarn: jest.SpyInstance;
@@ -401,6 +349,7 @@ function setup({
     mocks.AccountsController.listMultichainAccounts.mockImplementation(
       () => mocks.AccountsController.accounts,
     );
+
     messenger.registerActionHandler(
       'AccountsController:listMultichainAccounts',
       mocks.AccountsController.listMultichainAccounts,
@@ -474,8 +423,10 @@ function setup({
     );
   }
 
+  const accountTreeControllerMessenger =
+    getAccountTreeControllerMessenger(messenger);
   const controller = new AccountTreeController({
-    messenger: getAccountTreeControllerMessenger(messenger),
+    messenger: accountTreeControllerMessenger,
     state,
     ...(config && { config }),
   });
@@ -487,6 +438,7 @@ function setup({
   return {
     controller,
     messenger,
+    accountTreeControllerMessenger,
     spies: { consoleWarn: consoleWarnSpy },
     mocks,
   };
@@ -1802,12 +1754,15 @@ describe('AccountTreeController', () => {
     });
 
     it('updates AccountsController selected account (with EVM account) when selectedAccountGroup changes', () => {
-      const { controller, messenger } = setup({
+      const { controller, accountTreeControllerMessenger } = setup({
         accounts: [MOCK_HD_ACCOUNT_1, MOCK_HD_ACCOUNT_2],
         keyrings: [MOCK_HD_KEYRING_1, MOCK_HD_KEYRING_2],
       });
 
-      const setSelectedAccountSpy = jest.spyOn(messenger, 'call');
+      const setSelectedAccountSpy = jest.spyOn(
+        accountTreeControllerMessenger,
+        'call',
+      );
 
       controller.init();
 
@@ -1839,15 +1794,18 @@ describe('AccountTreeController', () => {
           },
         },
       } as const;
-      const { controller, messenger } = setup({
+      const { controller, accountTreeControllerMessenger } = setup({
         accounts: [
           MOCK_HD_ACCOUNT_1,
           nonEvmAccount2, // Wallet 2 > Account 1.
         ],
         keyrings: [MOCK_HD_KEYRING_1, MOCK_HD_KEYRING_2],
       });
 
-      const setSelectedAccountSpy = jest.spyOn(messenger, 'call');
+      const setSelectedAccountSpy = jest.spyOn(
+        accountTreeControllerMessenger,
+        'call',
+      );
 
       controller.init();
 
@@ -1969,18 +1927,14 @@ describe('AccountTreeController', () => {
     });
 
     it('falls back to first wallet first group when AccountsController returns EMPTY_ACCOUNT', () => {
-      const { controller, messenger } = setup({
+      const { controller } = setup({
         accounts: [MOCK_HD_ACCOUNT_1, MOCK_HD_ACCOUNT_2],
         keyrings: [MOCK_HD_KEYRING_1, MOCK_HD_KEYRING_2],
       });
 
-      // Unregister existing handler and register new one BEFORE init
-      messenger.unregisterActionHandler(
-        'AccountsController:getSelectedMultichainAccount',
-      );
-      messenger.registerActionHandler(
-        'AccountsController:getSelectedMultichainAccount',
-        () => EMPTY_ACCOUNT_MOCK,
+      // Mock action handler BEFORE init
+      mockGetSelectedMultichainAccountActionHandler.mockReturnValue(
+        EMPTY_ACCOUNT_MOCK,
       );
 
       controller.init();
@@ -1998,7 +1952,7 @@ describe('AccountTreeController', () => {
     });
 
     it('falls back to first wallet first group when selected account is not in tree', () => {
-      const { controller, messenger } = setup({
+      const { controller } = setup({
         accounts: [MOCK_HD_ACCOUNT_1, MOCK_HD_ACCOUNT_2],
         keyrings: [MOCK_HD_KEYRING_1, MOCK_HD_KEYRING_2],
       });
@@ -2009,12 +1963,8 @@ describe('AccountTreeController', () => {
         id: 'unknown-account-id',
       };
 
-      messenger.unregisterActionHandler(
-        'AccountsController:getSelectedMultichainAccount',
-      );
-      messenger.registerActionHandler(
-        'AccountsController:getSelectedMultichainAccount',
-        () => unknownAccount,
+      mockGetSelectedMultichainAccountActionHandler.mockReturnValue(
+        unknownAccount,
       );
 
       controller.init();
@@ -2032,18 +1982,14 @@ describe('AccountTreeController', () => {
     });
 
     it('returns empty string when no wallets exist and getSelectedMultichainAccount returns EMPTY_ACCOUNT', () => {
-      const { controller, messenger } = setup({
+      const { controller } = setup({
         accounts: [],
         keyrings: [],
       });
 
-      // Mock getSelectedMultichainAccount to return EMPTY_ACCOUNT_MOCK (id is '') BEFORE init
-      messenger.unregisterActionHandler(
-        'AccountsController:getSelectedMultichainAccount',
-      );
-      messenger.registerActionHandler(
-        'AccountsController:getSelectedMultichainAccount',
-        () => EMPTY_ACCOUNT_MOCK,
+      // Mock getSelectedAccount to return EMPTY_ACCOUNT_MOCK (id is '') BEFORE init
+      mockGetSelectedMultichainAccountActionHandler.mockReturnValue(
+        EMPTY_ACCOUNT_MOCK,
       );
 
       controller.init();
@@ -4328,7 +4274,7 @@ describe('AccountTreeController', () => {
         deriveStateFromMetadata(
           controller.state,
           controller.metadata,
-          'anonymous',
+          'includeInDebugSnapshot',
         ),
       ).toMatchInlineSnapshot(`Object {}`);
     });