diff --git a/packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.ts b/packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.ts index fc52d2c5d7..9442607e56 100644 --- a/packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.ts +++ b/packages/assets-controllers/src/MultichainBalancesController/MultichainBalancesController.ts @@ -19,12 +19,8 @@ import { HandlerType } from '@metamask/snaps-utils'; import type { Json, JsonRpcRequest } from '@metamask/utils'; import type { Draft } from 'immer'; -import { - BalancesTracker, - BALANCE_UPDATE_INTERVALS, - NETWORK_ASSETS_MAP, -} from '.'; -import { getScopeForAccount } from './utils'; +import { BalancesTracker, NETWORK_ASSETS_MAP } from '.'; +import { getScopeForAccount, getBlockTimeForAccount } from './utils'; const controllerName = 'MultichainBalancesController'; @@ -168,7 +164,7 @@ export class MultichainBalancesController 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, this.#getBlockTimeFor(account)); + this.#tracker.track(account.id, getBlockTimeForAccount(account.type)); } } @@ -197,20 +193,23 @@ export class MultichainBalancesController extends BaseController< } /** - * Gets the block time for a given account. + * Updates the balances of one account. This method doesn't return + * anything, but it updates the state of the controller. * - * @param account - The account to get the block time for. - * @returns The block time for the account. + * @param accountId - The account ID. */ - #getBlockTimeFor(account: InternalAccount): number { - if (account.type in BALANCE_UPDATE_INTERVALS) { - return BALANCE_UPDATE_INTERVALS[ - account.type as keyof typeof BALANCE_UPDATE_INTERVALS - ]; - } - throw new Error( - `Unsupported account type for balance tracking: ${account.type}`, - ); + async updateBalance(accountId: string): Promise { + // NOTE: No need to track the account here, since we start tracking those when + // the "AccountsController:accountAdded" is fired. + 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(): Promise { + await this.#tracker.updateBalances(); } /** @@ -242,17 +241,14 @@ export class MultichainBalancesController extends BaseController< * @returns The non-EVM account. */ #getAccount(accountId: string): InternalAccount { - const account: InternalAccount | undefined = - this.#listMultichainAccounts().find( - (multichainAccount) => multichainAccount.id === accountId, - ); + const account: InternalAccount | undefined = this.#listAccounts().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; } @@ -282,26 +278,6 @@ export class MultichainBalancesController extends BaseController< } } - /** - * 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 { - // NOTE: No need to track the account here, since we start tracking those when - // the "AccountsController:accountAdded" is fired. - 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(): Promise { - await this.#tracker.updateBalances(); - } - /** * Checks for non-EVM accounts. * @@ -327,7 +303,7 @@ export class MultichainBalancesController extends BaseController< return; } - this.#tracker.track(account.id, this.#getBlockTimeFor(account)); + this.#tracker.track(account.id, getBlockTimeForAccount(account.type)); // NOTE: Unfortunately, we cannot update the balance right away here, because // messenger's events are running synchronously and fetching the balance is // asynchronous. diff --git a/packages/assets-controllers/src/MultichainBalancesController/utils.test.ts b/packages/assets-controllers/src/MultichainBalancesController/utils.test.ts index 3e65f473a0..c566ad8374 100644 --- a/packages/assets-controllers/src/MultichainBalancesController/utils.test.ts +++ b/packages/assets-controllers/src/MultichainBalancesController/utils.test.ts @@ -8,11 +8,12 @@ import { KeyringTypes } from '@metamask/keyring-controller'; import { validate, Network } from 'bitcoin-address-validation'; import { v4 as uuidv4 } from 'uuid'; -import { MultichainNetworks } from './constants'; +import { MultichainNetworks, BALANCE_UPDATE_INTERVALS } from '.'; import { getScopeForBtcAddress, getScopeForSolAddress, getScopeForAccount, + getBlockTimeForAccount, } from './utils'; const mockBtcAccount = { @@ -169,3 +170,24 @@ describe('getScopeForAddress', () => { ); }); }); + +describe('getBlockTimeForAccount', () => { + it('returns the block time for a supported Bitcoin account', () => { + const blockTime = getBlockTimeForAccount(BtcAccountType.P2wpkh); + expect(blockTime).toBe(BALANCE_UPDATE_INTERVALS[BtcAccountType.P2wpkh]); + }); + + it('returns the block time for a supported Solana account', () => { + const blockTime = getBlockTimeForAccount(SolAccountType.DataAccount); + expect(blockTime).toBe( + BALANCE_UPDATE_INTERVALS[SolAccountType.DataAccount], + ); + }); + + it('throws an error for an unsupported account type', () => { + const unsupportedAccountType = 'unsupported-type'; + expect(() => getBlockTimeForAccount(unsupportedAccountType)).toThrow( + `Unsupported account type for balance tracking: ${unsupportedAccountType}`, + ); + }); +}); diff --git a/packages/assets-controllers/src/MultichainBalancesController/utils.ts b/packages/assets-controllers/src/MultichainBalancesController/utils.ts index 72728b2299..205cca8fc3 100644 --- a/packages/assets-controllers/src/MultichainBalancesController/utils.ts +++ b/packages/assets-controllers/src/MultichainBalancesController/utils.ts @@ -2,7 +2,7 @@ import { BtcAccountType, SolAccountType } from '@metamask/keyring-api'; import type { InternalAccount } from '@metamask/keyring-internal-api'; import { validate, Network } from 'bitcoin-address-validation'; -import { MultichainNetworks } from './constants'; +import { MultichainNetworks, BALANCE_UPDATE_INTERVALS } from './constants'; /** * Gets the scope for a specific and supported Bitcoin account. @@ -58,3 +58,20 @@ export const getScopeForAccount = (account: InternalAccount): string => { throw new Error(`Unsupported non-EVM account type: ${account.type}`); } }; + +/** + * Gets the block time for a given account. + * + * @param accountType - The account type to get the block time for. + * @returns The block time for the account. + */ +export const getBlockTimeForAccount = (accountType: string): number => { + if (accountType in BALANCE_UPDATE_INTERVALS) { + return BALANCE_UPDATE_INTERVALS[ + accountType as keyof typeof BALANCE_UPDATE_INTERVALS + ]; + } + throw new Error( + `Unsupported account type for balance tracking: ${accountType}`, + ); +};