refactor: migrate AccountTreeController to @metamask/messenger (#6380)

## Explanation

<!--
Thanks for your contribution! Take a moment to answer these questions so
that reviewers have the information they need to properly understand
your changes:

* What is the current state of things and why does it need to change?
* What is the solution your changes offer and how does it work?
* Are there any changes whose purpose might not obvious to those
unfamiliar with the domain?
* If your primary goal was to update one package but you found you had
to update another one along the way, why did you do so?
* If you had to upgrade a dependency, why did you do so?
-->
This PR migrates the `AccountTreeController` class to the new
`@metamask/messenger` comm system, as opposed to the one exported from
`@metamask/base-controller`.

## References

<!--
Are there any issues that this pull request is tied to?
Are there other links that reviewers should consult to understand these
changes better?
Are there client or consumer pull requests to adopt any breaking
changes?

For example:

* Fixes #12345
* Related to #67890
-->
* Related to #5626

## Checklist

- [ ] I've updated the test suite for new or updated code as appropriate
- [ ] I've updated documentation (JSDoc, Markdown, etc.) for new or
updated code as appropriate
- [ ] I've communicated my changes to consumers by [updating changelogs
for packages I've
changed](https://github.com/MetaMask/core/tree/main/docs/contributing.md#updating-changelogs),
highlighting breaking changes as necessary
- [ ] I've prepared draft pull requests for clients and consumer
packages to resolve any breaking changes

Author: mikesposito

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 {}`);
     });