fix(multichain): use accounts{Added,Removed} to fetch/clear balances

app/scripts/lib/accounts/BalancesController.test.ts

@@ -14,7 +14,7 @@ import {
   defaultState,
   BalancesControllerMessenger,
 } from './BalancesController';
-import { Poller } from './Poller';
+import { BalancesTracker } from './BalancesTracker';
 
 const mockBtcAccount = createMockInternalAccount({
   address: '',
@@ -54,8 +54,14 @@ const setupController = ({
   const balancesControllerMessenger: BalancesControllerMessenger =
     controllerMessenger.getRestricted({
       name: 'BalancesController',
-      allowedActions: ['SnapController:handleRequest'],
-      allowedEvents: ['AccountsController:stateChange'],
+      allowedActions: [
+        'SnapController:handleRequest',
+        'AccountsController:listMultichainAccounts',
+      ],
+      allowedEvents: [
+        'AccountsController:accountAdded',
+        'AccountsController:accountRemoved',
+      ],
     });
 
   const mockSnapHandleRequest = jest.fn();
@@ -66,20 +72,22 @@ const setupController = ({
     ),
   );
 
-  // TODO: remove when listMultichainAccounts action is available
-  const mockListMultichainAccounts = jest
-    .fn()
-    .mockReturnValue(mocks?.listMultichainAccounts ?? [mockBtcAccount]);
+  const mockListMultichainAccounts = jest.fn();
+  controllerMessenger.registerActionHandler(
+    'AccountsController:listMultichainAccounts',
+    mockListMultichainAccounts.mockReturnValue(
+      mocks?.listMultichainAccounts ?? [mockBtcAccount],
+    ),
+  );
 
   const controller = new BalancesController({
     messenger: balancesControllerMessenger,
     state,
-    // TODO: remove when listMultichainAccounts action is available
-    listMultichainAccounts: mockListMultichainAccounts,
   });
 
   return {
     controller,
+    messenger: controllerMessenger,
     mockSnapHandleRequest,
     mockListMultichainAccounts,
   };
@@ -91,19 +99,19 @@ describe('BalancesController', () => {
     expect(controller.state).toEqual({ balances: {} });
   });
 
-  it('starts polling when calling start', async () => {
-    const spyPoller = jest.spyOn(Poller.prototype, 'start');
+  it('starts tracking when calling start', async () => {
+    const spyTracker = jest.spyOn(BalancesTracker.prototype, 'start');
     const { controller } = setupController();
     await controller.start();
-    expect(spyPoller).toHaveBeenCalledTimes(1);
+    expect(spyTracker).toHaveBeenCalledTimes(1);
   });
 
-  it('stops polling when calling stop', async () => {
-    const spyPoller = jest.spyOn(Poller.prototype, 'stop');
+  it('stops tracking when calling stop', async () => {
+    const spyTracker = jest.spyOn(BalancesTracker.prototype, 'stop');
     const { controller } = setupController();
     await controller.start();
     await controller.stop();
-    expect(spyPoller).toHaveBeenCalledTimes(1);
+    expect(spyTracker).toHaveBeenCalledTimes(1);
   });
 
   it('update balances when calling updateBalances', async () => {
@@ -113,13 +121,49 @@ describe('BalancesController', () => {
 
     expect(controller.state).toEqual({
       balances: {
-        [mockBtcAccount.id]: {
-          'bip122:000000000933ea01ad0ee984209779ba/slip44:0': {
-            amount: '0.00000000',
-            unit: 'BTC',
-          },
+        [mockBtcAccount.id]: mockBalanceResult,
+      },
+    });
+  });
+
+  it.only('update balances when "AccountsController:accountAdded" is fired', async () => {
+    const { controller, messenger, mockListMultichainAccounts } =
+      setupController({
+        mocks: {
+          listMultichainAccounts: [],
         },
+      });
+
+    controller.start();
+    mockListMultichainAccounts.mockReturnValue([mockBtcAccount]);
+    messenger.publish('AccountsController:accountAdded', mockBtcAccount);
+    await controller.updateBalances();
+
+    expect(controller.state).toEqual({
+      balances: {
+        [mockBtcAccount.id]: mockBalanceResult,
       },
     });
   });
+
+  it.only('update balances when "AccountsController:accountRemoved" is fired', async () => {
+    const { controller, messenger, mockListMultichainAccounts } =
+      setupController();
+
+    controller.start();
+    await controller.updateBalances();
+    expect(controller.state).toEqual({
+      balances: {
+        [mockBtcAccount.id]: mockBalanceResult,
+      },
+    });
+
+    messenger.publish('AccountsController:accountRemoved', mockBtcAccount.id);
+    mockListMultichainAccounts.mockReturnValue([]);
+    await controller.updateBalances();
+
+    expect(controller.state).toEqual({
+      balances: {},
+    });
+  });
 });

app/scripts/lib/accounts/BalancesController.ts

@@ -19,11 +19,12 @@ import type { SnapId } from '@metamask/snaps-sdk';
 import { HandlerType } from '@metamask/snaps-utils';
 import type { Draft } from 'immer';
 import type {
-  AccountsControllerChangeEvent,
-  AccountsControllerState,
+  AccountsControllerAccountAddedEvent,
+  AccountsControllerAccountRemovedEvent,
+  AccountsControllerListMultichainAccountsAction,
 } from '@metamask/accounts-controller';
 import { isBtcMainnetAddress } from '../../../../shared/lib/multichain';
-import { Poller } from './Poller';
+import { BalancesTracker } from './BalancesTracker';
 
 const controllerName = 'BalancesController';
 
@@ -85,12 +86,16 @@ export type BalancesControllerEvents = BalancesControllerStateChange;
 /**
  * Actions that this controller is allowed to call.
  */
-export type AllowedActions = HandleSnapRequest;
+export type AllowedActions =
+  | HandleSnapRequest
+  | AccountsControllerListMultichainAccountsAction;
 
 /**
  * Events that this controller is allowed to subscribe.
  */
-export type AllowedEvents = AccountsControllerChangeEvent;
+export type AllowedEvents =
+  | AccountsControllerAccountAddedEvent
+  | AccountsControllerAccountRemovedEvent;
 
 /**
  * Messenger type for the BalancesController.
@@ -119,7 +124,7 @@ const balancesControllerMetadata = {
 
 const BTC_TESTNET_ASSETS = ['bip122:000000000933ea01ad0ee984209779ba/slip44:0'];
 const BTC_MAINNET_ASSETS = ['bip122:000000000019d6689c085ae165831e93/slip44:0'];
-export const BTC_AVG_BLOCK_TIME = 600000; // 10 minutes in milliseconds
+export const BTC_AVG_BLOCK_TIME = 10 * 60 * 1000; // 10 minutes in milliseconds
 
 /**
  * The BalancesController is responsible for fetching and caching account
@@ -130,19 +135,14 @@ export class BalancesController extends BaseController<
   BalancesControllerState,
   BalancesControllerMessenger
 > {
-  #poller: Poller;
-
-  // TODO: remove once action is implemented
-  #listMultichainAccounts: () => InternalAccount[];
+  #tracker: BalancesTracker;
 
   constructor({
     messenger,
     state,
-    listMultichainAccounts,
   }: {
     messenger: BalancesControllerMessenger;
     state: BalancesControllerState;
-    listMultichainAccounts: () => InternalAccount[];
   }) {
     super({
       messenger,
@@ -154,27 +154,62 @@ export class BalancesController extends BaseController<
       },
     });
 
+    this.#tracker = new BalancesTracker(async (accountId: string) => {
+      // The BalancesTracker only uses account IDs, so we have to get the associated account first
+      const account = this.#listMultichainAccounts().find(
+        (multichainAccount) => multichainAccount.id === accountId,
+      );
+      if (!account) {
+        throw new Error(`Unknown account: ${accountId}`);
+      }
+      // Just to double-check, we make sure we are using a non-EVM account
+      if (!this.#isNonEvmAccount(account)) {
+        throw new Error(`Account is not a non-EVM account: ${accountId}`);
+      }
+
+      await this.#updateBalance(account);
+    });
+
+    // Register all non-EVM accounts into the tracker
+    for (const account of this.#listAccounts()) {
+      if (this.#isNonEvmAccount(account)) {
+        this.#tracker.track(account.id, BTC_AVG_BLOCK_TIME);
+      }
+    }
+
     this.messagingSystem.subscribe(
-      'AccountsController:stateChange',
-      (newState) => this.#handleOnAccountsControllerChange(newState),
+      'AccountsController:accountAdded',
+      (account) => this.#handleOnAccountAdded(account),
+    );
+    this.messagingSystem.subscribe(
+      'AccountsController:accountRemoved',
+      (account) => this.#handleOnAccountRemoved(account),
     );
-
-    this.#listMultichainAccounts = listMultichainAccounts;
-    this.#poller = new Poller(() => this.updateBalances(), BTC_AVG_BLOCK_TIME);
   }
 
   /**
    * Starts the polling process.
    */
   async start(): Promise<void> {
-    this.#poller.start();
+    this.#tracker.start();
   }
 
   /**
    * Stops the polling process.
    */
   async stop(): Promise<void> {
-    this.#poller.stop();
+    this.#tracker.stop();
+  }
+
+  /**
+   * Lists the multichain accounts coming from the `AccountsController`.
+   *
+   * @returns A list of multichain accounts.
+   */
+  #listMultichainAccounts(): InternalAccount[] {
+    return this.messagingSystem.call(
+      'AccountsController:listMultichainAccounts',
+    );
   }
 
   /**
@@ -185,50 +220,102 @@ export class BalancesController extends BaseController<
    *
    * @returns A list of accounts that we should get balances for.
    */
-  async #listAccounts(): Promise<InternalAccount[]> {
+  #listAccounts(): InternalAccount[] {
     const accounts = this.#listMultichainAccounts();
 
     return accounts.filter((account) => account.type === BtcAccountType.P2wpkh);
   }
 
   /**
-   * Updates the balances of all supported accounts. This method doesn't return
+   * Updates the balances of one account. This method doesn't return
    * anything, but it updates the state of the controller.
+   *
+   * **NOTE**: This method assumes the given account is a non-EVM account
+   * associated with a non-EVM Snap.
+   *
+   * @param account - The account.
    */
-  async updateBalances() {
-    const accounts = await this.#listAccounts();
+  async #updateBalance(account: InternalAccount) {
     const partialState: BalancesControllerState = { balances: {} };
 
-    for (const account of accounts) {
-      if (account.metadata.snap) {
-        partialState.balances[account.id] = await this.#getBalances(
-          account.id,
-          account.metadata.snap.id,
-          isBtcMainnetAddress(account.address)
-            ? BTC_MAINNET_ASSETS
-            : BTC_TESTNET_ASSETS,
-        );
-      }
-    }
+    partialState.balances[account.id] = await this.#getBalances(
+      account.id,
+      account.metadata.snap.id,
+      isBtcMainnetAddress(account.address)
+        ? BTC_MAINNET_ASSETS
+        : BTC_TESTNET_ASSETS,
+    );
 
     this.update((state: Draft<BalancesControllerState>) => ({
       ...state,
-      ...partialState,
+      balances: {
+        ...state.balances,
+        ...partialState.balances,
+      },
     }));
   }
 
   /**
-   * Handles changes in the accounts state, specifically when new non-EVM accounts are added.
+   * 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) {
+    await this.#tracker.updateBalance(accountId);
+  }
+
+  /**
+   * Updates the balances of all supported accounts. This method doesn't return
+   * anything, but it updates the state of the controller.
+   */
+  async updateBalances() {
+    await this.#tracker.updateBalances();
+  }
+
+  /**
+   * Checks for non-EVM accounts.
    *
-   * @param newState - The new state of the accounts controller.
+   * @param account - The new account to be checked.
+   * @returns True if the account is a non-EVM account, false otherwise.
    */
-  #handleOnAccountsControllerChange(newState: AccountsControllerState) {
-    // If we have any new non-EVM accounts, we just update non-EVM balances
-    const newNonEvmAccounts = Object.values(
-      newState.internalAccounts.accounts,
-    ).filter((account) => !isEvmAccountType(account.type));
-    if (newNonEvmAccounts.length) {
-      this.updateBalances();
+  #isNonEvmAccount(account: InternalAccount): boolean {
+    return (
+      !isEvmAccountType(account.type) &&
+      // Non-EVM accounts are backed by a Snap for now
+      account.metadata.snap
+    );
+  }
+
+  /**
+   * Handles changes when a new account has been added.
+   *
+   * @param account - The new account being added.
+   */
+  async #handleOnAccountAdded(account: InternalAccount) {
+    if (!this.#isNonEvmAccount(account)) {
+      // Nothing to do here for EVM accounts
+      return;
+    }
+
+    this.#tracker.track(account.id, BTC_AVG_BLOCK_TIME);
+  }
+
+  /**
+   * Handles changes when a new account has been removed.
+   *
+   * @param accountId - The account ID being removed.
+   */
+  async #handleOnAccountRemoved(accountId: string) {
+    if (this.#tracker.isTracked(accountId)) {
+      this.#tracker.untrack(accountId);
+    }
+
+    if (accountId in this.state.balances) {
+      this.update((state: Draft<BalancesControllerState>) => {
+        delete state.balances[accountId];
+        return state;
+      });
     }
   }