chore: get assets list from MultichainAssetsController state (#5295)

## Explanation

We are getting the full list of assets for an account via the
MultichainAssetsController state, instead of only requesting the native
token.

## 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
-->

## Changelog

<!--
If you're making any consumer-facing changes, list those changes here as
if you were updating a changelog, using the template below as a guide.

(CATEGORY is one of BREAKING, ADDED, CHANGED, DEPRECATED, REMOVED, or
FIXED. For security-related issues, follow the Security Advisory
process.)

Please take care to name the exact pieces of the API you've added or
changed (e.g. types, interfaces, functions, or methods).

If there are any breaking changes, make sure to offer a solution for
consumers to follow once they upgrade to the changes.

Finally, if you're only making changes to development scripts or tests,
you may replace the template below with "None".
-->

### `@metamask/assets-controllers`

- **ADDED**: Adds the ability to request for all the assets that an
account can have


## Checklist

- [x] I've updated the test suite for new or updated code as appropriate
- [x] I've updated documentation (JSDoc, Markdown, etc.) for new or
updated code as appropriate
- [x] I've highlighted breaking changes using the "BREAKING" category
above as appropriate
- [x] I've prepared draft pull requests for clients and consumer
packages to resolve any breaking changes

---------

Co-authored-by: Charly Chevalier <charly.chevalier@consensys.net>

Author: zone-live

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.
    *