99 type AccountGroupId ,
1010} from '@metamask/account-api' ;
1111import type { AccountId } from '@metamask/accounts-controller' ;
12- import { Messenger , deriveStateFromMetadata } from '@metamask/base-controller' ;
12+ import { deriveStateFromMetadata } from '@metamask/base-controller/next ' ;
1313import {
1414 BtcAccountType ,
1515 EthAccountType ,
@@ -36,14 +36,11 @@ import type { BackupAndSyncAnalyticsEventPayload } from './backup-and-sync/analy
3636import { BackupAndSyncService } from './backup-and-sync/service' ;
3737import { isAccountGroupNameUnique } from './group' ;
3838import { getAccountWalletNameFromKeyringType } from './rules/keyring' ;
39+ import { type AccountTreeControllerState } from './types' ;
3940import {
40- type AccountTreeControllerMessenger ,
41- type AccountTreeControllerActions ,
42- type AccountTreeControllerEvents ,
43- type AccountTreeControllerState ,
44- type AllowedActions ,
45- type AllowedEvents ,
46- } from './types' ;
41+ getAccountTreeControllerMessenger ,
42+ getRootMessenger ,
43+ } from '../tests/mockMessenger' ;
4744
4845// Local mock of EMPTY_ACCOUNT to avoid circular dependency
4946const EMPTY_ACCOUNT_MOCK : InternalAccount = {
@@ -233,53 +230,7 @@ const MOCK_HARDWARE_ACCOUNT_1: InternalAccount = {
233230 } ,
234231} ;
235232
236- /**
237- * Creates a new root messenger instance for testing.
238- *
239- * @returns A new Messenger instance.
240- */
241- function getRootMessenger ( ) {
242- return new Messenger <
243- AccountTreeControllerActions | AllowedActions ,
244- AccountTreeControllerEvents | AllowedEvents
245- > ( ) ;
246- }
247-
248- /**
249- * Retrieves a restricted messenger for the AccountTreeController.
250- *
251- * @param messenger - The root messenger instance. Defaults to a new Messenger created by getRootMessenger().
252- * @returns The restricted messenger for the AccountTreeController.
253- */
254- function getAccountTreeControllerMessenger (
255- messenger = getRootMessenger ( ) ,
256- ) : AccountTreeControllerMessenger {
257- return messenger . getRestricted ( {
258- name : 'AccountTreeController' ,
259- allowedEvents : [
260- 'AccountsController:accountAdded' ,
261- 'AccountsController:accountRemoved' ,
262- 'AccountsController:selectedAccountChange' ,
263- 'UserStorageController:stateChange' ,
264- 'MultichainAccountService:walletStatusChange' ,
265- ] ,
266- allowedActions : [
267- 'AccountsController:listMultichainAccounts' ,
268- 'AccountsController:getAccount' ,
269- 'AccountsController:getSelectedMultichainAccount' ,
270- 'AccountsController:setSelectedAccount' ,
271- 'UserStorageController:getState' ,
272- 'UserStorageController:performGetStorage' ,
273- 'UserStorageController:performGetStorageAllFeatureEntries' ,
274- 'UserStorageController:performSetStorage' ,
275- 'UserStorageController:performBatchSetStorage' ,
276- 'AuthenticationController:getSessionProfile' ,
277- 'MultichainAccountService:createMultichainAccountGroup' ,
278- 'KeyringController:getState' ,
279- 'SnapController:get' ,
280- ] ,
281- } ) ;
282- }
233+ const mockGetSelectedMultichainAccountActionHandler = jest . fn ( ) ;
283234
284235/**
285236 * Sets up the AccountTreeController for testing.
@@ -317,10 +268,7 @@ function setup({
317268 } ,
318269} : {
319270 state ?: Partial < AccountTreeControllerState > ;
320- messenger ?: Messenger <
321- AccountTreeControllerActions | AllowedActions ,
322- AccountTreeControllerEvents | AllowedEvents
323- > ;
271+ messenger ?: ReturnType < typeof getRootMessenger > ;
324272 accounts ?: InternalAccount [ ] ;
325273 keyrings ?: KeyringObject [ ] ;
326274 config ?: {
@@ -338,9 +286,9 @@ function setup({
338286 } ;
339287} = { } ) : {
340288 controller : AccountTreeController ;
341- messenger : Messenger <
342- AccountTreeControllerActions | AllowedActions ,
343- AccountTreeControllerEvents | AllowedEvents
289+ messenger : ReturnType < typeof getRootMessenger > ;
290+ accountTreeControllerMessenger : ReturnType <
291+ typeof getAccountTreeControllerMessenger
344292 > ;
345293 spies : {
346294 consoleWarn : jest . SpyInstance ;
@@ -401,6 +349,7 @@ function setup({
401349 mocks . AccountsController . listMultichainAccounts . mockImplementation (
402350 ( ) => mocks . AccountsController . accounts ,
403351 ) ;
352+
404353 messenger . registerActionHandler (
405354 'AccountsController:listMultichainAccounts' ,
406355 mocks . AccountsController . listMultichainAccounts ,
@@ -474,8 +423,10 @@ function setup({
474423 ) ;
475424 }
476425
426+ const accountTreeControllerMessenger =
427+ getAccountTreeControllerMessenger ( messenger ) ;
477428 const controller = new AccountTreeController ( {
478- messenger : getAccountTreeControllerMessenger ( messenger ) ,
429+ messenger : accountTreeControllerMessenger ,
479430 state,
480431 ...( config && { config } ) ,
481432 } ) ;
@@ -487,6 +438,7 @@ function setup({
487438 return {
488439 controller,
489440 messenger,
441+ accountTreeControllerMessenger,
490442 spies : { consoleWarn : consoleWarnSpy } ,
491443 mocks,
492444 } ;
@@ -1802,12 +1754,15 @@ describe('AccountTreeController', () => {
18021754 } ) ;
18031755
18041756 it ( 'updates AccountsController selected account (with EVM account) when selectedAccountGroup changes' , ( ) => {
1805- const { controller, messenger } = setup ( {
1757+ const { controller, accountTreeControllerMessenger } = setup ( {
18061758 accounts : [ MOCK_HD_ACCOUNT_1 , MOCK_HD_ACCOUNT_2 ] ,
18071759 keyrings : [ MOCK_HD_KEYRING_1 , MOCK_HD_KEYRING_2 ] ,
18081760 } ) ;
18091761
1810- const setSelectedAccountSpy = jest . spyOn ( messenger , 'call' ) ;
1762+ const setSelectedAccountSpy = jest . spyOn (
1763+ accountTreeControllerMessenger ,
1764+ 'call' ,
1765+ ) ;
18111766
18121767 controller . init ( ) ;
18131768
@@ -1839,15 +1794,18 @@ describe('AccountTreeController', () => {
18391794 } ,
18401795 } ,
18411796 } as const ;
1842- const { controller, messenger } = setup ( {
1797+ const { controller, accountTreeControllerMessenger } = setup ( {
18431798 accounts : [
18441799 MOCK_HD_ACCOUNT_1 ,
18451800 nonEvmAccount2 , // Wallet 2 > Account 1.
18461801 ] ,
18471802 keyrings : [ MOCK_HD_KEYRING_1 , MOCK_HD_KEYRING_2 ] ,
18481803 } ) ;
18491804
1850- const setSelectedAccountSpy = jest . spyOn ( messenger , 'call' ) ;
1805+ const setSelectedAccountSpy = jest . spyOn (
1806+ accountTreeControllerMessenger ,
1807+ 'call' ,
1808+ ) ;
18511809
18521810 controller . init ( ) ;
18531811
@@ -1969,18 +1927,14 @@ describe('AccountTreeController', () => {
19691927 } ) ;
19701928
19711929 it ( 'falls back to first wallet first group when AccountsController returns EMPTY_ACCOUNT' , ( ) => {
1972- const { controller, messenger } = setup ( {
1930+ const { controller } = setup ( {
19731931 accounts : [ MOCK_HD_ACCOUNT_1 , MOCK_HD_ACCOUNT_2 ] ,
19741932 keyrings : [ MOCK_HD_KEYRING_1 , MOCK_HD_KEYRING_2 ] ,
19751933 } ) ;
19761934
1977- // Unregister existing handler and register new one BEFORE init
1978- messenger . unregisterActionHandler (
1979- 'AccountsController:getSelectedMultichainAccount' ,
1980- ) ;
1981- messenger . registerActionHandler (
1982- 'AccountsController:getSelectedMultichainAccount' ,
1983- ( ) => EMPTY_ACCOUNT_MOCK ,
1935+ // Mock action handler BEFORE init
1936+ mockGetSelectedMultichainAccountActionHandler . mockReturnValue (
1937+ EMPTY_ACCOUNT_MOCK ,
19841938 ) ;
19851939
19861940 controller . init ( ) ;
@@ -1998,7 +1952,7 @@ describe('AccountTreeController', () => {
19981952 } ) ;
19991953
20001954 it ( 'falls back to first wallet first group when selected account is not in tree' , ( ) => {
2001- const { controller, messenger } = setup ( {
1955+ const { controller } = setup ( {
20021956 accounts : [ MOCK_HD_ACCOUNT_1 , MOCK_HD_ACCOUNT_2 ] ,
20031957 keyrings : [ MOCK_HD_KEYRING_1 , MOCK_HD_KEYRING_2 ] ,
20041958 } ) ;
@@ -2009,12 +1963,8 @@ describe('AccountTreeController', () => {
20091963 id : 'unknown-account-id' ,
20101964 } ;
20111965
2012- messenger . unregisterActionHandler (
2013- 'AccountsController:getSelectedMultichainAccount' ,
2014- ) ;
2015- messenger . registerActionHandler (
2016- 'AccountsController:getSelectedMultichainAccount' ,
2017- ( ) => unknownAccount ,
1966+ mockGetSelectedMultichainAccountActionHandler . mockReturnValue (
1967+ unknownAccount ,
20181968 ) ;
20191969
20201970 controller . init ( ) ;
@@ -2032,18 +1982,14 @@ describe('AccountTreeController', () => {
20321982 } ) ;
20331983
20341984 it ( 'returns empty string when no wallets exist and getSelectedMultichainAccount returns EMPTY_ACCOUNT' , ( ) => {
2035- const { controller, messenger } = setup ( {
1985+ const { controller } = setup ( {
20361986 accounts : [ ] ,
20371987 keyrings : [ ] ,
20381988 } ) ;
20391989
2040- // Mock getSelectedMultichainAccount to return EMPTY_ACCOUNT_MOCK (id is '') BEFORE init
2041- messenger . unregisterActionHandler (
2042- 'AccountsController:getSelectedMultichainAccount' ,
2043- ) ;
2044- messenger . registerActionHandler (
2045- 'AccountsController:getSelectedMultichainAccount' ,
2046- ( ) => EMPTY_ACCOUNT_MOCK ,
1990+ // Mock getSelectedAccount to return EMPTY_ACCOUNT_MOCK (id is '') BEFORE init
1991+ mockGetSelectedMultichainAccountActionHandler . mockReturnValue (
1992+ EMPTY_ACCOUNT_MOCK ,
20471993 ) ;
20481994
20491995 controller . init ( ) ;
@@ -4328,7 +4274,7 @@ describe('AccountTreeController', () => {
43284274 deriveStateFromMetadata (
43294275 controller . state ,
43304276 controller . metadata ,
4331- 'anonymous ' ,
4277+ 'includeInDebugSnapshot ' ,
43324278 ) ,
43334279 ) . toMatchInlineSnapshot ( `Object {}` ) ;
43344280 } ) ;
0 commit comments