chore: get assets list from MultichainAssetsController state

packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.test.ts

@@ -15,14 +15,12 @@ import { KeyringTypes } from '@metamask/keyring-controller';
 import type { InternalAccount } from '@metamask/keyring-internal-api';
 import { v4 as uuidv4 } from 'uuid';
 
-import {
-  MultichainBalancesController,
-  getDefaultMultichainBalancesControllerState,
-} from './MultichainBalancesController';
+import { MultichainBalancesController } from '.';
 import type {
   MultichainBalancesControllerMessenger,
   MultichainBalancesControllerState,
-} from './MultichainBalancesController';
+} from '.';
+import { getDefaultMultichainBalancesControllerState } from './MultichainBalancesController';
 import type {
   ExtractAvailableAction,
   ExtractAvailableEvent,
@@ -122,6 +120,31 @@ function getRootMessenger(): Messenger<RootAction, RootEvent> {
   return new Messenger<RootAction, RootEvent>();
 }
 
+/**
+ * Constructs the restricted messenger for the MultichainBalancesController.
+ *
+ * @param messenger - The root messenger.
+ * @returns The unrestricted messenger suited for MultichainBalancesController.
+ */
+function getRestrictedMessenger(
+  messenger: Messenger<RootAction, RootEvent>,
+): MultichainBalancesControllerMessenger {
+  return messenger.getRestricted({
+    name: 'MultichainBalancesController',
+    allowedActions: [
+      'SnapController:handleRequest',
+      'AccountsController:listMultichainAccounts',
+      'MultichainAssetsController:getState',
+    ],
+    allowedEvents: [
+      'AccountsController:accountAdded',
+      'AccountsController:accountRemoved',
+      'AccountsController:accountBalancesUpdated',
+      'MultichainAssetsController:stateChange',
+    ],
+  });
+}
+
 const setupController = ({
   state = getDefaultMultichainBalancesControllerState(),
   mocks,
@@ -133,20 +156,7 @@ const setupController = ({
   };
 } = {}) => {
   const messenger = getRootMessenger();
-
-  const multichainBalancesMessenger: MultichainBalancesControllerMessenger =
-    messenger.getRestricted({
-      name: 'MultichainBalancesController',
-      allowedActions: [
-        'SnapController:handleRequest',
-        'AccountsController:listMultichainAccounts',
-      ],
-      allowedEvents: [
-        'AccountsController:accountAdded',
-        'AccountsController:accountRemoved',
-        'AccountsController:accountBalancesUpdated',
-      ],
-    });
+  const multichainBalancesMessenger = getRestrictedMessenger(messenger);
 
   const mockSnapHandleRequest = jest.fn();
   messenger.registerActionHandler(
@@ -164,6 +174,16 @@ const setupController = ({
     ),
   );
 
+  const mockGetAssetsState = jest.fn().mockReturnValue({
+    accountsAssets: {
+      [mockBtcAccount.id]: [mockBtcNativeAsset],
+    },
+  });
+  messenger.registerActionHandler(
+    'MultichainAssetsController:getState',
+    mockGetAssetsState,
+  );
+
   const controller = new MultichainBalancesController({
     messenger: multichainBalancesMessenger,
     state,
@@ -174,6 +194,7 @@ const setupController = ({
     messenger,
     mockSnapHandleRequest,
     mockListMultichainAccounts,
+    mockGetAssetsState,
   };
 };
 
@@ -193,7 +214,22 @@ async function waitForAllPromises(): Promise<void> {
 
 describe('BalancesController', () => {
   it('initialize with default state', () => {
-    const { controller } = setupController({});
+    const messenger = getRootMessenger();
+    const multichainBalancesMessenger = getRestrictedMessenger(messenger);
+
+    messenger.registerActionHandler('SnapController:handleRequest', jest.fn());
+    messenger.registerActionHandler(
+      'AccountsController:listMultichainAccounts',
+      jest.fn().mockReturnValue([]),
+    );
+    messenger.registerActionHandler(
+      'MultichainAssetsController:getState',
+      jest.fn(),
+    );
+
+    const controller = new MultichainBalancesController({
+      messenger: multichainBalancesMessenger,
+    });
     expect(controller.state).toStrictEqual({ balances: {} });
   });
 
@@ -206,26 +242,6 @@ describe('BalancesController', () => {
     );
   });
 
-  it('updates balances when "AccountsController:accountAdded" is fired', async () => {
-    const { controller, messenger, mockListMultichainAccounts } =
-      setupController({
-        mocks: {
-          listMultichainAccounts: [],
-        },
-      });
-
-    mockListMultichainAccounts.mockReturnValue([mockBtcAccount]);
-    messenger.publish('AccountsController:accountAdded', mockBtcAccount);
-
-    await waitForAllPromises();
-
-    expect(controller.state).toStrictEqual({
-      balances: {
-        [mockBtcAccount.id]: mockBalanceResult,
-      },
-    });
-  });
-
   it('updates balances when "AccountsController:accountRemoved" is fired', async () => {
     const { controller, messenger } = setupController();
 
@@ -292,25 +308,6 @@ describe('BalancesController', () => {
     expect(controller.state.balances).toStrictEqual({});
   });
 
-  it('handles errors gracefully when constructing the controller', async () => {
-    // This method will be used in the constructor of that controller.
-    const updateBalanceSpy = jest.spyOn(
-      MultichainBalancesController.prototype,
-      'updateBalance',
-    );
-    updateBalanceSpy.mockRejectedValue(
-      new Error('Something unexpected happen'),
-    );
-
-    const { controller } = setupController({
-      mocks: {
-        listMultichainAccounts: [mockBtcAccount],
-      },
-    });
-
-    expect(controller.state.balances).toStrictEqual({});
-  });
-
   it('handles errors when trying to upgrade the balance of a non-existing account', async () => {
     const { controller } = setupController({
       mocks: {
@@ -392,4 +389,41 @@ describe('BalancesController', () => {
       mockBalanceResult,
     );
   });
+
+  it('handles an account with no assets in MultichainAssetsController state', async () => {
+    const { controller, mockGetAssetsState } = setupController({
+      mocks: {
+        handleRequestReturnValue: {},
+      },
+    });
+
+    mockGetAssetsState.mockReturnValue({
+      accountsAssets: {},
+    });
+
+    await controller.updateBalance(mockBtcAccount.id);
+
+    expect(controller.state.balances[mockBtcAccount.id]).toStrictEqual({});
+  });
+
+  it('updates balances when receiving "MultichainAssetsController:stateChange" event', async () => {
+    const { controller, messenger } = setupController();
+
+    messenger.publish(
+      'MultichainAssetsController:stateChange',
+      {
+        assetsMetadata: {},
+        accountsAssets: {
+          [mockBtcAccount.id]: [mockBtcNativeAsset],
+        },
+      },
+      [],
+    );
+
+    await waitForAllPromises();
+
+    expect(controller.state.balances[mockBtcAccount.id]).toStrictEqual(
+      mockBalanceResult,
+    );
+  });
 });

packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.ts

@@ -24,8 +24,11 @@ import { HandlerType } from '@metamask/snaps-utils';
 import type { Json, JsonRpcRequest } from '@metamask/utils';
 import type { Draft } from 'immer';
 
-import { NETWORK_ASSETS_MAP } from '.';
-import { getScopeForAccount } from './utils';
+import type {
+  MultichainAssetsControllerGetStateAction,
+  MultichainAssetsControllerState,
+  MultichainAssetsControllerStateChangeEvent,
+} from '../MultichainAssetsController';
 
 const controllerName = 'MultichainBalancesController';
 
@@ -90,15 +93,18 @@ export type MultichainBalancesControllerEvents =
  */
 type AllowedActions =
   | HandleSnapRequest
-  | AccountsControllerListMultichainAccountsAction;
+  | AccountsControllerListMultichainAccountsAction
+  | MultichainAssetsControllerGetStateAction;
 
 /**
  * Events that this controller is allowed to subscribe.
  */
 type AllowedEvents =
   | AccountsControllerAccountAddedEvent
   | AccountsControllerAccountRemovedEvent
-  | AccountsControllerAccountBalancesUpdatesEvent;
+  | AccountsControllerAccountBalancesUpdatesEvent
+  | MultichainAssetsControllerStateChangeEvent;
+
 /**
  * Messenger type for the MultichainBalancesController.
  */
@@ -152,18 +158,11 @@ export class MultichainBalancesController extends BaseController<
 
     // Fetch initial balances for all non-EVM accounts
     for (const account of this.#listAccounts()) {
-      this.updateBalance(account.id).catch((error) => {
-        console.error(
-          `Failed to fetch initial balance for account ${account.id}:`,
-          error,
-        );
-      });
+      // Fetching the balance is asynchronous and we cannot use `await` here.
+      // eslint-disable-next-line no-void
+      void this.updateBalance(account.id);
     }
 
-    this.messagingSystem.subscribe(
-      'AccountsController:accountAdded',
-      (account: InternalAccount) => this.#handleOnAccountAdded(account),
-    );
     this.messagingSystem.subscribe(
       'AccountsController:accountRemoved',
       (account: string) => this.#handleOnAccountRemoved(account),
@@ -173,40 +172,68 @@ export class MultichainBalancesController extends BaseController<
       (balanceUpdate: AccountBalancesUpdatedEventPayload) =>
         this.#handleOnAccountBalancesUpdated(balanceUpdate),
     );
+    // TODO: Maybe add a MultichainAssetsController:accountAssetListUpdated event instead of using the entire state.
+    // Since MultichainAssetsController already listens for the AccountsController:accountAdded, we can rely in it for that event
+    // and not listen for it also here, in this controller, since it would be redundant
+    this.messagingSystem.subscribe(
+      'MultichainAssetsController:stateChange',
+      async (assetsState: MultichainAssetsControllerState) => {
+        for (const accountId of Object.keys(assetsState.accountsAssets)) {
+          await this.#updateBalance(
+            accountId,
+            assetsState.accountsAssets[accountId],
+          );
+        }
+      },
+    );
   }
 
   /**
    * Updates the balances of one account. This method doesn't return
    * anything, but it updates the state of the controller.
    *
    * @param accountId - The account ID.
+   * @param assets - The list of asset types for this account to upadte.
    */
-  async updateBalance(accountId: string): Promise<void> {
+  async #updateBalance(
+    accountId: string,
+    assets: CaipAssetType[],
+  ): Promise<void> {
     try {
       const account = this.#getAccount(accountId);
 
       if (account.metadata.snap) {
-        const scope = getScopeForAccount(account);
-        const assetTypes = NETWORK_ASSETS_MAP[scope];
-
         const accountBalance = await this.#getBalances(
           account.id,
           account.metadata.snap.id,
-          assetTypes,
+          assets,
         );
 
         this.update((state: Draft<MultichainBalancesControllerState>) => {
           state.balances[accountId] = accountBalance;
         });
       }
     } catch (error) {
+      // FIXME: Maybe we shouldn't catch all errors here since this method is also being
+      // used in the public methods. This means if something else uses `updateBalance` it
+      // won't be able to catch and gets the error itself...
       console.error(
         `Failed to fetch balances for account ${accountId}:`,
         error,
       );
     }
   }
 
+  /**
+   * Updates the balances of one account. This method doesn't return
+   * anything, but it updates the state of the controller.
+   *
+   * @param accountId - The account ID.
+   */
+  async updateBalance(accountId: string): Promise<void> {
+    await this.#updateBalance(accountId, this.#listAccountAssets(accountId));
+  }
+
   /**
    * Lists the multichain accounts coming from the `AccountsController`.
    *
@@ -229,6 +256,21 @@ export class MultichainBalancesController extends BaseController<
     return accounts.filter((account) => this.#isNonEvmAccount(account));
   }
 
+  /**
+   * Lists the accounts assets.
+   *
+   * @param accountId - The account ID.
+   * @returns The list of assets for this account, returns an empty list if none.
+   */
+  #listAccountAssets(accountId: string): CaipAssetType[] {
+    // TODO: Add an action `MultichainAssetsController:getAccountAssets` maybe?
+    const assetsState = this.messagingSystem.call(
+      'MultichainAssetsController:getState',
+    );
+
+    return assetsState.accountsAssets[accountId] ?? [];
+  }
+
   /**
    * Get a non-EVM account from its ID.
    *
@@ -261,19 +303,6 @@ export class MultichainBalancesController extends BaseController<
     );
   }
 
-  /**
-   * Handles changes when a new account has been added.
-   *
-   * @param account - The new account being added.
-   */
-  async #handleOnAccountAdded(account: InternalAccount): Promise<void> {
-    if (!this.#isNonEvmAccount(account)) {
-      return;
-    }
-
-    await this.updateBalance(account.id);
-  }
-
   /**
    * Handles balance updates received from the AccountsController.
    *