From 262616bb215acde6f91f5f2cf7676828a2509061 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Wed, 17 Jul 2024 12:02:48 +0200 Subject: [PATCH 01/19] fix(multichain): use accounts{Added,Removed} to fetch/clear balances --- .../lib/accounts/BalancesController.ts | 96 +++++++++++++++---- app/scripts/metamask-controller.js | 2 +- 2 files changed, 77 insertions(+), 21 deletions(-) diff --git a/app/scripts/lib/accounts/BalancesController.ts b/app/scripts/lib/accounts/BalancesController.ts index ab1eb8c6cfe6..d33a2b1ac9f6 100644 --- a/app/scripts/lib/accounts/BalancesController.ts +++ b/app/scripts/lib/accounts/BalancesController.ts @@ -19,8 +19,8 @@ import type { SnapId } from '@metamask/snaps-sdk'; import { HandlerType } from '@metamask/snaps-utils'; import type { Draft } from 'immer'; import type { - AccountsControllerChangeEvent, - AccountsControllerState, + AccountsControllerAccountAddedEvent, + AccountsControllerAccountRemovedEvent, } from '@metamask/accounts-controller'; import { isBtcMainnetAddress } from '../../../../shared/lib/multichain'; import { Poller } from './Poller'; @@ -90,7 +90,9 @@ export type AllowedActions = HandleSnapRequest; /** * Events that this controller is allowed to subscribe. */ -export type AllowedEvents = AccountsControllerChangeEvent; +export type AllowedEvents = + | AccountsControllerAccountAddedEvent + | AccountsControllerAccountRemovedEvent; /** * Messenger type for the BalancesController. @@ -155,8 +157,12 @@ export class BalancesController extends BaseController< }); 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; @@ -191,6 +197,20 @@ export class BalancesController extends BaseController< return accounts.filter((account) => account.type === BtcAccountType.P2wpkh); } + /** + * Updates the balances of all supported accounts. This method doesn't return + * anything, but it updates the state of the controller. + */ + async #getBalancesForAccount(account: InternalAccount) { + return await this.#getBalances( + account.id, + account.metadata.snap.id, + isBtcMainnetAddress(account.address) + ? BTC_MAINNET_ASSETS + : BTC_TESTNET_ASSETS, + ); + } + /** * Updates the balances of all supported accounts. This method doesn't return * anything, but it updates the state of the controller. @@ -201,12 +221,8 @@ export class BalancesController extends BaseController< 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.#getBalancesForAccount( + account, ); } } @@ -218,17 +234,57 @@ export class BalancesController extends BaseController< } /** - * Handles changes in the accounts state, specifically when new non-EVM accounts are added. + * Checks for non-EVM accounts. + * + * @param account - The new account to be checked. + * @returns True if the account is a non-EVM account, false otherwise. + */ + #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 newState - The new state of the accounts controller. + * @param account - The new account being added. */ - #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(); + async #handleOnAccountAdded(account: InternalAccount) { + if (!this.#isNonEvmAccount(account)) { + // Nothing to do here for EVM accounts + return; + } + + const partialState: BalancesControllerState = { balances: {} }; + + // Update the balance only for the newly added account + partialState.balances[account.id] = await this.#getBalancesForAccount( + account, + ); + + this.update((state: Draft) => ({ + ...state, + ...partialState, + })); + } + + /** + * Handles changes when a new account has been removed. + * + * @param account - The new account being removed. + */ + async #handleOnAccountRemoved(accountId: string) { + // We still check if the accounts has a balance here to avoid non-necessary state + // update + if (accountId in this.state.balances) { + this.update((state: Draft) => { + delete state.balances[accountId]; + + return state; + }); } } diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js index 59dfc9d96540..bf3e241bcf18 100644 --- a/app/scripts/metamask-controller.js +++ b/app/scripts/metamask-controller.js @@ -876,7 +876,7 @@ export default class MetamaskController extends EventEmitter { const multichainBalancesControllerMessenger = this.controllerMessenger.getRestricted({ name: 'BalancesController', - allowedEvents: ['AccountsController:stateChange'], + allowedEvents: ['AccountsController:accountAdded', 'AccountsController:accountRemoved'], allowedActions: ['SnapController:handleRequest'], }); From 9e959df25e2c83d7635869f8fcd52ed64b20fd8e Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Wed, 17 Jul 2024 16:23:34 +0200 Subject: [PATCH 02/19] feat(multichain): now tracks non-EVM balances to relax excessive fetches --- .../lib/accounts/BalancesController.ts | 106 +++++++++------- app/scripts/lib/accounts/BalancesTracker.ts | 119 ++++++++++++++++++ app/scripts/metamask-controller.js | 6 +- 3 files changed, 187 insertions(+), 44 deletions(-) create mode 100644 app/scripts/lib/accounts/BalancesTracker.ts diff --git a/app/scripts/lib/accounts/BalancesController.ts b/app/scripts/lib/accounts/BalancesController.ts index d33a2b1ac9f6..530ca92d1507 100644 --- a/app/scripts/lib/accounts/BalancesController.ts +++ b/app/scripts/lib/accounts/BalancesController.ts @@ -21,9 +21,11 @@ import type { Draft } from 'immer'; import type { AccountsControllerAccountAddedEvent, AccountsControllerAccountRemovedEvent, + AccountsControllerChangeEvent, + AccountsControllerState, } from '@metamask/accounts-controller'; import { isBtcMainnetAddress } from '../../../../shared/lib/multichain'; -import { Poller } from './Poller'; +import { BalancesTracker } from './BalancesTracker'; const controllerName = 'BalancesController'; @@ -92,7 +94,8 @@ export type AllowedActions = HandleSnapRequest; */ export type AllowedEvents = | AccountsControllerAccountAddedEvent - | AccountsControllerAccountRemovedEvent; + | AccountsControllerAccountRemovedEvent + | AccountsControllerChangeEvent; /** * Messenger type for the BalancesController. @@ -121,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 @@ -132,7 +135,7 @@ export class BalancesController extends BaseController< BalancesControllerState, BalancesControllerMessenger > { - #poller: Poller; + #tracker: BalancesTracker; // TODO: remove once action is implemented #listMultichainAccounts: () => InternalAccount[]; @@ -156,6 +159,23 @@ export class BalancesController extends BaseController< }, }); + this.#listMultichainAccounts = listMultichainAccounts; + + 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}`); + } + + await this.#updateBalancesForAccount(account); + }); + for (const account of this.#listAccounts()) { + this.#tracker.track(account.id, BTC_AVG_BLOCK_TIME); + } + this.messagingSystem.subscribe( 'AccountsController:accountAdded', (account) => this.#handleOnAccountAdded(account), @@ -164,23 +184,34 @@ export class BalancesController extends BaseController< 'AccountsController:accountRemoved', (account) => this.#handleOnAccountRemoved(account), ); - - this.#listMultichainAccounts = listMultichainAccounts; - this.#poller = new Poller(() => this.updateBalances(), BTC_AVG_BLOCK_TIME); + this.messagingSystem.subscribe( + 'AccountsController:stateChange', + (_newState: AccountsControllerState) => { + // The tracker won't refresh the balance if it's not required, so there's + // very little overhead of using it here. + // + // However, updating the balances here allow us to fetch the balance of any new + // created account directly (we start tracking in `:accountAdded` event handler). + // + // In this case, the tracker will fetch the balance (for the first time) of those + // new accounts. + this.#tracker.updateBalances(); + }, + ); } /** * Starts the polling process. */ async start(): Promise { - this.#poller.start(); + this.#tracker.start(); } /** * Stops the polling process. */ async stop(): Promise { - this.#poller.stop(); + this.#tracker.stop(); } /** @@ -191,24 +222,36 @@ export class BalancesController extends BaseController< * * @returns A list of accounts that we should get balances for. */ - async #listAccounts(): Promise { + #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. + * + * @param account - The account. */ - async #getBalancesForAccount(account: InternalAccount) { - return await this.#getBalances( + async #updateBalancesForAccount(account: InternalAccount) { + const partialState: BalancesControllerState = { balances: {} }; + + 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) => ({ + ...state, + balances: { + ...state.balances, + ...partialState.balances, + }, + })); } /** @@ -216,21 +259,7 @@ export class BalancesController extends BaseController< * anything, but it updates the state of the controller. */ async updateBalances() { - const accounts = await this.#listAccounts(); - const partialState: BalancesControllerState = { balances: {} }; - - for (const account of accounts) { - if (account.metadata.snap) { - partialState.balances[account.id] = await this.#getBalancesForAccount( - account, - ); - } - } - - this.update((state: Draft) => ({ - ...state, - ...partialState, - })); + await this.#tracker.updateBalances(); } /** @@ -258,31 +287,22 @@ export class BalancesController extends BaseController< return; } - const partialState: BalancesControllerState = { balances: {} }; - - // Update the balance only for the newly added account - partialState.balances[account.id] = await this.#getBalancesForAccount( - account, - ); - - this.update((state: Draft) => ({ - ...state, - ...partialState, - })); + this.#tracker.track(account.id, BTC_AVG_BLOCK_TIME); } /** * Handles changes when a new account has been removed. * - * @param account - The new account being removed. + * @param accountId - The account ID being removed. */ async #handleOnAccountRemoved(accountId: string) { - // We still check if the accounts has a balance here to avoid non-necessary state - // update + if (this.#tracker.isTracked(accountId)) { + this.#tracker.untrack(accountId); + } + if (accountId in this.state.balances) { this.update((state: Draft) => { delete state.balances[accountId]; - return state; }); } diff --git a/app/scripts/lib/accounts/BalancesTracker.ts b/app/scripts/lib/accounts/BalancesTracker.ts new file mode 100644 index 000000000000..eac60ab6854d --- /dev/null +++ b/app/scripts/lib/accounts/BalancesTracker.ts @@ -0,0 +1,119 @@ +import { Poller } from './Poller'; + +type BalanceInfo = { + lastUpdated: number; + blockTime: number; +}; + +const BALANCES_TRACKING_INTERVAL = 30 * 1000; // Every 30s in milliseconds. + +export class BalancesTracker { + #poller: Poller; + + #updateBalance: (accountId: string) => Promise; + + #balances: Record = {}; + + constructor(updateBalanceCallback: (accountId: string) => Promise) { + this.#updateBalance = updateBalanceCallback; + + this.#poller = new Poller(() => { + this.updateBalances(); + }, BALANCES_TRACKING_INTERVAL); + } + + /** + * Starts the tracking process. + */ + async start(): Promise { + this.#poller.start(); + } + + /** + * Stops the tracking process. + */ + async stop(): Promise { + this.#poller.stop(); + } + + /** + * Checks if an account ID is being tracked. + * + * @param accountId - The account ID. + * @returns True if the account is being tracker, false otherwise. + */ + isTracked(accountId: string) { + return accountId in this.#balances; + } + + /** + * Asserts that an account ID is being tracked. + * + * @param accountId - The account ID. + * @throws If the account ID is not being tracked. + */ + assertBeingTracked(accountId: string) { + if (!this.isTracked(accountId)) { + throw new Error(`Account is not being tracked: ${accountId}`); + } + } + + /** + * Starts tracking a new account ID. This method has no effect on already tracked + * accounts. + * + * @param accountId - The account ID. + * @param blockTime - The block time (used when refreshing the account balances). + */ + track(accountId: string, blockTime: number) { + // Do not overwrite current info if already being tracked! + if (!this.isTracked(accountId)) { + this.#balances[accountId] = { + lastUpdated: 0, + blockTime, + }; + } + } + + /** + * Stops tracking a tracked account ID. + * + * @param accountId - The account ID. + * @throws If the account ID is not being tracked. + */ + untrack(accountId: string) { + this.assertBeingTracked(accountId); + delete this.#balances[accountId]; + } + + /** + * Update the balances for a tracked account ID. + * + * @param accountId - The account ID. + * @throws If the account ID is not being tracked. + */ + async updateBalance(accountId: string) { + this.assertBeingTracked(accountId); + await this.#updateBalance(accountId); + this.#balances[accountId].lastUpdated = Date.now(); + } + + /** + * Update the balances of all tracked accounts (only if the balances + * is considered outdated). + */ + async updateBalances() { + for (const [accountId, info] of Object.entries(this.#balances)) { + // We check if the balance is outdated (by comparing to the block time associated + // with this kind of account). + // + // This might not be super accurate, but we could probably compute this differently + // and try to sync with the "real block time"! + const isOutdated = Date.now() - info.lastUpdated >= info.blockTime; + + if (isOutdated) { + await this.updateBalance(accountId); + } + } + } +} diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js index bf3e241bcf18..1b27087e766d 100644 --- a/app/scripts/metamask-controller.js +++ b/app/scripts/metamask-controller.js @@ -876,7 +876,11 @@ export default class MetamaskController extends EventEmitter { const multichainBalancesControllerMessenger = this.controllerMessenger.getRestricted({ name: 'BalancesController', - allowedEvents: ['AccountsController:accountAdded', 'AccountsController:accountRemoved'], + allowedEvents: [ + 'AccountsController:accountAdded', + 'AccountsController:accountRemoved', + 'AccountsController:stateChange', + ], allowedActions: ['SnapController:handleRequest'], }); From 17a2e96de64b37bb8afccf9dce2daf34e51d1ae2 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Wed, 17 Jul 2024 18:51:50 +0200 Subject: [PATCH 03/19] test(multichain): fix + add more tests for BalancesController --- .../lib/accounts/BalancesController.test.ts | 67 +++++++++++++++---- 1 file changed, 54 insertions(+), 13 deletions(-) diff --git a/app/scripts/lib/accounts/BalancesController.test.ts b/app/scripts/lib/accounts/BalancesController.test.ts index 6bfe0bfabfc6..a65047b767b8 100644 --- a/app/scripts/lib/accounts/BalancesController.test.ts +++ b/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: '', @@ -55,7 +55,11 @@ const setupController = ({ controllerMessenger.getRestricted({ name: 'BalancesController', allowedActions: ['SnapController:handleRequest'], - allowedEvents: ['AccountsController:stateChange'], + allowedEvents: [ + 'AccountsController:accountAdded', + 'AccountsController:accountRemoved', + 'AccountsController:stateChange', + ], }); const mockSnapHandleRequest = jest.fn(); @@ -80,6 +84,7 @@ const setupController = ({ return { controller, + messenger: controllerMessenger, mockSnapHandleRequest, mockListMultichainAccounts, }; @@ -91,19 +96,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 +118,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: {}, + }); }); }); From 00cc6bea012f02a25b57bdc9451b7231065e88b2 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Wed, 17 Jul 2024 18:54:54 +0200 Subject: [PATCH 04/19] test(multichain): add tests for BalancesTracker --- .../lib/accounts/BalancesTracker.test.ts | 128 ++++++++++++++++++ 1 file changed, 128 insertions(+) create mode 100644 app/scripts/lib/accounts/BalancesTracker.test.ts diff --git a/app/scripts/lib/accounts/BalancesTracker.test.ts b/app/scripts/lib/accounts/BalancesTracker.test.ts new file mode 100644 index 000000000000..060cdf477530 --- /dev/null +++ b/app/scripts/lib/accounts/BalancesTracker.test.ts @@ -0,0 +1,128 @@ +import { BtcAccountType } from '@metamask/keyring-api'; +import { createMockInternalAccount } from '../../../../test/jest/mocks'; +import { Poller } from './Poller'; +import { BalancesTracker } from './BalancesTracker'; + +const mockBtcAccount = createMockInternalAccount({ + address: '', + name: 'Btc Account', + // @ts-expect-error - account type may be btc or eth, mock file is not typed + type: BtcAccountType.P2wpkh, + // @ts-expect-error - snap options is not typed and defaults to undefined + snapOptions: { + id: 'mock-btc-snap', + name: 'mock-btc-snap', + enabled: true, + }, +}); + +const mockBalanceResult = { + 'bip122:000000000933ea01ad0ee984209779ba/slip44:0': { + amount: '0.00000000', + unit: 'BTC', + }, +}; + +const MOCK_TIMESTAMP = 1709983353; + +function setupTracker() { + const mockUpdateBalance = jest.fn(); + const tracker = new BalancesTracker(mockUpdateBalance); + + return { + tracker, + mockUpdateBalance, + }; +} + +describe('BalancesTracker', () => { + it('starts polling when calling start', async () => { + const { tracker } = setupTracker(); + const spyPoller = jest.spyOn(Poller.prototype, 'start'); + + await tracker.start(); + expect(spyPoller).toHaveBeenCalledTimes(1); + }); + + it('stops polling when calling stop', async () => { + const { tracker } = setupTracker(); + const spyPoller = jest.spyOn(Poller.prototype, 'stop'); + + await tracker.start(); + await tracker.stop(); + expect(spyPoller).toHaveBeenCalledTimes(1); + }); + + it('is not tracking if none accounts have been registered', async () => { + const { tracker, mockUpdateBalance } = setupTracker(); + + await tracker.start(); + await tracker.updateBalances(); + + expect(mockUpdateBalance).not.toHaveBeenCalled(); + }); + + it('tracks account balances', async () => { + const { tracker, mockUpdateBalance } = setupTracker(); + + await tracker.start(); + // We must track account IDs explicitly + tracker.track(mockBtcAccount.id, 0); + // Trigger balances refresh (not waiting for the Poller here) + await tracker.updateBalances(); + + expect(mockUpdateBalance).toHaveBeenCalledWith(mockBtcAccount.id); + }); + + it('untracks account balances', async () => { + const { tracker, mockUpdateBalance } = setupTracker(); + + await tracker.start(); + tracker.track(mockBtcAccount.id, 0); + await tracker.updateBalances(); + expect(mockUpdateBalance).toHaveBeenCalledWith(mockBtcAccount.id); + + tracker.untrack(mockBtcAccount.id); + await tracker.updateBalances(); + expect(mockUpdateBalance).toHaveBeenCalledTimes(1); // No second call after untracking + }); + + it('tracks account after being registered', async () => { + const { tracker } = setupTracker(); + + await tracker.start(); + tracker.track(mockBtcAccount.id, 0); + expect(tracker.isTracked(mockBtcAccount.id)).toBe(true); + }); + + it('does not track account if not registered', async () => { + const { tracker } = setupTracker(); + + await tracker.start(); + expect(tracker.isTracked(mockBtcAccount.id)).toBe(false); + }); + + it('does not refresh balance if they are considered up-to-date', async () => { + const { tracker, mockUpdateBalance } = setupTracker(); + + const blockTime = 10 * 60 * 1000; // 10 minutes in milliseconds. + jest + .spyOn(global.Date, 'now') + .mockImplementation(() => new Date(MOCK_TIMESTAMP).getTime()); + + await tracker.start(); + tracker.track(mockBtcAccount.id, blockTime); + await tracker.updateBalances(); + expect(mockUpdateBalance).toHaveBeenCalledTimes(1); + + await tracker.updateBalances(); + expect(mockUpdateBalance).toHaveBeenCalledTimes(1); // No second call since the balances is already still up-to-date + + jest + .spyOn(global.Date, 'now') + .mockImplementation(() => new Date(MOCK_TIMESTAMP + blockTime).getTime()); + + await tracker.updateBalances(); + expect(mockUpdateBalance).toHaveBeenCalledTimes(2); // Now the balance will update + }); +}); From d9c01bb5a8663614600c7cf4d71179ae1f1e4418 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Wed, 17 Jul 2024 22:28:17 +0200 Subject: [PATCH 05/19] chore: lint --- app/scripts/lib/accounts/BalancesTracker.test.ts | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/app/scripts/lib/accounts/BalancesTracker.test.ts b/app/scripts/lib/accounts/BalancesTracker.test.ts index 060cdf477530..27c862aa4460 100644 --- a/app/scripts/lib/accounts/BalancesTracker.test.ts +++ b/app/scripts/lib/accounts/BalancesTracker.test.ts @@ -3,6 +3,8 @@ import { createMockInternalAccount } from '../../../../test/jest/mocks'; import { Poller } from './Poller'; import { BalancesTracker } from './BalancesTracker'; +const MOCK_TIMESTAMP = 1709983353; + const mockBtcAccount = createMockInternalAccount({ address: '', name: 'Btc Account', @@ -16,15 +18,6 @@ const mockBtcAccount = createMockInternalAccount({ }, }); -const mockBalanceResult = { - 'bip122:000000000933ea01ad0ee984209779ba/slip44:0': { - amount: '0.00000000', - unit: 'BTC', - }, -}; - -const MOCK_TIMESTAMP = 1709983353; - function setupTracker() { const mockUpdateBalance = jest.fn(); const tracker = new BalancesTracker(mockUpdateBalance); From bd3dda91e2740a5213bd90bc46a51386190245c4 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Thu, 18 Jul 2024 10:42:16 +0200 Subject: [PATCH 06/19] refactor(multichain): use action rather than AccountsController.listMultichainAccounts callback --- .../lib/accounts/BalancesController.test.ts | 18 +++++++++------ .../lib/accounts/BalancesController.ts | 23 ++++++++++++------- app/scripts/metamask-controller.js | 10 ++++---- 3 files changed, 30 insertions(+), 21 deletions(-) diff --git a/app/scripts/lib/accounts/BalancesController.test.ts b/app/scripts/lib/accounts/BalancesController.test.ts index a65047b767b8..41d30ffa3bf5 100644 --- a/app/scripts/lib/accounts/BalancesController.test.ts +++ b/app/scripts/lib/accounts/BalancesController.test.ts @@ -54,7 +54,10 @@ const setupController = ({ const balancesControllerMessenger: BalancesControllerMessenger = controllerMessenger.getRestricted({ name: 'BalancesController', - allowedActions: ['SnapController:handleRequest'], + allowedActions: [ + 'SnapController:handleRequest', + 'AccountsController:listMultichainAccounts', + ], allowedEvents: [ 'AccountsController:accountAdded', 'AccountsController:accountRemoved', @@ -70,16 +73,17 @@ 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 { diff --git a/app/scripts/lib/accounts/BalancesController.ts b/app/scripts/lib/accounts/BalancesController.ts index 530ca92d1507..35c0c3a12573 100644 --- a/app/scripts/lib/accounts/BalancesController.ts +++ b/app/scripts/lib/accounts/BalancesController.ts @@ -22,6 +22,7 @@ import type { AccountsControllerAccountAddedEvent, AccountsControllerAccountRemovedEvent, AccountsControllerChangeEvent, + AccountsControllerListMultichainAccountsAction, AccountsControllerState, } from '@metamask/accounts-controller'; import { isBtcMainnetAddress } from '../../../../shared/lib/multichain'; @@ -87,7 +88,9 @@ 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. @@ -137,17 +140,12 @@ export class BalancesController extends BaseController< > { #tracker: BalancesTracker; - // TODO: remove once action is implemented - #listMultichainAccounts: () => InternalAccount[]; - constructor({ messenger, state, - listMultichainAccounts, }: { messenger: BalancesControllerMessenger; state: BalancesControllerState; - listMultichainAccounts: () => InternalAccount[]; }) { super({ messenger, @@ -159,8 +157,6 @@ export class BalancesController extends BaseController< }, }); - this.#listMultichainAccounts = listMultichainAccounts; - 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( @@ -214,6 +210,17 @@ export class BalancesController extends BaseController< 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', + ); + } + /** * Lists the accounts that we should get balances for. * diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js index 1b27087e766d..59e4ea8209bb 100644 --- a/app/scripts/metamask-controller.js +++ b/app/scripts/metamask-controller.js @@ -881,17 +881,15 @@ export default class MetamaskController extends EventEmitter { 'AccountsController:accountRemoved', 'AccountsController:stateChange', ], - allowedActions: ['SnapController:handleRequest'], + allowedActions: [ + 'AccountsController:listMultichainAccounts', + 'SnapController:handleRequest', + ], }); this.multichainBalancesController = new MultichainBalancesController({ messenger: multichainBalancesControllerMessenger, state: {}, - // TODO: remove when listMultichainAccounts action is available - listMultichainAccounts: - this.accountsController.listMultichainAccounts.bind( - this.accountsController, - ), }); const multichainRatesControllerMessenger = From ad712afc08fb46c7bde87b94160cbcd47328830c Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Thu, 18 Jul 2024 13:14:39 +0200 Subject: [PATCH 07/19] refactor(multichain): add public BalancesController.updateBalance method --- .../lib/accounts/BalancesController.ts | 23 +++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/app/scripts/lib/accounts/BalancesController.ts b/app/scripts/lib/accounts/BalancesController.ts index 35c0c3a12573..e15b538a8f8c 100644 --- a/app/scripts/lib/accounts/BalancesController.ts +++ b/app/scripts/lib/accounts/BalancesController.ts @@ -166,7 +166,9 @@ export class BalancesController extends BaseController< throw new Error(`Unknown account: ${accountId}`); } - await this.#updateBalancesForAccount(account); + // Just to double-check, we call updateBalance rather than #updateBalance to make sure + // we are using a non-EVM account + await this.updateBalance(account); }); for (const account of this.#listAccounts()) { this.#tracker.track(account.id, BTC_AVG_BLOCK_TIME); @@ -239,9 +241,12 @@ export class BalancesController extends BaseController< * 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 #updateBalancesForAccount(account: InternalAccount) { + async #updateBalance(account: InternalAccount) { const partialState: BalancesControllerState = { balances: {} }; partialState.balances[account.id] = await this.#getBalances( @@ -261,6 +266,20 @@ export class BalancesController extends BaseController< })); } + /** + * Updates the balances of one account. This method doesn't return + * anything, but it updates the state of the controller. + * + * @param account - The account. + */ + async updateBalance(account: InternalAccount) { + if (!this.#isNonEvmAccount(account)) { + throw new Error('Cannot update balance of EVM accounts from here'); + } + + this.#updateBalance(account); + } + /** * Updates the balances of all supported accounts. This method doesn't return * anything, but it updates the state of the controller. From c8e21df5590e9f27fbbff0cde42ccb24a3dbab9a Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Thu, 18 Jul 2024 13:15:08 +0200 Subject: [PATCH 08/19] fix(multichain): only tracks non-EVM accounts when constructing BalancesController --- app/scripts/lib/accounts/BalancesController.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/app/scripts/lib/accounts/BalancesController.ts b/app/scripts/lib/accounts/BalancesController.ts index e15b538a8f8c..f2d7eab41593 100644 --- a/app/scripts/lib/accounts/BalancesController.ts +++ b/app/scripts/lib/accounts/BalancesController.ts @@ -171,7 +171,9 @@ export class BalancesController extends BaseController< await this.updateBalance(account); }); for (const account of this.#listAccounts()) { - this.#tracker.track(account.id, BTC_AVG_BLOCK_TIME); + if (this.#isNonEvmAccount(account)) { + this.#tracker.track(account.id, BTC_AVG_BLOCK_TIME); + } } this.messagingSystem.subscribe( From 7008a686d41caa3776f7432df5f68122cc08799a Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Thu, 18 Jul 2024 13:22:46 +0200 Subject: [PATCH 09/19] test(multichain): fix metamask-controller.test.js with new BalancesController logic --- app/scripts/metamask-controller.test.js | 39 +++++++++++++++++++++++-- 1 file changed, 36 insertions(+), 3 deletions(-) diff --git a/app/scripts/metamask-controller.test.js b/app/scripts/metamask-controller.test.js index 1ad12736fe2a..45e2c3bb1b19 100644 --- a/app/scripts/metamask-controller.test.js +++ b/app/scripts/metamask-controller.test.js @@ -2240,12 +2240,27 @@ describe('MetaMaskController', () => { type: BtcAccountType.P2wpkh, methods: [BtcMethod.SendMany], address: 'bc1qar0srrr7xfkvy5l643lydnw9re59gtzzwf5mdq', + // We need to have a "Snap account" account here, since the MultichainBalancesController will + // filter it out otherwise! + metadata: { + name: 'Bitcoin Account', + importTime: Date.now(), + keyring: { + type: KeyringType.snap, + }, + snap: { + id: 'npm:@metamask/bitcoin-wallet-snap', + }, + }, }; let localMetamaskController; beforeEach(() => { jest.useFakeTimers(); jest.spyOn(MultichainBalancesController.prototype, 'updateBalances'); + jest + .spyOn(MultichainBalancesController.prototype, 'updateBalance') + .mockResolvedValue(); localMetamaskController = new MetaMaskController({ showUserConfirmation: noop, encryptor: mockEncryptor, @@ -2284,11 +2299,29 @@ describe('MetaMaskController', () => { }); it('calls updateBalances after the interval has passed', async () => { - jest.advanceTimersByTime(BTC_AVG_BLOCK_TIME); - // 2 calls because 1 is during startup + // 1st call is during startup: + // updatesBalances is going to call updateBalance for the only non-EVM + // account that we have expect( localMetamaskController.multichainBalancesController.updateBalances, - ).toHaveBeenCalledTimes(2); + ).toHaveBeenCalledTimes(1); + expect( + localMetamaskController.multichainBalancesController.updateBalance, + ).toHaveBeenCalledTimes(1); + expect( + localMetamaskController.multichainBalancesController.updateBalance, + ).toHaveBeenCalledWith(mockNonEvmAccount); + + // Wait for "block time", so balances will have to be refreshed + jest.advanceTimersByTime(BTC_AVG_BLOCK_TIME); + + // 2nd call because balances are considered "outdated": + expect( + localMetamaskController.multichainBalancesController.updateBalance, + ).toHaveBeenCalledTimes(2); // 1 (startup) + 1 (refresh) + expect( + localMetamaskController.multichainBalancesController.updateBalance, + ).toHaveBeenLastCalledWith(mockNonEvmAccount); }); }); }); From 06c946b5eb420a02c0cf725db0480c06ffaab5e6 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Thu, 18 Jul 2024 16:40:37 +0200 Subject: [PATCH 10/19] fix(multichain): do not subscribe to AccountsController:stateChange when fetching balances --- .../lib/accounts/BalancesController.test.ts | 1 - .../lib/accounts/BalancesController.ts | 19 +------------------ app/scripts/metamask-controller.js | 1 - 3 files changed, 1 insertion(+), 20 deletions(-) diff --git a/app/scripts/lib/accounts/BalancesController.test.ts b/app/scripts/lib/accounts/BalancesController.test.ts index 41d30ffa3bf5..a596c9aae43b 100644 --- a/app/scripts/lib/accounts/BalancesController.test.ts +++ b/app/scripts/lib/accounts/BalancesController.test.ts @@ -61,7 +61,6 @@ const setupController = ({ allowedEvents: [ 'AccountsController:accountAdded', 'AccountsController:accountRemoved', - 'AccountsController:stateChange', ], }); diff --git a/app/scripts/lib/accounts/BalancesController.ts b/app/scripts/lib/accounts/BalancesController.ts index f2d7eab41593..2834985eb0f4 100644 --- a/app/scripts/lib/accounts/BalancesController.ts +++ b/app/scripts/lib/accounts/BalancesController.ts @@ -21,9 +21,7 @@ import type { Draft } from 'immer'; import type { AccountsControllerAccountAddedEvent, AccountsControllerAccountRemovedEvent, - AccountsControllerChangeEvent, AccountsControllerListMultichainAccountsAction, - AccountsControllerState, } from '@metamask/accounts-controller'; import { isBtcMainnetAddress } from '../../../../shared/lib/multichain'; import { BalancesTracker } from './BalancesTracker'; @@ -97,8 +95,7 @@ export type AllowedActions = */ export type AllowedEvents = | AccountsControllerAccountAddedEvent - | AccountsControllerAccountRemovedEvent - | AccountsControllerChangeEvent; + | AccountsControllerAccountRemovedEvent; /** * Messenger type for the BalancesController. @@ -184,20 +181,6 @@ export class BalancesController extends BaseController< 'AccountsController:accountRemoved', (account) => this.#handleOnAccountRemoved(account), ); - this.messagingSystem.subscribe( - 'AccountsController:stateChange', - (_newState: AccountsControllerState) => { - // The tracker won't refresh the balance if it's not required, so there's - // very little overhead of using it here. - // - // However, updating the balances here allow us to fetch the balance of any new - // created account directly (we start tracking in `:accountAdded` event handler). - // - // In this case, the tracker will fetch the balance (for the first time) of those - // new accounts. - this.#tracker.updateBalances(); - }, - ); } /** diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js index 59e4ea8209bb..b73bc29d4561 100644 --- a/app/scripts/metamask-controller.js +++ b/app/scripts/metamask-controller.js @@ -879,7 +879,6 @@ export default class MetamaskController extends EventEmitter { allowedEvents: [ 'AccountsController:accountAdded', 'AccountsController:accountRemoved', - 'AccountsController:stateChange', ], allowedActions: [ 'AccountsController:listMultichainAccounts', From 502872eaaec00ced12d9d51d11d4172035545542 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Thu, 18 Jul 2024 17:25:57 +0200 Subject: [PATCH 11/19] feat(multichain): force fetch balance upon account creation (in CreateBtcAccount) --- .../lib/accounts/BalancesController.ts | 20 +++++++++---------- app/scripts/metamask-controller.js | 7 +++++++ .../create-btc-account/create-btc-account.tsx | 4 ++++ ui/store/actions.ts | 12 +++++++++++ 4 files changed, 33 insertions(+), 10 deletions(-) diff --git a/app/scripts/lib/accounts/BalancesController.ts b/app/scripts/lib/accounts/BalancesController.ts index 2834985eb0f4..87a6441570e3 100644 --- a/app/scripts/lib/accounts/BalancesController.ts +++ b/app/scripts/lib/accounts/BalancesController.ts @@ -162,11 +162,15 @@ export class BalancesController extends BaseController< 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}`); + } - // Just to double-check, we call updateBalance rather than #updateBalance to make sure - // we are using a non-EVM account - await this.updateBalance(account); + 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); @@ -255,14 +259,10 @@ export class BalancesController extends BaseController< * Updates the balances of one account. This method doesn't return * anything, but it updates the state of the controller. * - * @param account - The account. + * @param accountId - The account ID. */ - async updateBalance(account: InternalAccount) { - if (!this.#isNonEvmAccount(account)) { - throw new Error('Cannot update balance of EVM accounts from here'); - } - - this.#updateBalance(account); + async updateBalance(accountId: string) { + await this.#tracker.updateBalance(accountId); } /** diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js index b73bc29d4561..93beb9f7ac44 100644 --- a/app/scripts/metamask-controller.js +++ b/app/scripts/metamask-controller.js @@ -3750,6 +3750,13 @@ export default class MetamaskController extends EventEmitter { ), setName: this.nameController.setName.bind(this.nameController), + // MultichainBalancesController + multichainUpdateBalance: (accountId) => + this.multichainBalancesController.updateBalance(accountId), + + multichainUpdateBalances: () => + this.multichainBalancesController.updateBalances(), + // Transaction Decode decodeTransactionData: (request) => decodeTransactionData({ diff --git a/ui/components/multichain/create-btc-account/create-btc-account.tsx b/ui/components/multichain/create-btc-account/create-btc-account.tsx index 29c7b8f345b3..fb92ae973ae4 100644 --- a/ui/components/multichain/create-btc-account/create-btc-account.tsx +++ b/ui/components/multichain/create-btc-account/create-btc-account.tsx @@ -7,6 +7,7 @@ import { BitcoinWalletSnapSender } from '../../../../app/scripts/lib/snap-keyrin import { setAccountLabel, forceUpdateMetamaskState, + multichainUpdateBalance, } from '../../../store/actions'; type CreateBtcAccountOptions = { @@ -51,6 +52,9 @@ export const CreateBtcAccount = ({ dispatch(setAccountLabel(account.address, name)); } + // Force update the balances + await multichainUpdateBalance(account.id); + await onActionComplete(true); }; diff --git a/ui/store/actions.ts b/ui/store/actions.ts index 7493ee03ba3f..e62b00a5dced 100644 --- a/ui/store/actions.ts +++ b/ui/store/actions.ts @@ -5660,3 +5660,15 @@ export async function decodeTransactionData({ }, ]); } + +export async function multichainUpdateBalance( + accountId: string, +): Promise { + return await submitRequestToBackground('multichainUpdateBalance', [ + accountId, + ]); +} + +export async function multichainUpdateBalances(): Promise { + return await submitRequestToBackground('multichainUpdateBalances', []); +} From f973430e2891f841722e70b4fac7b5aaf6cccbb6 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Thu, 18 Jul 2024 17:46:11 +0200 Subject: [PATCH 12/19] feat(multichain): update fetch info when fetching single account balances --- app/scripts/lib/accounts/BalancesTracker.ts | 27 +++++++++++---------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/app/scripts/lib/accounts/BalancesTracker.ts b/app/scripts/lib/accounts/BalancesTracker.ts index eac60ab6854d..55ff0728690c 100644 --- a/app/scripts/lib/accounts/BalancesTracker.ts +++ b/app/scripts/lib/accounts/BalancesTracker.ts @@ -94,8 +94,18 @@ export class BalancesTracker { */ async updateBalance(accountId: string) { this.assertBeingTracked(accountId); - await this.#updateBalance(accountId); - this.#balances[accountId].lastUpdated = Date.now(); + + // We check if the balance is outdated (by comparing to the block time associated + // with this kind of account). + // + // This might not be super accurate, but we could probably compute this differently + // and try to sync with the "real block time"! + const info = this.#balances[accountId]; + const isOutdated = Date.now() - info.lastUpdated >= info.blockTime; + if (isOutdated) { + await this.#updateBalance(accountId); + this.#balances[accountId].lastUpdated = Date.now(); + } } /** @@ -103,17 +113,8 @@ export class BalancesTracker { * is considered outdated). */ async updateBalances() { - for (const [accountId, info] of Object.entries(this.#balances)) { - // We check if the balance is outdated (by comparing to the block time associated - // with this kind of account). - // - // This might not be super accurate, but we could probably compute this differently - // and try to sync with the "real block time"! - const isOutdated = Date.now() - info.lastUpdated >= info.blockTime; - - if (isOutdated) { - await this.updateBalance(accountId); - } + for (const accountId of Object.keys(this.#balances)) { + await this.updateBalance(accountId); } } } From a3bd765af8de05d4e6adb6554c6632c03223347c Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Thu, 18 Jul 2024 17:46:32 +0200 Subject: [PATCH 13/19] fix(multichain): fix BalancesController initial state --- app/scripts/metamask-controller.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js index 93beb9f7ac44..3d739072d7cb 100644 --- a/app/scripts/metamask-controller.js +++ b/app/scripts/metamask-controller.js @@ -888,7 +888,7 @@ export default class MetamaskController extends EventEmitter { this.multichainBalancesController = new MultichainBalancesController({ messenger: multichainBalancesControllerMessenger, - state: {}, + state: initState.MultichainBalancesController, }); const multichainRatesControllerMessenger = From 1e805b23266b3f6ee0e283f455b260e700e20dab Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Thu, 18 Jul 2024 18:47:14 +0200 Subject: [PATCH 14/19] fix(multichain): fix btc balances tests --- app/scripts/metamask-controller.test.js | 30 +++++++++++-------- .../create-btc-account.test.tsx | 11 +++++++ 2 files changed, 29 insertions(+), 12 deletions(-) diff --git a/app/scripts/metamask-controller.test.js b/app/scripts/metamask-controller.test.js index 45e2c3bb1b19..9acfa0525726 100644 --- a/app/scripts/metamask-controller.test.js +++ b/app/scripts/metamask-controller.test.js @@ -43,6 +43,7 @@ import { BalancesController as MultichainBalancesController, BTC_AVG_BLOCK_TIME, } from './lib/accounts/BalancesController'; +import { BalancesTracker as MultichainBalancesTracker } from './lib/accounts/BalancesTracker'; import { deferredPromise } from './lib/util'; import MetaMaskController from './metamask-controller'; @@ -2254,6 +2255,7 @@ describe('MetaMaskController', () => { }, }; let localMetamaskController; + let spyBalancesTrackerUpdateBalance; beforeEach(() => { jest.useFakeTimers(); @@ -2261,6 +2263,9 @@ describe('MetaMaskController', () => { jest .spyOn(MultichainBalancesController.prototype, 'updateBalance') .mockResolvedValue(); + spyBalancesTrackerUpdateBalance = jest + .spyOn(MultichainBalancesTracker.prototype, 'updateBalance') + .mockResolvedValue(); localMetamaskController = new MetaMaskController({ showUserConfirmation: noop, encryptor: mockEncryptor, @@ -2305,23 +2310,24 @@ describe('MetaMaskController', () => { expect( localMetamaskController.multichainBalancesController.updateBalances, ).toHaveBeenCalledTimes(1); - expect( - localMetamaskController.multichainBalancesController.updateBalance, - ).toHaveBeenCalledTimes(1); - expect( - localMetamaskController.multichainBalancesController.updateBalance, - ).toHaveBeenCalledWith(mockNonEvmAccount); + expect(spyBalancesTrackerUpdateBalance).toHaveBeenCalledTimes(1); + expect(spyBalancesTrackerUpdateBalance).toHaveBeenCalledWith( + mockNonEvmAccount.id, + ); // Wait for "block time", so balances will have to be refreshed jest.advanceTimersByTime(BTC_AVG_BLOCK_TIME); - // 2nd call because balances are considered "outdated": - expect( - localMetamaskController.multichainBalancesController.updateBalance, - ).toHaveBeenCalledTimes(2); // 1 (startup) + 1 (refresh) + // Check that we tried to fetch the balances more than once + // NOTE: For now, this method might be called a lot more than just twice, but this + // method has some internal logic to prevent fetching the balance too often if we + // consider the balance to be "up-to-date" expect( - localMetamaskController.multichainBalancesController.updateBalance, - ).toHaveBeenLastCalledWith(mockNonEvmAccount); + spyBalancesTrackerUpdateBalance.mock.calls.length, + ).toBeGreaterThan(1); + expect(spyBalancesTrackerUpdateBalance).toHaveBeenLastCalledWith( + mockNonEvmAccount.id, + ); }); }); }); diff --git a/ui/components/multichain/create-btc-account/create-btc-account.test.tsx b/ui/components/multichain/create-btc-account/create-btc-account.test.tsx index 23fbe9226707..6c64bb7d553b 100644 --- a/ui/components/multichain/create-btc-account/create-btc-account.test.tsx +++ b/ui/components/multichain/create-btc-account/create-btc-account.test.tsx @@ -33,12 +33,22 @@ const mockBtcAccount = { methods: [BtcMethod.SendMany], }; const mockBitcoinWalletSnapSend = jest.fn().mockReturnValue(mockBtcAccount); +const mockMultichainUpdateBalance = jest.fn().mockReturnValue({ + [mockBtcAccount.address]: { + [`${MultichainNetworks.BITCOIN_TESTNET}/slip44:0`]: { + amount: '0.00000000', + unit: 'BTC', + }, + }, +}); const mockSetAccountLabel = jest.fn().mockReturnValue({ type: 'TYPE' }); jest.mock('../../../store/actions', () => ({ forceUpdateMetamaskState: jest.fn(), setAccountLabel: (address: string, label: string) => mockSetAccountLabel(address, label), + multichainUpdateBalance: (accountId: string) => + mockMultichainUpdateBalance(accountId), })); jest.mock( @@ -85,6 +95,7 @@ describe('CreateBtcAccount', () => { newAccountName, ), ); + await waitFor(() => expect(mockMultichainUpdateBalance).toHaveBeenCalled()); await waitFor(() => expect(onActionComplete).toHaveBeenCalled()); }); From ecd1b6aa75ab074e2cf0f87965e9a2f9d5d18e72 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Fri, 19 Jul 2024 13:10:33 +0200 Subject: [PATCH 15/19] feat(multichain): add BALANCES_UPDATE_TIME (lower than average block time) --- app/scripts/lib/accounts/BalancesController.ts | 8 ++++++-- app/scripts/metamask-controller.test.js | 4 ++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/app/scripts/lib/accounts/BalancesController.ts b/app/scripts/lib/accounts/BalancesController.ts index 87a6441570e3..f6e50148b198 100644 --- a/app/scripts/lib/accounts/BalancesController.ts +++ b/app/scripts/lib/accounts/BalancesController.ts @@ -124,7 +124,11 @@ const balancesControllerMetadata = { const BTC_TESTNET_ASSETS = ['bip122:000000000933ea01ad0ee984209779ba/slip44:0']; const BTC_MAINNET_ASSETS = ['bip122:000000000019d6689c085ae165831e93/slip44:0']; -export const BTC_AVG_BLOCK_TIME = 10 * 60 * 1000; // 10 minutes in milliseconds +const BTC_AVG_BLOCK_TIME = 10 * 60 * 1000; // 10 minutes in milliseconds + +// NOTE: We set an interval of half the average block time to mitigate when our interval +// is de-synchronized with the actual block time. +export const BALANCES_UPDATE_TIME = BTC_AVG_BLOCK_TIME / 2; /** * The BalancesController is responsible for fetching and caching account @@ -173,7 +177,7 @@ export class BalancesController extends BaseController< // 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.#tracker.track(account.id, BALANCES_UPDATE_TIME); } } diff --git a/app/scripts/metamask-controller.test.js b/app/scripts/metamask-controller.test.js index 9acfa0525726..3aeac2866219 100644 --- a/app/scripts/metamask-controller.test.js +++ b/app/scripts/metamask-controller.test.js @@ -41,7 +41,7 @@ import { ETH_EOA_METHODS } from '../../shared/constants/eth-methods'; import { createMockInternalAccount } from '../../test/jest/mocks'; import { BalancesController as MultichainBalancesController, - BTC_AVG_BLOCK_TIME, + BALANCES_UPDATE_TIME as MULTICHAIN_BALANCES_UPDATE_TIME, } from './lib/accounts/BalancesController'; import { BalancesTracker as MultichainBalancesTracker } from './lib/accounts/BalancesTracker'; import { deferredPromise } from './lib/util'; @@ -2316,7 +2316,7 @@ describe('MetaMaskController', () => { ); // Wait for "block time", so balances will have to be refreshed - jest.advanceTimersByTime(BTC_AVG_BLOCK_TIME); + jest.advanceTimersByTime(MULTICHAIN_BALANCES_UPDATE_TIME); // Check that we tried to fetch the balances more than once // NOTE: For now, this method might be called a lot more than just twice, but this From 175abe5eaa63ec18c13a9ca72a0f99199aa1d955 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Fri, 19 Jul 2024 13:15:42 +0200 Subject: [PATCH 16/19] refactor(multichain): use Promise.allSettled when fecthing all balances --- app/scripts/lib/accounts/BalancesTracker.ts | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/app/scripts/lib/accounts/BalancesTracker.ts b/app/scripts/lib/accounts/BalancesTracker.ts index 55ff0728690c..48ecd6f84cca 100644 --- a/app/scripts/lib/accounts/BalancesTracker.ts +++ b/app/scripts/lib/accounts/BalancesTracker.ts @@ -113,8 +113,10 @@ export class BalancesTracker { * is considered outdated). */ async updateBalances() { - for (const accountId of Object.keys(this.#balances)) { - await this.updateBalance(accountId); - } + await Promise.allSettled( + Object.keys(this.#balances).map(async (accountId) => { + await this.updateBalance(accountId); + }), + ); } } From 3dba8f7a9f363458494ce7459d779c57a39e8a57 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Fri, 19 Jul 2024 13:40:27 +0200 Subject: [PATCH 17/19] refactor(multichain): use named function for BalancesTracker's callback --- .../lib/accounts/BalancesController.ts | 45 ++++++++++--------- 1 file changed, 25 insertions(+), 20 deletions(-) diff --git a/app/scripts/lib/accounts/BalancesController.ts b/app/scripts/lib/accounts/BalancesController.ts index f6e50148b198..5e76c4ad5a7f 100644 --- a/app/scripts/lib/accounts/BalancesController.ts +++ b/app/scripts/lib/accounts/BalancesController.ts @@ -158,21 +158,9 @@ 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); - }); + this.#tracker = new BalancesTracker( + async (accountId: string) => await this.#updateBalance(accountId), + ); // Register all non-EVM accounts into the tracker for (const account of this.#listAccounts()) { @@ -230,16 +218,33 @@ export class BalancesController extends BaseController< return accounts.filter((account) => account.type === BtcAccountType.P2wpkh); } + /** + * Get a non-EVM account from its ID. + * + * @param accountId - The account ID. + */ + #getAccount(accountId: string): InternalAccount { + const account: InternalAccount = this.#listMultichainAccounts().find( + (multichainAccount) => multichainAccount.id === accountId, + ); + + if (!account) { + throw new Error(`Unknown account: ${accountId}`); + } + if (!this.#isNonEvmAccount(account)) { + throw new Error(`Account is not a non-EVM account: ${accountId}`); + } + return account; + } + /** * 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. + * @param accountId - The account ID. */ - async #updateBalance(account: InternalAccount) { + async #updateBalance(accountId: string) { + const account = this.#getAccount(accountId); const partialState: BalancesControllerState = { balances: {} }; partialState.balances[account.id] = await this.#getBalances( From 5be1e464ca397426b05207827cc1f952949685aa Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Fri, 19 Jul 2024 13:40:41 +0200 Subject: [PATCH 18/19] test(multichain): remove use of it.only --- app/scripts/lib/accounts/BalancesController.test.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/app/scripts/lib/accounts/BalancesController.test.ts b/app/scripts/lib/accounts/BalancesController.test.ts index a596c9aae43b..036b29f1da76 100644 --- a/app/scripts/lib/accounts/BalancesController.test.ts +++ b/app/scripts/lib/accounts/BalancesController.test.ts @@ -126,7 +126,7 @@ describe('BalancesController', () => { }); }); - it.only('update balances when "AccountsController:accountAdded" is fired', async () => { + it('update balances when "AccountsController:accountAdded" is fired', async () => { const { controller, messenger, mockListMultichainAccounts } = setupController({ mocks: { @@ -146,7 +146,7 @@ describe('BalancesController', () => { }); }); - it.only('update balances when "AccountsController:accountRemoved" is fired', async () => { + it('update balances when "AccountsController:accountRemoved" is fired', async () => { const { controller, messenger, mockListMultichainAccounts } = setupController(); From 9637c8e654d5024faf987eb2e843c8b510b9bcd7 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Fri, 19 Jul 2024 13:50:13 +0200 Subject: [PATCH 19/19] refactor(btc): prevent error from bubbling up when fetching BTC balance after account creation --- .../create-btc-account/create-btc-account.tsx | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/ui/components/multichain/create-btc-account/create-btc-account.tsx b/ui/components/multichain/create-btc-account/create-btc-account.tsx index fb92ae973ae4..98b4619daeca 100644 --- a/ui/components/multichain/create-btc-account/create-btc-account.tsx +++ b/ui/components/multichain/create-btc-account/create-btc-account.tsx @@ -52,10 +52,21 @@ export const CreateBtcAccount = ({ dispatch(setAccountLabel(account.address, name)); } - // Force update the balances - await multichainUpdateBalance(account.id); - + // This will close up the name dialog await onActionComplete(true); + + // Force update the balances + try { + await multichainUpdateBalance(account.id); + } catch (error) { + // To avoid breaking the flow entirely, we do catch any error that might happens while fetching + // the balance. + // Worst case scenario, the balance will be updated during a future tick of the + // MultichainBalancesTracker! + console.warn( + `Unable to fetch Bitcoin balance: ${(error as Error).message}`, + ); + } }; const getNextAvailableAccountName = async (_accounts: InternalAccount[]) => {