From c3fcd19a528fd15d4f3db01e1d0c5a7923de65b4 Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Thu, 19 Mar 2026 12:00:01 +0000 Subject: [PATCH 1/5] feat: add batching around update balance requests We have many things (events, interactions) invoking balance update requests. This adds some batching and aggregate logic to minimise the number of requests to process. --- .../src/TokenBalancesController.test.ts | 2191 ++++++++++------- .../src/TokenBalancesController.ts | 71 +- .../src/__fixtures__/test-utils.ts | 35 + .../src/utils/create-batch-handler.test.ts | 121 + .../src/utils/create-batch-handler.ts | 39 + 5 files changed, 1554 insertions(+), 903 deletions(-) create mode 100644 packages/assets-controllers/src/__fixtures__/test-utils.ts create mode 100644 packages/assets-controllers/src/utils/create-batch-handler.test.ts create mode 100644 packages/assets-controllers/src/utils/create-batch-handler.ts diff --git a/packages/assets-controllers/src/TokenBalancesController.test.ts b/packages/assets-controllers/src/TokenBalancesController.test.ts index 950258c0858..09f765882d5 100644 --- a/packages/assets-controllers/src/TokenBalancesController.test.ts +++ b/packages/assets-controllers/src/TokenBalancesController.test.ts @@ -17,6 +17,7 @@ import BN from 'bn.js'; import type nock from 'nock'; import { mockAPI_accountsAPI_MultichainAccountBalances as mockAPIAccountsAPIMultichainAccountBalancesCamelCase } from './__fixtures__/account-api-v4-mocks'; +import { waitFor } from './__fixtures__/test-utils'; import { AccountsApiBalanceFetcher } from './multi-chain-accounts-service/api-balance-fetcher'; import * as multicall from './multicall'; import { RpcBalanceFetcher } from './rpc-service/rpc-balance-fetcher'; @@ -26,10 +27,13 @@ import type { ChecksumAddress, TokenBalancesControllerState, TokenBalances, + UpdateBalancesOptions, } from './TokenBalancesController'; import { TokenBalancesController, + UPDATE_BALANCES_BATCH_MS, caipChainIdToHex, + mergeUpdateBalancesOptions, parseAssetType, } from './TokenBalancesController'; import type { TokensControllerState } from './TokensController'; @@ -324,10 +328,217 @@ describe('Utility Functions', () => { }); }); +describe('mergeUpdateBalancesOptions', () => { + const arrangeChainIdInput = ( + chainIds: ChainIdHex[] | undefined, + ): UpdateBalancesOptions => ({ + chainIds, + tokenAddresses: undefined, + queryAllAccounts: undefined, + }); + + const arrangeTokenAddressesInput = ( + tokenAddresses: string[] | undefined, + ): UpdateBalancesOptions => ({ + chainIds: undefined, + tokenAddresses, + queryAllAccounts: undefined, + }); + + const arrangeQueryAllAccountsInput = ( + queryAllAccounts: boolean | undefined, + ): UpdateBalancesOptions => ({ + chainIds: undefined, + tokenAddresses: undefined, + queryAllAccounts, + }); + + const chainIdTestCases: { + testName: string; + balanceInputs: [UpdateBalancesOptions, UpdateBalancesOptions]; + expectedOutput: Partial; + }[] = [ + { + testName: 'merges chainIds as union when both specify chainIds', + balanceInputs: [ + arrangeChainIdInput(['0x1']), + arrangeChainIdInput(['0x89', '0x1']), + ], + expectedOutput: { chainIds: ['0x1', '0x89'] }, + }, + { + testName: 'returns undefined chainIds when first option has no chainIds', + balanceInputs: [ + arrangeChainIdInput(undefined), + arrangeChainIdInput(['0x89']), + ], + expectedOutput: { chainIds: undefined }, + }, + { + testName: 'returns undefined chainIds when second option has no chainIds', + balanceInputs: [ + arrangeChainIdInput(['0x1']), + arrangeChainIdInput(undefined), + ], + expectedOutput: { chainIds: undefined }, + }, + { + testName: 'returns undefined chainIds when neither option has chainIds', + balanceInputs: [ + arrangeChainIdInput(undefined), + arrangeChainIdInput(undefined), + ], + expectedOutput: { chainIds: undefined }, + }, + ]; + + const tokenAddressesTestCases: { + testName: string; + balanceInputs: [UpdateBalancesOptions, UpdateBalancesOptions]; + expectedOutput: Partial; + }[] = [ + { + testName: + 'merges tokenAddresses as union when both specify tokenAddresses', + balanceInputs: [ + arrangeTokenAddressesInput(['0xabc', '0xdef']), + arrangeTokenAddressesInput(['0xdef', '0x123']), + ], + expectedOutput: { + tokenAddresses: ['0xabc', '0xdef', '0x123'], + }, + }, + { + testName: + 'returns undefined tokenAddresses when first option has no tokenAddresses', + balanceInputs: [ + arrangeTokenAddressesInput(undefined), + arrangeTokenAddressesInput(['0xdef']), + ], + expectedOutput: { tokenAddresses: undefined }, + }, + { + testName: + 'returns undefined tokenAddresses when second option has no tokenAddresses', + balanceInputs: [ + arrangeTokenAddressesInput(['0xabc']), + arrangeTokenAddressesInput(undefined), + ], + expectedOutput: { tokenAddresses: undefined }, + }, + ]; + + const queryAllAccountsTestCases: { + testName: string; + balanceInputs: [UpdateBalancesOptions, UpdateBalancesOptions]; + expectedOutput: Partial; + }[] = [ + { + testName: 'returns true when first is true (true, false)', + balanceInputs: [ + arrangeQueryAllAccountsInput(true), + arrangeQueryAllAccountsInput(false), + ], + expectedOutput: { queryAllAccounts: true }, + }, + { + testName: 'returns true when second is true (false, true)', + balanceInputs: [ + arrangeQueryAllAccountsInput(false), + arrangeQueryAllAccountsInput(true), + ], + expectedOutput: { queryAllAccounts: true }, + }, + { + testName: 'returns true when both have queryAllAccounts true', + balanceInputs: [ + arrangeQueryAllAccountsInput(true), + arrangeQueryAllAccountsInput(true), + ], + expectedOutput: { queryAllAccounts: true }, + }, + { + testName: 'returns true when second is true and first is undefined', + balanceInputs: [ + arrangeQueryAllAccountsInput(undefined), + arrangeQueryAllAccountsInput(true), + ], + expectedOutput: { queryAllAccounts: true }, + }, + { + testName: 'returns false when both have queryAllAccounts false', + balanceInputs: [ + arrangeQueryAllAccountsInput(false), + arrangeQueryAllAccountsInput(false), + ], + expectedOutput: { queryAllAccounts: false }, + }, + { + testName: 'returns false when both are undefined', + balanceInputs: [ + arrangeQueryAllAccountsInput(undefined), + arrangeQueryAllAccountsInput(undefined), + ], + expectedOutput: { queryAllAccounts: false }, + }, + { + testName: 'returns false when second is false and first is undefined', + balanceInputs: [ + arrangeQueryAllAccountsInput(undefined), + arrangeQueryAllAccountsInput(false), + ], + expectedOutput: { queryAllAccounts: false }, + }, + ]; + + it.each(chainIdTestCases)( + '$testName', + ({ balanceInputs, expectedOutput }) => { + expect( + mergeUpdateBalancesOptions(balanceInputs[0], balanceInputs[1]), + ).toStrictEqual(expect.objectContaining(expectedOutput)); + }, + ); + + it.each(tokenAddressesTestCases)( + '$testName', + ({ balanceInputs, expectedOutput }) => { + expect( + mergeUpdateBalancesOptions(balanceInputs[0], balanceInputs[1]), + ).toStrictEqual(expect.objectContaining(expectedOutput)); + }, + ); + + it.each(queryAllAccountsTestCases)( + '$testName', + ({ balanceInputs, expectedOutput }) => { + expect( + mergeUpdateBalancesOptions(balanceInputs[0], balanceInputs[1]), + ).toStrictEqual(expect.objectContaining(expectedOutput)); + }, + ); + + it('merges all fields together when both options are fully specified', () => { + const a: UpdateBalancesOptions = { + chainIds: ['0x1'], + tokenAddresses: ['0xaaa'], + queryAllAccounts: false, + }; + const b: UpdateBalancesOptions = { + chainIds: ['0x89', '0x1'], + tokenAddresses: ['0xbbb', '0xaaa'], + queryAllAccounts: true, + }; + expect(mergeUpdateBalancesOptions(a, b)).toStrictEqual({ + chainIds: ['0x1', '0x89'], + tokenAddresses: ['0xaaa', '0xbbb'], + queryAllAccounts: true, + }); + }); +}); + describe('TokenBalancesController', () => { beforeEach(() => { - jest.useFakeTimers({ doNotFake: ['nextTick', 'queueMicrotask'] }); - // Mock safelyExecuteWithTimeout to execute the operation normally by default mockedSafelyExecuteWithTimeout.mockImplementation( async (operation: () => Promise) => { @@ -525,22 +736,27 @@ describe('TokenBalancesController', () => { }); it('should poll and update balances in the right interval', async () => { - const pollSpy = jest.spyOn( - TokenBalancesController.prototype, - '_executePoll', - ); + jest.useFakeTimers({ doNotFake: ['nextTick', 'queueMicrotask'] }); + try { + const pollSpy = jest.spyOn( + TokenBalancesController.prototype, + '_executePoll', + ); - const interval = 10; - const { controller } = setupController({ config: { interval } }); + const interval = 10; + const { controller } = setupController({ config: { interval } }); - controller.startPolling({ chainIds: ['0x1'] }); + controller.startPolling({ chainIds: ['0x1'] }); - await jestAdvanceTime({ duration: 1 }); - expect(pollSpy).toHaveBeenCalled(); - expect(pollSpy).not.toHaveBeenCalledTimes(2); + await jestAdvanceTime({ duration: 1 }); + expect(pollSpy).toHaveBeenCalled(); + expect(pollSpy).not.toHaveBeenCalledTimes(2); - await jestAdvanceTime({ duration: interval * 1.5 }); - expect(pollSpy).toHaveBeenCalledTimes(2); + await jestAdvanceTime({ duration: interval * 1.5 }); + expect(pollSpy).toHaveBeenCalledTimes(2); + } finally { + jest.useRealTimers(); + } }); it('should update balances on poll', async () => { @@ -578,14 +794,16 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); }); @@ -624,14 +842,16 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); } }); @@ -649,7 +869,10 @@ describe('TokenBalancesController', () => { chainIds: [chainId], queryAllAccounts: true, }); - expect(controller.state.tokenBalances).toStrictEqual({}); + + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({}); + }); const balance = 123456; jest @@ -680,16 +903,16 @@ describe('TokenBalancesController', () => { [], ); - await jestAdvanceTime({ duration: 1 }); - - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); }); @@ -735,15 +958,17 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Verify initial balance is set - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + // Verify initial balance is set + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); // Publish an update with no tokens @@ -757,14 +982,14 @@ describe('TokenBalancesController', () => { [], ); - await jestAdvanceTime({ duration: 1 }); - - // Verify balance was removed - expect(updateSpy).toHaveBeenCalledTimes(2); - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: {}, // Empty balances object - }, + await waitFor(() => { + // Verify balance was removed + expect(updateSpy).toHaveBeenCalledTimes(2); + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: {}, // Empty balances object + }, + }); }); }); it('skips removing balances when incoming chainIds are not in the current chainIds list for tokenBalances', async () => { @@ -805,15 +1030,17 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Verify initial balance is set - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + // Verify initial balance is set + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); // Publish an update with no tokens @@ -827,17 +1054,17 @@ describe('TokenBalancesController', () => { [], ); - await jestAdvanceTime({ duration: 1 }); - - expect(updateSpy).toHaveBeenCalledTimes(2); - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + expect(updateSpy).toHaveBeenCalledTimes(2); + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); }); @@ -879,15 +1106,17 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Verify initial balance is set - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + // Verify initial balance is set + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); // Publish an update with no tokens @@ -907,18 +1136,18 @@ describe('TokenBalancesController', () => { [], ); - await jestAdvanceTime({ duration: 1 }); - - // Verify initial balances are still there - expect(updateSpy).toHaveBeenCalledTimes(1); // should be called only once when we first updated the balances and not twice - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + // Verify initial balances are still there + expect(updateSpy).toHaveBeenCalledTimes(1); // should be called only once when we first updated the balances and not twice + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); }); @@ -979,21 +1208,23 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - expect(controller.state.tokenBalances).toStrictEqual({ - [account1]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance1), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({ + [account1]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance1), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, - [account2]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance2), - [STAKING_CONTRACT_ADDRESS]: '0x0', + [account2]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance2), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); }); @@ -1047,43 +1278,59 @@ describe('TokenBalancesController', () => { await controller.updateBalances({ chainIds: [chainId] }); - // Verify tracked token balance was updated - expect( - controller.state.tokenBalances[accountAddress]?.[chainId]?.[trackedToken], - ).toBe(toHex(trackedBalance)); + await waitFor(() => { + // Verify tracked token balance was updated + expect( + controller.state.tokenBalances[accountAddress]?.[chainId]?.[ + trackedToken + ], + ).toBe(toHex(trackedBalance)); - // Verify ignored token balance was updated (ignored tokens should still be tracked) - expect( - controller.state.tokenBalances[accountAddress]?.[chainId]?.[ignoredToken], - ).toBe(toHex(ignoredBalance)); + // Verify ignored token balance was updated (ignored tokens should still be tracked) + expect( + controller.state.tokenBalances[accountAddress]?.[chainId]?.[ + ignoredToken + ], + ).toBe(toHex(ignoredBalance)); - // Verify untracked token balance was NOT updated - expect( - controller.state.tokenBalances[accountAddress]?.[chainId]?.[ - untrackedToken - ], - ).toBeUndefined(); + // Verify untracked token balance was NOT updated + expect( + controller.state.tokenBalances[accountAddress]?.[chainId]?.[ + untrackedToken + ], + ).toBeUndefined(); - // Verify native token is always updated regardless of tracking - expect( - controller.state.tokenBalances[accountAddress]?.[chainId]?.[ - NATIVE_TOKEN_ADDRESS - ], - ).toBe('0x0'); + // Verify native token is always updated regardless of tracking + expect( + controller.state.tokenBalances[accountAddress]?.[chainId]?.[ + NATIVE_TOKEN_ADDRESS + ], + ).toBe('0x0'); + }); }); it('should always update native token balances regardless of tracking status', async () => { const chainId = '0x1'; const accountAddress = '0x0000000000000000000000000000000000000000'; + const tokenAddress = '0x0000000000000000000000000000000000000001'; - // Setup with no tracked tokens + // One tracked token so the controller fetches; we assert native is always updated const tokens = { allDetectedTokens: {}, - allTokens: {}, + allTokens: { + [chainId]: { + [accountAddress]: [ + { address: tokenAddress, symbol: 'T', decimals: 18 }, + ], + }, + }, allIgnoredTokens: {}, }; - const { controller } = setupController({ tokens }); + const { controller } = setupController({ + tokens, + listAccounts: [createMockInternalAccount({ address: accountAddress })], + }); const nativeBalance = new BN('1000000000000000000'); // 1 ETH @@ -1094,6 +1341,9 @@ describe('TokenBalancesController', () => { [NATIVE_TOKEN_ADDRESS]: { [accountAddress]: nativeBalance, }, + [tokenAddress]: { + [accountAddress]: new BN(0), + }, }, stakedBalances: { [accountAddress]: new BN(0), @@ -1102,12 +1352,14 @@ describe('TokenBalancesController', () => { await controller.updateBalances({ chainIds: [chainId] }); - // Verify native token balance was updated even though no tokens are tracked - expect( - controller.state.tokenBalances[accountAddress]?.[chainId]?.[ - NATIVE_TOKEN_ADDRESS - ], - ).toBe(toHex(nativeBalance)); + await waitFor(() => { + // Verify native token balance was updated + expect( + controller.state.tokenBalances[accountAddress]?.[chainId]?.[ + NATIVE_TOKEN_ADDRESS + ], + ).toBe(toHex(nativeBalance)); + }); }); it('should filter untracked tokens from balance updates', async () => { @@ -1150,17 +1402,22 @@ describe('TokenBalancesController', () => { await controller.updateBalances({ chainIds: [chainId] }); - // Verify tracked token balance was updated - expect( - controller.state.tokenBalances[accountAddress]?.[chainId]?.[trackedToken], - ).toBe(toHex(trackedBalance)); + await waitFor(() => { + // Verify tracked token balance was updated - // Verify untracked token balance was NOT updated - expect( - controller.state.tokenBalances[accountAddress]?.[chainId]?.[ - untrackedToken - ], - ).toBeUndefined(); + expect( + controller.state.tokenBalances[accountAddress]?.[chainId]?.[ + trackedToken + ], + ).toBe(toHex(trackedBalance)); + + // Verify untracked token balance was NOT updated + expect( + controller.state.tokenBalances[accountAddress]?.[chainId]?.[ + untrackedToken + ], + ).toBeUndefined(); + }); }); it('does not update balances when multi-account balances is enabled and all returned values did not change', async () => { @@ -1206,21 +1463,23 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - expect(controller.state.tokenBalances).toStrictEqual({ - [account1]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance1), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({ + [account1]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance1), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, - [account2]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance2), - [STAKING_CONTRACT_ADDRESS]: '0x0', + [account2]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance2), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); await controller._executePoll({ @@ -1228,8 +1487,10 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Should only update once since the values haven't changed - expect(updateSpy).toHaveBeenCalledTimes(1); + await waitFor(() => { + // Should only update once since the values haven't changed + expect(updateSpy).toHaveBeenCalledTimes(1); + }); }); it('does not update balances when multi-account balances is enabled and multi-account contract failed', async () => { @@ -1268,14 +1529,16 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - expect(controller.state.tokenBalances).toStrictEqual({ - [account1]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: '0x0', - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({ + [account1]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: '0x0', + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); await controller._executePoll({ @@ -1283,7 +1546,9 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - expect(updateSpy).toHaveBeenCalledTimes(1); // Called once because native/staking balances are added + await waitFor(() => { + expect(updateSpy).toHaveBeenCalledTimes(1); // Called once because native/staking balances are added + }); }); it('updates balances when multi-account balances is enabled and some returned values changed', async () => { @@ -1324,27 +1589,29 @@ describe('TokenBalancesController', () => { }, }, }); - - await controller._executePoll({ - chainIds: [chainId], - queryAllAccounts: true, - }); - - expect(controller.state.tokenBalances).toStrictEqual({ - [account1]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance1), - [STAKING_CONTRACT_ADDRESS]: '0x0', - }, - }, - [account2]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance2), - [STAKING_CONTRACT_ADDRESS]: '0x0', - }, - }, + + await controller._executePoll({ + chainIds: [chainId], + queryAllAccounts: true, + }); + + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({ + [account1]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance1), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, + }, + [account2]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance2), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, + }, + }); }); jest @@ -1364,21 +1631,23 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - expect(controller.state.tokenBalances).toStrictEqual({ - [account1]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance1), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({ + [account1]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance1), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, - [account2]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance3), - [STAKING_CONTRACT_ADDRESS]: '0x0', + [account2]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance3), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); expect(updateSpy).toHaveBeenCalledTimes(2); @@ -1427,15 +1696,17 @@ describe('TokenBalancesController', () => { queryAllAccounts: false, }); - // Should only contain balance for selected account - expect(controller.state.tokenBalances).toStrictEqual({ - [selectedAccount]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + // Should only contain balance for selected account + expect(controller.state.tokenBalances).toStrictEqual({ + [selectedAccount]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); }); @@ -1582,35 +1853,37 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(balance), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(balance), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, - [accountAddress2]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress2]: toHex(balance2), - [STAKING_CONTRACT_ADDRESS]: '0x0', + [accountAddress2]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress2]: toHex(balance2), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); messenger.publish('KeyringController:accountRemoved', account.address); - await jestAdvanceTime({ duration: 1 }); - - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress2]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress2]: toHex(balance2), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress2]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress2]: toHex(balance2), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); }); }); @@ -1735,8 +2008,10 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Verify the new multicall function was called - expect(mockGetTokenBalances).toHaveBeenCalled(); + await waitFor(() => { + // Verify the new multicall function was called + expect(mockGetTokenBalances).toHaveBeenCalled(); + }); }); it('should use queryAllAccounts when provided', async () => { @@ -1779,13 +2054,15 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Verify RPC fetcher was called with queryAllAccounts: true - expect(mockRpcFetch).toHaveBeenCalledWith( - expect.objectContaining({ - chainIds: ['0x1'], - queryAllAccounts: true, - }), - ); + await waitFor(() => { + // Verify RPC fetcher was called with queryAllAccounts: true + expect(mockRpcFetch).toHaveBeenCalledWith( + expect.objectContaining({ + chainIds: ['0x1'], + queryAllAccounts: true, + }), + ); + }); mockRpcFetch.mockRestore(); }); @@ -1827,11 +2104,15 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Verify the controller is properly configured - expect(controller).toBeDefined(); + await waitFor(() => { + // Verify the controller is properly configured + expect(controller).toBeDefined(); - // Verify multicall was attempted - expect(multicall.getTokenBalancesForMultipleAddresses).toHaveBeenCalled(); + // Verify multicall was attempted + expect( + multicall.getTokenBalancesForMultipleAddresses, + ).toHaveBeenCalled(); + }); }); it('should handle different constructor options', () => { @@ -1895,20 +2176,22 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Verify that staked balances are included in the state change event (even if zero) - expect(publishSpy).toHaveBeenCalledWith( - 'TokenBalancesController:stateChange', - expect.objectContaining({ - tokenBalances: { - [accountAddress]: { - [chainId]: expect.objectContaining({ - [STAKING_CONTRACT_ADDRESS]: '0x0', // Zero staked balance should be included - }), + await waitFor(() => { + // Verify that staked balances are included in the state change event (even if zero) + expect(publishSpy).toHaveBeenCalledWith( + 'TokenBalancesController:stateChange', + expect.objectContaining({ + tokenBalances: { + [accountAddress]: { + [chainId]: expect.objectContaining({ + [STAKING_CONTRACT_ADDRESS]: '0x0', // Zero staked balance should be included + }), + }, }, - }, - }), - expect.any(Array), - ); + }), + expect.any(Array), + ); + }); }); }); @@ -1955,14 +2238,16 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Only successful token should be in state - expect( - controller.state.tokenBalances[accountAddress][chainId], - ).toStrictEqual({ - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress1]: toHex(100), - [tokenAddress2]: '0x0', - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + // Only successful token should be in state + expect( + controller.state.tokenBalances[accountAddress][chainId], + ).toStrictEqual({ + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress1]: toHex(100), + [tokenAddress2]: '0x0', + [STAKING_CONTRACT_ADDRESS]: '0x0', + }); }); }); }); @@ -2003,14 +2288,16 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Verify both tokens are in state - expect( - controller.state.tokenBalances[accountAddress][chainId], - ).toStrictEqual({ - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress1]: toHex(100), - [tokenAddress2]: toHex(200), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + // Verify both tokens are in state + expect( + controller.state.tokenBalances[accountAddress][chainId], + ).toStrictEqual({ + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress1]: toHex(100), + [tokenAddress2]: toHex(200), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }); }); // For this test, we just verify the basic functionality without testing @@ -2047,8 +2334,10 @@ describe('TokenBalancesController', () => { // This should not throw and should return early await controller.updateBalances({ queryAllAccounts: true }); - // Verify no balances were fetched - expect(controller.state.tokenBalances).toStrictEqual({}); + await waitFor(() => { + // Verify no balances were fetched + expect(controller.state.tokenBalances).toStrictEqual({}); + }); }); it('handles case when no balances are aggregated', async () => { @@ -2071,8 +2360,10 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Verify no state update occurred - expect(controller.state.tokenBalances).toStrictEqual({}); + await waitFor(() => { + // Verify no state update occurred + expect(controller.state.tokenBalances).toStrictEqual({}); + }); }); it('handles case when no network configuration is found', async () => { @@ -2089,8 +2380,10 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Verify no balances were fetched - expect(controller.state.tokenBalances).toStrictEqual({}); + await waitFor(() => { + // Verify no balances were fetched + expect(controller.state.tokenBalances).toStrictEqual({}); + }); }); it('update native balance when fetch is successful', async () => { @@ -2136,14 +2429,16 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Verify no balances were fetched - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [tokenAddress]: toHex(100), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + // Verify no balances were fetched + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [tokenAddress]: toHex(100), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); }); @@ -2200,20 +2495,22 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Verify that: - // - tokenAddress1 has its actual fetched balance - // - tokenAddress2, tokenAddress3, and detectedTokenAddress have balance 0 - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress1]: toHex(123456), // Actual fetched balance - [tokenAddress2]: '0x0', // Zero balance for missing token - [tokenAddress3]: '0x0', // Zero balance for missing token - [detectedTokenAddress]: '0x0', // Zero balance for missing detected token - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + // Verify that: + // - tokenAddress1 has its actual fetched balance + // - tokenAddress2, tokenAddress3, and detectedTokenAddress have balance 0 + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress1]: toHex(123456), // Actual fetched balance + [tokenAddress2]: '0x0', // Zero balance for missing token + [tokenAddress3]: '0x0', // Zero balance for missing token + [detectedTokenAddress]: '0x0', // Zero balance for missing detected token + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); }); @@ -2256,16 +2553,18 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Verify all tokens have zero balance - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress1]: '0x0', // Zero balance when fetch fails - [tokenAddress2]: '0x0', // Zero balance when fetch fails - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + // Verify all tokens have zero balance + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress1]: '0x0', // Zero balance when fetch fails + [tokenAddress2]: '0x0', // Zero balance when fetch fails + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); }); @@ -2316,22 +2615,24 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Verify both accounts have their respective tokens with appropriate balances - expect(controller.state.tokenBalances).toStrictEqual({ - [account1]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress1]: toHex(500), // Actual fetched balance - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + // Verify both accounts have their respective tokens with appropriate balances + expect(controller.state.tokenBalances).toStrictEqual({ + [account1]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress1]: toHex(500), // Actual fetched balance + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, - [account2]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress2]: '0x0', // Zero balance for missing token - [STAKING_CONTRACT_ADDRESS]: '0x0', + [account2]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress2]: '0x0', // Zero balance for missing token + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); }); @@ -2373,14 +2674,16 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(new BN('1000000000000000000')), - [STAKING_CONTRACT_ADDRESS]: toHex(stakedBalance), + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(new BN('1000000000000000000')), + [STAKING_CONTRACT_ADDRESS]: toHex(stakedBalance), + }, }, - }, + }); }); }); @@ -2433,21 +2736,23 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - expect(controller.state.tokenBalances).toStrictEqual({ - [account1]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(new BN('1000000000000000000')), - [STAKING_CONTRACT_ADDRESS]: toHex(new BN('3000000000000000000')), + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({ + [account1]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(new BN('1000000000000000000')), + [STAKING_CONTRACT_ADDRESS]: toHex(new BN('3000000000000000000')), + }, }, - }, - [account2]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(new BN('2000000000000000000')), - [STAKING_CONTRACT_ADDRESS]: toHex(new BN('4000000000000000000')), + [account2]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(new BN('2000000000000000000')), + [STAKING_CONTRACT_ADDRESS]: toHex(new BN('4000000000000000000')), + }, }, - }, + }); }); }); @@ -2487,14 +2792,16 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(new BN('1000000000000000000')), - [STAKING_CONTRACT_ADDRESS]: '0x0', // Zero balance + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(new BN('1000000000000000000')), + [STAKING_CONTRACT_ADDRESS]: '0x0', // Zero balance + }, }, - }, + }); }); }); @@ -2532,14 +2839,16 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(new BN('1000000000000000000')), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(new BN('1000000000000000000')), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); }); @@ -2582,14 +2891,16 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress]: toHex(new BN('1000000000000000000')), - // No staking contract address for unsupported chain + await waitFor(() => { + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress]: toHex(new BN('1000000000000000000')), + // No staking contract address for unsupported chain + }, }, - }, + }); }); }); }); @@ -2640,9 +2951,11 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // With safelyExecuteWithTimeout, errors are logged as console.error - // and the operation continues gracefully - expect(consoleErrorSpy).toHaveBeenCalledWith(mockError); + await waitFor(() => { + // With safelyExecuteWithTimeout, errors are logged as console.error + // and the operation continues gracefully + expect(consoleErrorSpy).toHaveBeenCalledWith(mockError); + }); // Restore mocks multicallSpy.mockRestore(); @@ -2650,55 +2963,60 @@ describe('TokenBalancesController', () => { }); it('should log error when updateBalances fails after token change', async () => { - const chainId = '0x1'; - const accountAddress = '0x0000000000000000000000000000000000000000'; - const tokenAddress = '0x0000000000000000000000000000000000000001'; - const mockError = new Error('UpdateBalances failed'); + jest.useFakeTimers({ doNotFake: ['nextTick', 'queueMicrotask'] }); + try { + const chainId = '0x1'; + const accountAddress = '0x0000000000000000000000000000000000000000'; + const tokenAddress = '0x0000000000000000000000000000000000000001'; + const mockError = new Error('UpdateBalances failed'); - // Spy on console.warn - const consoleWarnSpy = jest.spyOn(console, 'warn').mockImplementation(); + // Spy on console.warn + const consoleWarnSpy = jest.spyOn(console, 'warn').mockImplementation(); - const { controller, messenger } = setupController(); + const { controller, messenger } = setupController(); - // Mock updateBalances to throw an error - const updateBalancesSpy = jest - .spyOn(controller, 'updateBalances') - .mockRejectedValue(mockError); + // Mock updateBalances to throw an error + const updateBalancesSpy = jest + .spyOn(controller, 'updateBalances') + .mockRejectedValue(mockError); - // Publish a token change that should trigger updateBalances - messenger.publish( - 'TokensController:stateChange', - { - allDetectedTokens: {}, - allIgnoredTokens: {}, - allTokens: { - [chainId]: { - [accountAddress]: [ - { address: tokenAddress, decimals: 0, symbol: 'S' }, - ], + // Publish a token change that should trigger updateBalances + messenger.publish( + 'TokensController:stateChange', + { + allDetectedTokens: {}, + allIgnoredTokens: {}, + allTokens: { + [chainId]: { + [accountAddress]: [ + { address: tokenAddress, decimals: 0, symbol: 'S' }, + ], + }, }, }, - }, - [], - ); + [], + ); - await jestAdvanceTime({ duration: 1 }); + await jestAdvanceTime({ duration: 1 }); - // Verify updateBalances was called - expect(updateBalancesSpy).toHaveBeenCalled(); + // Verify updateBalances was called + expect(updateBalancesSpy).toHaveBeenCalled(); - // Wait a bit more for the catch block to execute - await jestAdvanceTime({ duration: 1 }); + // Wait a bit more for the catch block to execute + await jestAdvanceTime({ duration: 1 }); - // Verify the error was logged - expect(consoleWarnSpy).toHaveBeenCalledWith( - 'Error updating balances after token change:', - mockError, - ); + // Verify the error was logged + expect(consoleWarnSpy).toHaveBeenCalledWith( + 'Error updating balances after token change:', + mockError, + ); - // Restore the original method - updateBalancesSpy.mockRestore(); - consoleWarnSpy.mockRestore(); + // Restore the original method + updateBalancesSpy.mockRestore(); + consoleWarnSpy.mockRestore(); + } finally { + jest.useRealTimers(); + } }); it('should handle timeout scenario', async () => { @@ -2726,9 +3044,6 @@ describe('TokenBalancesController', () => { const { controller } = setupController({ tokens }); - // Use fake timers for precise control - jest.useFakeTimers(); - // Mock safelyExecuteWithTimeout to simulate timeout by returning undefined mockedSafelyExecuteWithTimeout.mockImplementation( async () => undefined, // Simulates timeout behavior @@ -2742,27 +3057,23 @@ describe('TokenBalancesController', () => { stakedBalances: {}, }); - try { - // Start the balance update - should complete gracefully despite timeout - await controller.updateBalances({ - chainIds: [chainId], - queryAllAccounts: true, - }); + // Start the balance update - should complete gracefully despite timeout + await controller.updateBalances({ + chainIds: [chainId], + queryAllAccounts: true, + }); + await waitFor(() => { // With safelyExecuteWithTimeout, timeouts are handled gracefully // The system should continue operating without throwing errors // No specific timeout error message should be logged at controller level // Verify that the update completed without errors expect(controller.state.tokenBalances).toBeDefined(); + }); - // Restore mocks - multicallSpy.mockRestore(); - consoleWarnSpy.mockRestore(); - } finally { - // Always restore timers - jest.useRealTimers(); - } + multicallSpy.mockRestore(); + consoleWarnSpy.mockRestore(); }); }); @@ -2814,15 +3125,17 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Should only have one entry with proper checksum address - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddressProperChecksum]: '0x186a0', // Only checksum version exists - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + // Should only have one entry with proper checksum address + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddressProperChecksum]: '0x186a0', // Only checksum version exists + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); // Verify no duplicate entries exist @@ -2888,16 +3201,18 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // All addresses should be normalized to proper checksum format - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddress1Checksum]: toHex(500), - [tokenAddress2Checksum]: toHex(1000), - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + // All addresses should be normalized to proper checksum format + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddress1Checksum]: toHex(500), + [tokenAddress2Checksum]: toHex(1000), + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); }); @@ -2943,15 +3258,17 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Should only have one normalized entry with proper checksum - expect(controller.state.tokenBalances).toStrictEqual({ - [accountAddress]: { - [chainId]: { - [NATIVE_TOKEN_ADDRESS]: '0x0', - [tokenAddressChecksum]: '0x186a0', // Only checksum version exists - [STAKING_CONTRACT_ADDRESS]: '0x0', + await waitFor(() => { + // Should only have one normalized entry with proper checksum + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddress]: { + [chainId]: { + [NATIVE_TOKEN_ADDRESS]: '0x0', + [tokenAddressChecksum]: '0x186a0', // Only checksum version exists + [STAKING_CONTRACT_ADDRESS]: '0x0', + }, }, - }, + }); }); // Verify no case variations exist as separate keys @@ -3004,11 +3321,13 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // Should have balances set for the account and chain - expect(controller.state.tokenBalances[accountAddress]).toBeDefined(); - expect( - controller.state.tokenBalances[accountAddress][chainId], - ).toBeDefined(); + await waitFor(() => { + // Should have balances set for the account and chain + expect(controller.state.tokenBalances[accountAddress]).toBeDefined(); + expect( + controller.state.tokenBalances[accountAddress][chainId], + ).toBeDefined(); + }); const chainBalances = controller.state.tokenBalances[accountAddress][chainId]; @@ -3085,20 +3404,21 @@ describe('TokenBalancesController', () => { chainIds: [chainId], queryAllAccounts: false, }); - - // Verify that getTokenBalancesForMultipleAddresses was called with only the selected account - expect(mockGetTokenBalances).toHaveBeenCalledWith( - [ - { - accountAddress: selectedAccount, - tokenAddresses: [tokenAddress, NATIVE_TOKEN_ADDRESS], - }, - ], - chainId, - expect.any(Object), // provider - true, // include native - true, // include staked - ); + await waitFor(() => { + // Verify that getTokenBalancesForMultipleAddresses was called with only the selected account + expect(mockGetTokenBalances).toHaveBeenCalledWith( + [ + { + accountAddress: selectedAccount, + tokenAddresses: [tokenAddress, NATIVE_TOKEN_ADDRESS], + }, + ], + chainId, + expect.any(Object), // provider + true, // include native + true, // include staked + ); + }); // Should only contain balance for selected account when queryMultipleAccounts is false expect(controller.state.tokenBalances).toStrictEqual({ @@ -3162,6 +3482,13 @@ describe('TokenBalancesController', () => { }); describe('Per-chain polling intervals', () => { + beforeEach(() => { + jest.useFakeTimers({ doNotFake: ['nextTick', 'queueMicrotask'] }); + }); + afterEach(() => { + jest.useRealTimers(); + }); + it('should use default interval when no chain-specific config is provided', () => { const defaultInterval = 30000; const { controller } = setupController({ @@ -4065,147 +4392,162 @@ describe('TokenBalancesController', () => { describe('Error handling and edge cases', () => { it('should handle polling errors gracefully', async () => { - const chainId = '0x1'; - const accountAddress = '0x0000000000000000000000000000000000000000'; - const tokenAddress = '0x0000000000000000000000000000000000000001'; + jest.useFakeTimers({ doNotFake: ['nextTick', 'queueMicrotask'] }); + try { + const chainId = '0x1'; + const accountAddress = '0x0000000000000000000000000000000000000000'; + const tokenAddress = '0x0000000000000000000000000000000000000001'; - const tokens = { - allDetectedTokens: {}, - allTokens: { - [chainId]: { - [accountAddress]: [ - { address: tokenAddress, symbol: 'TEST', decimals: 18 }, - ], + const tokens = { + allDetectedTokens: {}, + allTokens: { + [chainId]: { + [accountAddress]: [ + { address: tokenAddress, symbol: 'TEST', decimals: 18 }, + ], + }, }, - }, - }; + }; - const consoleSpy = jest.spyOn(console, 'warn').mockImplementation(); + const consoleSpy = jest.spyOn(console, 'warn').mockImplementation(); - const { controller } = setupController({ - tokens, - config: { interval: 100 }, - }); + const { controller } = setupController({ + tokens, + config: { interval: 100 }, + }); - // Mock _executePoll to throw an error - const pollSpy = jest - .spyOn(controller, '_executePoll') - .mockRejectedValue(new Error('Polling failed')); + // Mock _executePoll to throw an error + const pollSpy = jest + .spyOn(controller, '_executePoll') + .mockRejectedValue(new Error('Polling failed')); - controller.startPolling({ chainIds: ['0x1'] }); + controller.startPolling({ chainIds: ['0x1'] }); - // Wait for initial poll and error - await jestAdvanceTime({ duration: 1 }); + // Wait for initial poll and error + await jestAdvanceTime({ duration: 1 }); - // Wait for interval poll and error - await jestAdvanceTime({ duration: 100 }); + // Wait for interval poll and error + await jestAdvanceTime({ duration: 100 }); - // Should have attempted polls despite errors - expect(pollSpy).toHaveBeenCalledTimes(2); + // Should have attempted polls despite errors + expect(pollSpy).toHaveBeenCalledTimes(2); - controller.stopAllPolling(); - consoleSpy.mockRestore(); + controller.stopAllPolling(); + consoleSpy.mockRestore(); + } finally { + jest.useRealTimers(); + } }); it('should handle updateBalances errors in token change handler', async () => { - const chainId = '0x1'; - const accountAddress = '0x0000000000000000000000000000000000000000'; - const tokenAddress = '0x0000000000000000000000000000000000000001'; + jest.useFakeTimers({ doNotFake: ['nextTick', 'queueMicrotask'] }); + try { + const chainId = '0x1'; + const accountAddress = '0x0000000000000000000000000000000000000000'; + const tokenAddress = '0x0000000000000000000000000000000000000001'; - const tokens = { - allDetectedTokens: {}, - allTokens: { - [chainId]: { - [accountAddress]: [ - { address: tokenAddress, symbol: 'TEST', decimals: 18 }, - ], + const tokens = { + allDetectedTokens: {}, + allTokens: { + [chainId]: { + [accountAddress]: [ + { address: tokenAddress, symbol: 'TEST', decimals: 18 }, + ], + }, }, - }, - }; + }; - const consoleSpy = jest.spyOn(console, 'warn').mockImplementation(); + const consoleSpy = jest.spyOn(console, 'warn').mockImplementation(); - const { controller, messenger } = setupController({ - tokens, - }); + const { controller, messenger } = setupController({ + tokens, + }); - // Mock updateBalances to throw an error - const updateBalancesSpy = jest - .spyOn(controller, 'updateBalances') - .mockRejectedValue(new Error('Update failed')); + // Mock updateBalances to throw an error + const updateBalancesSpy = jest + .spyOn(controller, 'updateBalances') + .mockRejectedValue(new Error('Update failed')); - // Simulate token change that triggers balance update - const newTokens = { - ...tokens, - allTokens: { - [chainId]: { - [accountAddress]: [ - { address: tokenAddress, symbol: 'TEST', decimals: 18 }, - { - address: '0x0000000000000000000000000000000000000002', - symbol: 'NEW', - decimals: 18, - }, - ], + // Simulate token change that triggers balance update + const newTokens = { + ...tokens, + allTokens: { + [chainId]: { + [accountAddress]: [ + { address: tokenAddress, symbol: 'TEST', decimals: 18 }, + { + address: '0x0000000000000000000000000000000000000002', + symbol: 'NEW', + decimals: 18, + }, + ], + }, }, - }, - allIgnoredTokens: {}, - ignoredTokens: [], - detectedTokens: [], - tokens: [], - }; + allIgnoredTokens: {}, + ignoredTokens: [], + detectedTokens: [], + tokens: [], + }; - // Trigger token change by publishing state change - messenger.publish('TokensController:stateChange', newTokens, [ - { op: 'replace', path: [], value: newTokens }, - ]); + // Trigger token change by publishing state change + messenger.publish('TokensController:stateChange', newTokens, [ + { op: 'replace', path: [], value: newTokens }, + ]); - // Wait for async error handling - await jestAdvanceTime({ duration: 1 }); + // Wait for async error handling + await jestAdvanceTime({ duration: 1 }); - expect(updateBalancesSpy).toHaveBeenCalled(); - expect(consoleSpy).toHaveBeenCalledWith( - 'Error updating balances after token change:', - expect.any(Error), - ); + expect(updateBalancesSpy).toHaveBeenCalled(); + expect(consoleSpy).toHaveBeenCalledWith( + 'Error updating balances after token change:', + expect.any(Error), + ); - consoleSpy.mockRestore(); + consoleSpy.mockRestore(); + } finally { + jest.useRealTimers(); + } }); it('should handle malformed JSON in _stopPollingByPollingTokenSetId gracefully', async () => { - const { controller } = setupController(); + jest.useFakeTimers({ doNotFake: ['nextTick', 'queueMicrotask'] }); + try { + const { controller } = setupController(); - // Start polling to create an active session - controller.startPolling({ chainIds: ['0x1', '0x2'] }); + // Start polling to create an active session + controller.startPolling({ chainIds: ['0x1', '0x2'] }); - const consoleSpy = jest.spyOn(console, 'warn').mockImplementation(); + const consoleSpy = jest.spyOn(console, 'warn').mockImplementation(); - // Call with malformed JSON - this should trigger the fallback behavior - const malformedTokenSetId = '{invalid json}'; - controller._stopPollingByPollingTokenSetId(malformedTokenSetId); + // Call with malformed JSON - this should trigger the fallback behavior + const malformedTokenSetId = '{invalid json}'; + controller._stopPollingByPollingTokenSetId(malformedTokenSetId); - // Should log the error - expect(consoleSpy).toHaveBeenCalledWith( - 'Failed to parse tokenSetId, stopping all polling:', - expect.any(SyntaxError), - ); + // Should log the error + expect(consoleSpy).toHaveBeenCalledWith( + 'Failed to parse tokenSetId, stopping all polling:', + expect.any(SyntaxError), + ); - // Verify that controller can recover by starting new polling session successfully - // This demonstrates that the fallback stop-all-polling behavior worked - const updateBalancesSpy = jest - .spyOn(controller, 'updateBalances') - .mockResolvedValue(); + // Verify that controller can recover by starting new polling session successfully + // This demonstrates that the fallback stop-all-polling behavior worked + const updateBalancesSpy = jest + .spyOn(controller, 'updateBalances') + .mockResolvedValue(); - // Start new polling session - should work normally after error recovery - controller.startPolling({ chainIds: ['0x1'] }); + // Start new polling session - should work normally after error recovery + controller.startPolling({ chainIds: ['0x1'] }); - // Wait for any immediate polling to complete - await jestAdvanceTime({ duration: 1 }); + // Wait for any immediate polling to complete + await jestAdvanceTime({ duration: 1 }); - // Clean up - controller.stopAllPolling(); - consoleSpy.mockRestore(); - updateBalancesSpy.mockRestore(); + // Clean up + controller.stopAllPolling(); + consoleSpy.mockRestore(); + updateBalancesSpy.mockRestore(); + } finally { + jest.useRealTimers(); + } }); it('should properly destroy controller and cleanup resources', () => { @@ -4262,14 +4604,16 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - // With safelyExecuteWithTimeout timeout simulation, the system should continue operating - // The controller should have initialized the token with 0 balance despite timeout - expect(controller.state.tokenBalances).toStrictEqual({ - '0x0000000000000000000000000000000000000000': { - '0x1': { - '0x0000000000000000000000000000000000000001': '0x0', + await waitFor(() => { + // With safelyExecuteWithTimeout timeout simulation, the system should continue operating + // The controller should have initialized the token with 0 balance despite timeout + expect(controller.state.tokenBalances).toStrictEqual({ + '0x0000000000000000000000000000000000000000': { + '0x1': { + '0x0000000000000000000000000000000000000001': '0x0', + }, }, - }, + }); }); // Restore the mock to its default behavior @@ -4417,9 +4761,9 @@ describe('TokenBalancesController', () => { // First call: external services disabled, should use RPC fetcher await controller.updateBalances({ chainIds: [chainId] }); - - // RPC fetcher should have been called since external services are disabled - expect(multicallSpy).toHaveBeenCalled(); + await waitFor(() => { + expect(multicallSpy).toHaveBeenCalled(); + }); multicallSpy.mockClear(); // Now enable external services - this should be respected dynamically @@ -4430,11 +4774,13 @@ describe('TokenBalancesController', () => { // (though it may still fall back to RPC if the API call fails in test) await controller.updateBalances({ chainIds: [chainId] }); - // The test verifies that the allowExternalServices function is evaluated - // dynamically by checking that the controller was constructed successfully - // and that balance updates work in both states - expect(controller).toBeDefined(); - expect(controller.state.tokenBalances).toBeDefined(); + await waitFor(() => { + // The test verifies that the allowExternalServices function is evaluated + // dynamically by checking that the controller was constructed successfully + // and that balance updates work in both states + expect(controller).toBeDefined(); + expect(controller.state.tokenBalances).toBeDefined(); + }); multicallSpy.mockRestore(); }); @@ -4545,8 +4891,9 @@ describe('TokenBalancesController', () => { // This should trigger the continue statement (line 440) when no chains are supported await controller.updateBalances({ chainIds: [chainId] }); - - expect(mockSupports).toHaveBeenCalledWith(chainId); + await waitFor(() => { + expect(mockSupports).toHaveBeenCalledWith(chainId); + }); mockSupports.mockRestore(); }); @@ -4634,13 +4981,25 @@ describe('TokenBalancesController', () => { const originalFetch = global.fetch; global.fetch = mockGlobalFetch; - // Create controller with accountsApiChainIds to enable AccountsApi fetcher + // Create controller with accountsApiChainIds to enable AccountsApi fetcher; tokens so we fetch for chainId1 + const tokenAddress = '0x0000000000000000000000000000000000000001'; const { controller } = setupController({ config: { accountsApiChainIds: () => [chainId1, chainId2], // This enables AccountsApi for these chains allowExternalServices: () => true, }, listAccounts: [account], + tokens: { + allTokens: { + [chainId1]: { + [accountAddress]: [ + { address: tokenAddress, symbol: 'T', decimals: 18 }, + ], + }, + }, + allDetectedTokens: {}, + allIgnoredTokens: {}, + }, }); // Reset mocks after controller creation @@ -4650,13 +5009,16 @@ describe('TokenBalancesController', () => { // Test Case 1: Execute line 517 -> line 320 with chainId returned by accountsApiChainIds() mockSupports.mockReturnValue(true); await controller.updateBalances({ chainIds: [chainId1] }); // This triggers line 517 -> line 320 - - // Verify line 320 logic was executed (originalFetcher.supports was called) - expect(mockSupports).toHaveBeenCalledWith(chainId1); + await waitFor(() => { + expect(mockSupports).toHaveBeenCalledWith(chainId1); + }); // Test Case 2: Execute line 517 -> line 320 with chainId NOT returned by accountsApiChainIds() mockSupports.mockClear(); await controller.updateBalances({ chainIds: [chainId3] }); // This triggers line 517 -> line 320 + await new Promise((resolve) => + setTimeout(resolve, UPDATE_BALANCES_BATCH_MS + 100), + ); // Allow debounce + async to complete // Should NOT have called originalFetcher.supports because chainId3 is not returned by accountsApiChainIds() // This tests the short-circuit evaluation on line 322: this.#accountsApiChainIds().includes(chainId) @@ -4721,9 +5083,6 @@ describe('TokenBalancesController', () => { ], }); - // Wait for async update - await flushPromises(); - // Verify balance was updated (account addresses are lowercase in state) const checksumTokenAddress = tokenAddress; expect( @@ -4767,34 +5126,33 @@ describe('TokenBalancesController', () => { }, postBalance: { amount: '0xde0b6b3a7640000', // 1 ETH in wei - }, - transfers: [], - }, - ], - }); - - // Wait for async update - await flushPromises(); - - // Verify native balance was updated in TokenBalancesController (account addresses are lowercase in state) - const lowercaseAddr = accountAddress.toLowerCase(); - expect( - controller.state.tokenBalances[lowercaseAddr as ChecksumAddress]?.[ - chainId - ]?.[NATIVE_TOKEN_ADDRESS], - ).toBe('0xde0b6b3a7640000'); - - // Verify AccountTrackerController was called - expect(updateNativeBalancesSpy).toHaveBeenCalledWith( - 'AccountTrackerController:updateNativeBalances', - [ - { - address: lowercaseAddr, - chainId, - balance: '0xde0b6b3a7640000', + }, + transfers: [], }, ], - ); + }); + + await waitFor(() => { + // Verify native balance was updated in TokenBalancesController (account addresses are lowercase in state) + const lowercaseAddr = accountAddress.toLowerCase(); + expect( + controller.state.tokenBalances[lowercaseAddr as ChecksumAddress]?.[ + chainId + ]?.[NATIVE_TOKEN_ADDRESS], + ).toBe('0xde0b6b3a7640000'); + + // Verify AccountTrackerController was called + expect(updateNativeBalancesSpy).toHaveBeenCalledWith( + 'AccountTrackerController:updateNativeBalances', + [ + { + address: lowercaseAddr, + chainId, + balance: '0xde0b6b3a7640000', + }, + ], + ); + }); }); it('should handle balance update errors and trigger fallback polling', async () => { @@ -4828,11 +5186,10 @@ describe('TokenBalancesController', () => { ], }); - // Wait for async update - await flushPromises(); - - // Verify fallback polling was triggered - expect(updateBalancesSpy).toHaveBeenCalledWith({ chainIds: ['0x1'] }); + await waitFor(() => { + // Verify fallback polling was triggered + expect(updateBalancesSpy).toHaveBeenCalledWith({ chainIds: ['0x1'] }); + }); }); it('should handle unsupported asset types and trigger fallback polling', async () => { @@ -4864,12 +5221,10 @@ describe('TokenBalancesController', () => { }, ], }); - - // Wait for async update - await flushPromises(); - - // Verify fallback polling was triggered - expect(updateBalancesSpy).toHaveBeenCalledWith({ chainIds: ['0x1'] }); + await waitFor(() => { + // Verify fallback polling was triggered + expect(updateBalancesSpy).toHaveBeenCalledWith({ chainIds: ['0x1'] }); + }); }); it('should handle status change to "up" and increase polling interval', async () => { @@ -5060,20 +5415,19 @@ describe('TokenBalancesController', () => { ], }); - // Wait for async update - await flushPromises(); - - // Verify both balances were updated (account addresses are lowercase in state) - expect( - controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ - '0x1' - ]?.[token1], - ).toBe('0xf4240'); - expect( - controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ - '0x1' - ]?.[token2], - ).toBe('0x1e8480'); + await waitFor(() => { + // Verify both balances were updated (account addresses are lowercase in state) + expect( + controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ + '0x1' + ]?.[token1], + ).toBe('0xf4240'); + expect( + controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ + '0x1' + ]?.[token2], + ).toBe('0x1e8480'); + }); }); it('should handle invalid token addresses and trigger fallback polling', async () => { @@ -5103,10 +5457,9 @@ describe('TokenBalancesController', () => { }, ], }); - - await flushPromises(); - - expect(updateBalancesSpy).toHaveBeenCalledWith({ chainIds: ['0x1'] }); + await waitFor(() => { + expect(updateBalancesSpy).toHaveBeenCalledWith({ chainIds: ['0x1'] }); + }); }); it('should handle status changes with hex chain ID format', async () => { @@ -5183,23 +5536,21 @@ describe('TokenBalancesController', () => { }, ], }); + await waitFor(() => { + // Verify addDetectedTokensViaWs was called with the new token addresses and chainId + expect(addTokensSpy).toHaveBeenCalledWith({ + tokensSlice: [newTokenAddress], + chainId, + }); - // Wait for async processing - await flushPromises(); - - // Verify addDetectedTokensViaWs was called with the new token addresses and chainId - expect(addTokensSpy).toHaveBeenCalledWith({ - tokensSlice: [newTokenAddress], - chainId, + // Verify balance was updated from websocket (account addresses are lowercase in state) + const lowercaseAddr2 = accountAddress.toLowerCase(); + expect( + controller.state.tokenBalances[lowercaseAddr2 as ChecksumAddress]?.[ + chainId + ]?.[newTokenAddress], + ).toBe('0xf4240'); }); - - // Verify balance was updated from websocket (account addresses are lowercase in state) - const lowercaseAddr2 = accountAddress.toLowerCase(); - expect( - controller.state.tokenBalances[lowercaseAddr2 as ChecksumAddress]?.[ - chainId - ]?.[newTokenAddress], - ).toBe('0xf4240'); }); it('should process tracked tokens from allTokens without calling addTokens', async () => { @@ -5261,18 +5612,17 @@ describe('TokenBalancesController', () => { ], }); - // Wait for async processing - await flushPromises(); - - // Verify addTokens was NOT called since token is already tracked - expect(addTokensSpy).not.toHaveBeenCalled(); + await waitFor(() => { + // Verify addTokens was NOT called since token is already tracked + expect(addTokensSpy).not.toHaveBeenCalled(); - // Verify balance was updated (account addresses are lowercase in state) - expect( - controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ - chainId - ]?.[trackedTokenAddress], - ).toBe('0xf4240'); + // Verify balance was updated (account addresses are lowercase in state) + expect( + controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ + chainId + ]?.[trackedTokenAddress], + ).toBe('0xf4240'); + }); }); it('should process ignored tokens from allIgnoredTokens without calling addTokens', async () => { @@ -5328,18 +5678,17 @@ describe('TokenBalancesController', () => { ], }); - // Wait for async processing - await flushPromises(); - - // Verify addTokens was NOT called since token is ignored (tracked) - expect(addTokensSpy).not.toHaveBeenCalled(); + await waitFor(() => { + // Verify addTokens was NOT called since token is ignored (tracked) + expect(addTokensSpy).not.toHaveBeenCalled(); - // Verify balance was still updated (ignored tokens should still have balances tracked, account addresses are lowercase in state) - expect( - controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ - chainId - ]?.[ignoredTokenAddress], - ).toBe('0xf4240'); + // Verify balance was still updated (ignored tokens should still have balances tracked, account addresses are lowercase in state) + expect( + controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ + chainId + ]?.[ignoredTokenAddress], + ).toBe('0xf4240'); + }); }); it('should handle native tokens without checking if they are tracked', async () => { @@ -5389,19 +5738,18 @@ describe('TokenBalancesController', () => { ], }); - // Wait for async processing - await flushPromises(); - - // Verify addTokens was NOT called for native token - expect(addTokensSpy).not.toHaveBeenCalled(); + await waitFor(() => { + // Verify addTokens was NOT called for native token + expect(addTokensSpy).not.toHaveBeenCalled(); - // Verify native balance was updated (account addresses are lowercase in state) - const lowercaseAddr3 = accountAddress.toLowerCase(); - expect( - controller.state.tokenBalances[lowercaseAddr3 as ChecksumAddress]?.[ - chainId - ]?.[NATIVE_TOKEN_ADDRESS], - ).toBe('0xde0b6b3a7640000'); + // Verify native balance was updated (account addresses are lowercase in state) + const lowercaseAddr3 = accountAddress.toLowerCase(); + expect( + controller.state.tokenBalances[lowercaseAddr3 as ChecksumAddress]?.[ + chainId + ]?.[NATIVE_TOKEN_ADDRESS], + ).toBe('0xde0b6b3a7640000'); + }); }); it('should handle addTokens errors and trigger fallback polling', async () => { @@ -5459,17 +5807,16 @@ describe('TokenBalancesController', () => { ], }); - // Wait for async processing - await flushPromises(); - - // Verify error was logged - expect(consoleSpy).toHaveBeenCalledWith( - 'Error updating balances from AccountActivityService for chain eip155:1, account 0x1234567890123456789012345678901234567890:', - expect.any(Error), - ); + await waitFor(() => { + // Verify error was logged + expect(consoleSpy).toHaveBeenCalledWith( + 'Error updating balances from AccountActivityService for chain eip155:1, account 0x1234567890123456789012345678901234567890:', + expect.any(Error), + ); - // Verify fallback polling was triggered (once in addTokens error handler) - expect(updateBalancesSpy).toHaveBeenCalledWith({ chainIds: ['0x1'] }); + // Verify fallback polling was triggered (once in addTokens error handler) + expect(updateBalancesSpy).toHaveBeenCalledWith({ chainIds: ['0x1'] }); + }); consoleSpy.mockRestore(); }); @@ -5545,26 +5892,25 @@ describe('TokenBalancesController', () => { ], }); - // Wait for async processing - await flushPromises(); + await waitFor(() => { + // Verify addTokens was called only for the untracked token with networkClientId + expect(addTokensSpy).toHaveBeenCalledWith({ + tokensSlice: [untrackedToken], + chainId, + }); - // Verify addTokens was called only for the untracked token with networkClientId - expect(addTokensSpy).toHaveBeenCalledWith({ - tokensSlice: [untrackedToken], - chainId, + // Verify both token balances were updated from websocket (account addresses are lowercase in state) + expect( + controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ + chainId + ]?.[trackedToken], + ).toBe('0xf4240'); + expect( + controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ + chainId + ]?.[untrackedToken], + ).toBe('0x1e8480'); }); - - // Verify both token balances were updated from websocket (account addresses are lowercase in state) - expect( - controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ - chainId - ]?.[trackedToken], - ).toBe('0xf4240'); - expect( - controller.state.tokenBalances[lowercaseAddress as ChecksumAddress]?.[ - chainId - ]?.[untrackedToken], - ).toBe('0x1e8480'); }); it('should cleanup debouncing timer on destroy', () => { @@ -5601,6 +5947,7 @@ describe('TokenBalancesController', () => { mockAPIAccountsAPIMultichainAccountBalancesCamelCase(accountAddress); const account = createMockInternalAccount({ address: accountAddress }); + const tokenAddress = '0x2791bca1f2de4661ed88a30c99a7a9449aa84174'; // USDC from mock response const { controller } = setupController({ config: { @@ -5608,6 +5955,17 @@ describe('TokenBalancesController', () => { allowExternalServices: () => true, }, listAccounts: [account], + tokens: { + allTokens: { + [chainId]: { + [accountAddress]: [ + { address: tokenAddress, symbol: 'USDC', decimals: 6 }, + ], + }, + }, + allDetectedTokens: {}, + allIgnoredTokens: {}, + }, }); return { @@ -5624,10 +5982,12 @@ describe('TokenBalancesController', () => { queryAllAccounts: true, }); - expect(controller.state.tokenBalances[accountAddress]).toBeDefined(); - expect( - controller.state.tokenBalances[checksumAccountAddress], - ).toBeUndefined(); + await waitFor(() => { + expect(controller.state.tokenBalances[accountAddress]).toBeDefined(); + expect( + controller.state.tokenBalances[checksumAccountAddress], + ).toBeUndefined(); + }); expect(mockAccountsAPI.isDone()).toBe(true); }); @@ -5788,6 +6148,13 @@ describe('TokenBalancesController', () => { }); describe('polling behavior', () => { + beforeEach(() => { + jest.useFakeTimers({ doNotFake: ['nextTick', 'queueMicrotask'] }); + }); + afterEach(() => { + jest.useRealTimers(); + }); + it('should not poll when controller polling is not active', async () => { const { controller } = setupController({ config: { @@ -5977,6 +6344,13 @@ describe('TokenBalancesController', () => { }); describe('polling timer management', () => { + beforeEach(() => { + jest.useFakeTimers({ doNotFake: ['nextTick', 'queueMicrotask'] }); + }); + afterEach(() => { + jest.useRealTimers(); + }); + it('should clear existing timer when setting new one for same interval', async () => { // This test verifies line 586 const { controller } = setupController({ @@ -5999,119 +6373,129 @@ describe('TokenBalancesController', () => { }); it('should handle immediate polling errors gracefully', async () => { - // This test verifies that errors in updateBalances are caught by the polling error handler - const consoleWarnSpy = jest - .spyOn(console, 'warn') - .mockImplementation(() => undefined); + jest.useFakeTimers({ doNotFake: ['nextTick', 'queueMicrotask'] }); + try { + // This test verifies that errors in updateBalances are caught by the polling error handler + const consoleWarnSpy = jest + .spyOn(console, 'warn') + .mockImplementation(() => undefined); - const selectedAccount = createMockInternalAccount({ - address: '0x0000000000000000000000000000000000000001', - }); + const selectedAccount = createMockInternalAccount({ + address: '0x0000000000000000000000000000000000000001', + }); - const { controller, messenger } = setupController({ - config: { - accountsApiChainIds: () => [], - }, - listAccounts: [selectedAccount], - tokens: { - allTokens: { - '0x1': { - [selectedAccount.address]: [ - { - address: '0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48', - symbol: 'USDC', - decimals: 6, - }, - ], + const { controller, messenger } = setupController({ + config: { + accountsApiChainIds: () => [], + }, + listAccounts: [selectedAccount], + tokens: { + allTokens: { + '0x1': { + [selectedAccount.address]: [ + { + address: '0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48', + symbol: 'USDC', + decimals: 6, + }, + ], + }, }, + allDetectedTokens: {}, + allIgnoredTokens: {}, }, - allDetectedTokens: {}, - allIgnoredTokens: {}, - }, - }); + }); - // Unregister handler and re-register to cause an error in updateBalances - // Breaking AccountsController:getSelectedAccount causes error before #fetchAllBalances - messenger.unregisterActionHandler( - 'AccountsController:getSelectedAccount', - ); - messenger.registerActionHandler( - 'AccountsController:getSelectedAccount', - () => { - throw new Error('Account error'); - }, - ); + // Unregister handler and re-register to cause an error in updateBalances + // Breaking AccountsController:getSelectedAccount causes error before #fetchAllBalances + messenger.unregisterActionHandler( + 'AccountsController:getSelectedAccount', + ); + messenger.registerActionHandler( + 'AccountsController:getSelectedAccount', + () => { + throw new Error('Account error'); + }, + ); - controller.startPolling({ chainIds: ['0x1'] }); + controller.startPolling({ chainIds: ['0x1'] }); - await jest.advanceTimersByTimeAsync(100); + await jest.advanceTimersByTimeAsync(100); - expect(consoleWarnSpy).toHaveBeenCalledWith( - expect.stringContaining('Polling failed'), - expect.anything(), - ); + expect(consoleWarnSpy).toHaveBeenCalledWith( + expect.stringContaining('Polling failed'), + expect.anything(), + ); - consoleWarnSpy.mockRestore(); + consoleWarnSpy.mockRestore(); + } finally { + jest.useRealTimers(); + } }); it('should handle interval polling errors gracefully', async () => { - // This test verifies that errors in interval polling are caught and logged - const consoleWarnSpy = jest - .spyOn(console, 'warn') - .mockImplementation(() => undefined); + jest.useFakeTimers({ doNotFake: ['nextTick', 'queueMicrotask'] }); + try { + // This test verifies that errors in interval polling are caught and logged + const consoleWarnSpy = jest + .spyOn(console, 'warn') + .mockImplementation(() => undefined); - const selectedAccount = createMockInternalAccount({ - address: '0x0000000000000000000000000000000000000001', - }); + const selectedAccount = createMockInternalAccount({ + address: '0x0000000000000000000000000000000000000001', + }); - const { controller, messenger } = setupController({ - config: { - accountsApiChainIds: () => [], - interval: 1000, - }, - listAccounts: [selectedAccount], - tokens: { - allTokens: { - '0x1': { - [selectedAccount.address]: [ - { - address: '0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48', - symbol: 'USDC', - decimals: 6, - }, - ], + const { controller, messenger } = setupController({ + config: { + accountsApiChainIds: () => [], + interval: 1000, + }, + listAccounts: [selectedAccount], + tokens: { + allTokens: { + '0x1': { + [selectedAccount.address]: [ + { + address: '0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48', + symbol: 'USDC', + decimals: 6, + }, + ], + }, }, + allDetectedTokens: {}, + allIgnoredTokens: {}, }, - allDetectedTokens: {}, - allIgnoredTokens: {}, - }, - }); + }); - controller.startPolling({ chainIds: ['0x1'] }); + controller.startPolling({ chainIds: ['0x1'] }); - await jest.advanceTimersByTimeAsync(100); + await jest.advanceTimersByTimeAsync(100); - // Now break the handler to cause errors on subsequent polls - // Breaking AccountsController:getSelectedAccount causes error before #fetchAllBalances - messenger.unregisterActionHandler( - 'AccountsController:getSelectedAccount', - ); - messenger.registerActionHandler( - 'AccountsController:getSelectedAccount', - () => { - throw new Error('Account error'); - }, - ); + // Now break the handler to cause errors on subsequent polls + // Breaking AccountsController:getSelectedAccount causes error before #fetchAllBalances + messenger.unregisterActionHandler( + 'AccountsController:getSelectedAccount', + ); + messenger.registerActionHandler( + 'AccountsController:getSelectedAccount', + () => { + throw new Error('Account error'); + }, + ); - // Wait for interval polling to trigger - await jest.advanceTimersByTimeAsync(1500); + // Wait for interval polling to trigger + await jest.advanceTimersByTimeAsync(1500); - expect(consoleWarnSpy).toHaveBeenCalledWith( - expect.stringContaining('Polling failed'), - expect.anything(), - ); + expect(consoleWarnSpy).toHaveBeenCalledWith( + expect.stringContaining('Polling failed'), + expect.anything(), + ); - consoleWarnSpy.mockRestore(); + consoleWarnSpy.mockRestore(); + } finally { + jest.useRealTimers(); + } }); }); @@ -6202,12 +6586,24 @@ describe('TokenBalancesController', () => { const selectedAccount = createMockInternalAccount({ address: '0x1234567890123456789012345678901234567890', }); + const tokenAddress = '0x0000000000000000000000000000000000000001'; const { controller, messenger } = setupController({ listAccounts: [selectedAccount], config: { accountsApiChainIds: () => [], }, + tokens: { + allTokens: { + '0x1': { + [selectedAccount.address]: [ + { address: tokenAddress, symbol: 'T', decimals: 18 }, + ], + }, + }, + allDetectedTokens: {}, + allIgnoredTokens: {}, + }, }); // Lock and then unlock @@ -6224,9 +6620,9 @@ describe('TokenBalancesController', () => { // updateBalances should proceed after unlock await controller.updateBalances({ chainIds: ['0x1'] }); - - // Verify fetch WAS called because isActive is true - expect(fetchSpy).toHaveBeenCalled(); + await waitFor(() => { + expect(fetchSpy).toHaveBeenCalled(); + }); fetchSpy.mockRestore(); }); @@ -6404,15 +6800,14 @@ describe('TokenBalancesController', () => { chainIds: [chainId], tokenAddresses: [token1, token2], }); - - const balances = - controller.state.tokenBalances[accountAddress as ChecksumAddress]?.[ - chainId - ]; - - expect(balances).toBeDefined(); - expect(balances?.[token1 as ChecksumAddress]).toBeDefined(); - expect(balances?.[token2 as ChecksumAddress]).toBeDefined(); + await waitFor(() => { + const balances = + controller.state.tokenBalances[accountAddress as ChecksumAddress]?.[ + chainId + ]; + expect(balances?.[token1 as ChecksumAddress]).toBeDefined(); + expect(balances?.[token2 as ChecksumAddress]).toBeDefined(); + }); // token3 should also be present because multicall returns all tokens // The filtering happens at the fetcher level, not the state update level }); @@ -6451,15 +6846,15 @@ describe('TokenBalancesController', () => { chainIds: [chainId], queryAllAccounts: true, }); - - const balances = - controller.state.tokenBalances[accountAddress as ChecksumAddress]?.[ - chainId - ]; - - // Both tokens should have their returned balances - expect(balances?.[token1 as ChecksumAddress]).toBe(toHex(100)); - expect(balances?.[token2 as ChecksumAddress]).toBe(toHex(200)); + await waitFor(() => { + const balances = + controller.state.tokenBalances[accountAddress as ChecksumAddress]?.[ + chainId + ]; + // Both tokens should have their returned balances + expect(balances?.[token1 as ChecksumAddress]).toBe(toHex(100)); + expect(balances?.[token2 as ChecksumAddress]).toBe(toHex(200)); + }); }); it('should not call addDetectedTokensViaWs for empty token arrays (line 1082)', async () => { @@ -6611,14 +7006,13 @@ describe('TokenBalancesController', () => { // Start polling - the poll function catches errors and logs them controller.startPolling({ chainIds: ['0x1'] }); - // Wait for the promise to be caught - await flushPromises(); - - // Verify warning was logged (either immediate or interval polling message) - expect(consoleWarnSpy).toHaveBeenCalledWith( - expect.stringContaining('Polling failed'), - expect.any(Error), - ); + await waitFor(() => { + // Verify warning was logged (either immediate or interval polling message) + expect(consoleWarnSpy).toHaveBeenCalledWith( + expect.stringContaining('Polling failed'), + expect.any(Error), + ); + }); controller.stopAllPolling(); consoleWarnSpy.mockRestore(); @@ -6668,15 +7062,16 @@ describe('TokenBalancesController', () => { chainIds: [chainId], queryAllAccounts: true, }); - - // Verify detectTokens was called with forceRpc for unprocessed chains - expect(messengerCallSpy).toHaveBeenCalledWith( - 'TokenDetectionController:detectTokens', - { - chainIds: ['0x89'], - forceRpc: true, - }, - ); + await waitFor(() => { + // Verify detectTokens was called with forceRpc for unprocessed chains + expect(messengerCallSpy).toHaveBeenCalledWith( + 'TokenDetectionController:detectTokens', + { + chainIds: ['0x89'], + forceRpc: true, + }, + ); + }); messengerCallSpy.mockRestore(); }); @@ -6746,18 +7141,19 @@ describe('TokenBalancesController', () => { chainIds: [chainId], queryAllAccounts: true, }); - - expect(apiFetchSpy).toHaveBeenCalled(); - expect(rpcFetchSpy).toHaveBeenCalledWith( - expect.objectContaining({ - chainIds: [chainId], - unprocessedTokens: { - [accountAddress]: { - [chainId]: [token1], + await waitFor(() => { + expect(apiFetchSpy).toHaveBeenCalled(); + expect(rpcFetchSpy).toHaveBeenCalledWith( + expect.objectContaining({ + chainIds: [chainId], + unprocessedTokens: { + [accountAddress]: { + [chainId]: [token1], + }, }, - }, - }), - ); + }), + ); + }); expect( controller.state.tokenBalances[accountAddress as ChecksumAddress]?.[ @@ -6810,11 +7206,12 @@ describe('TokenBalancesController', () => { chainIds: [chainId], queryAllAccounts: true, }); - - // Verify warning was logged - expect(consoleWarnSpy).toHaveBeenCalledWith( - expect.stringContaining('Balance fetcher failed'), - ); + await waitFor(() => { + // Verify warning was logged + expect(consoleWarnSpy).toHaveBeenCalledWith( + expect.stringContaining('Balance fetcher failed'), + ); + }); // Verify detectTokens was called with forceRpc when fetcher fails expect(messengerCallSpy).toHaveBeenCalledWith( @@ -6874,18 +7271,19 @@ describe('TokenBalancesController', () => { chainIds: [chainId], queryAllAccounts: true, }); + await waitFor(() => { + const balances = + controller.state.tokenBalances[accountAddress as ChecksumAddress]?.[ + chainId + ]; + const token1Checksum = toChecksumHexAddress(token1) as ChecksumAddress; + const token2Checksum = toChecksumHexAddress(token2) as ChecksumAddress; - const balances = - controller.state.tokenBalances[accountAddress as ChecksumAddress]?.[ - chainId - ]; - const token1Checksum = toChecksumHexAddress(token1) as ChecksumAddress; - const token2Checksum = toChecksumHexAddress(token2) as ChecksumAddress; - - // token1 should be present with balance (success=true) - expect(balances?.[token1Checksum]).toBe(toHex(100)); - // token2 should NOT be present (success=false) - expect(balances?.[token2Checksum]).toBeUndefined(); + // token1 should be present with balance (success=true) + expect(balances?.[token1Checksum]).toBe(toHex(100)); + // token2 should NOT be present (success=false) + expect(balances?.[token2Checksum]).toBeUndefined(); + }); }); it('should skip balances with undefined value (line 963)', async () => { @@ -6933,18 +7331,19 @@ describe('TokenBalancesController', () => { chainIds: [chainId], queryAllAccounts: true, }); + await waitFor(() => { + const balances = + controller.state.tokenBalances[accountAddress as ChecksumAddress]?.[ + chainId + ]; + const token1Checksum = toChecksumHexAddress(token1) as ChecksumAddress; + const token2Checksum = toChecksumHexAddress(token2) as ChecksumAddress; - const balances = - controller.state.tokenBalances[accountAddress as ChecksumAddress]?.[ - chainId - ]; - const token1Checksum = toChecksumHexAddress(token1) as ChecksumAddress; - const token2Checksum = toChecksumHexAddress(token2) as ChecksumAddress; - - // token1 should be present with balance - expect(balances?.[token1Checksum]).toBe(toHex(100)); - // token2 should NOT be present (value=undefined) - expect(balances?.[token2Checksum]).toBeUndefined(); + // token1 should be present with balance + expect(balances?.[token1Checksum]).toBe(toHex(100)); + // token2 should NOT be present (value=undefined) + expect(balances?.[token2Checksum]).toBeUndefined(); + }); }); }); }); diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts index b5b1a67cb8e..d0e02823ed0 100644 --- a/packages/assets-controllers/src/TokenBalancesController.ts +++ b/packages/assets-controllers/src/TokenBalancesController.ts @@ -54,7 +54,7 @@ import { parseCaipChainId, } from '@metamask/utils'; import { produce } from 'immer'; -import { isEqual } from 'lodash'; +import { isEqual, union } from 'lodash'; import type { AccountTrackerControllerGetStateAction } from './AccountTrackerController'; import type { @@ -80,6 +80,7 @@ import type { TokensControllerState, TokensControllerStateChangeEvent, } from './TokensController'; +import { createBatchedHandler } from './utils/create-batch-handler'; export type ChainIdHex = Hex; export type ChecksumAddress = Hex; @@ -88,6 +89,40 @@ const CONTROLLER = 'TokenBalancesController' as const; const DEFAULT_INTERVAL_MS = 30_000; // 30 seconds const DEFAULT_WEBSOCKET_ACTIVE_POLLING_INTERVAL_MS = 300_000; // 5 minutes +/** Debounce wait (ms) for coalescing rapid updateBalances calls before flush */ +export const UPDATE_BALANCES_BATCH_MS = 50; + +export type UpdateBalancesOptions = { + chainIds?: ChainIdHex[]; + tokenAddresses?: string[]; + queryAllAccounts?: boolean; +}; + +/** + * Merges two UpdateBalancesOptions per queue-and-merge rules: + * - chainIds: union when both specify; undefined or empty means "all" so result becomes undefined. + * - tokenAddresses: union when both specify; undefined or empty means "all" so result becomes undefined. + * - queryAllAccounts: true if either is true. + * Exported for tests. + * + * @param a - First options (e.g. accumulated). + * @param b - Second options to merge in. + * @returns New merged options. + */ +export function mergeUpdateBalancesOptions( + a: UpdateBalancesOptions, + b: UpdateBalancesOptions, +): UpdateBalancesOptions { + // We will take + const chainIds = a.chainIds && b.chainIds && union(a.chainIds, b.chainIds); + const tokenAddresses = + a.tokenAddresses && + b.tokenAddresses && + union(a.tokenAddresses, b.tokenAddresses); + const queryAllAccounts = + Boolean(a.queryAllAccounts) || Boolean(b.queryAllAccounts); + return { chainIds, tokenAddresses, queryAllAccounts }; +} const metadata: StateMetadata = { tokenBalances: { includeInStateLogs: false, @@ -312,6 +347,8 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ pendingChanges: new Map(), }; + readonly #batchedUpdateBalances: (options: UpdateBalancesOptions) => void; + constructor({ messenger, interval = DEFAULT_INTERVAL_MS, @@ -373,6 +410,23 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ this.#isUnlocked = isUnlocked; this.#subscribeToControllers(); + this.#batchedUpdateBalances = createBatchedHandler( + (buffer) => + buffer.reduce( + (acc, opts) => mergeUpdateBalancesOptions(acc, opts), + {}, + ), + UPDATE_BALANCES_BATCH_MS, + (merged: UpdateBalancesOptions): Promise => + this.#executeUpdateBalances(merged).catch((error) => { + // With batched updateBalances, errors occur in the debounced flush. + // Log as polling failure so callers (e.g. interval polling) see consistent error reporting. + const chainIds = merged.chainIds ?? []; + const chainsLabel = + chainIds.length > 0 ? chainIds.join(', ') : 'unknown chains'; + console.warn(`Polling failed for chains ${chainsLabel}:`, error); + }), + ); messenger.registerMethodActionHandlers(this, MESSENGER_EXPOSED_METHODS); } @@ -685,15 +739,18 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ } } - async updateBalances({ + async updateBalances(options: UpdateBalancesOptions = {}): Promise { + if (!this.isActive) { + return; + } + this.#batchedUpdateBalances(options); + } + + async #executeUpdateBalances({ chainIds, tokenAddresses, queryAllAccounts = false, - }: { - chainIds?: ChainIdHex[]; - tokenAddresses?: string[]; - queryAllAccounts?: boolean; - } = {}): Promise { + }: UpdateBalancesOptions = {}): Promise { if (!this.isActive) { return; } diff --git a/packages/assets-controllers/src/__fixtures__/test-utils.ts b/packages/assets-controllers/src/__fixtures__/test-utils.ts new file mode 100644 index 00000000000..f97f266894f --- /dev/null +++ b/packages/assets-controllers/src/__fixtures__/test-utils.ts @@ -0,0 +1,35 @@ +type WaitForOptions = { + intervalMs?: number; + timeoutMs?: number; +}; + +/** + * Testing Utility - waitFor. Waits for and checks (at an interval) if assertion is reached. + * + * @param assertionFn - assertion function + * @param options - set wait for options + * @returns promise that you need to await in tests + */ +export const waitFor = async ( + assertionFn: () => void, + options: WaitForOptions = {}, +): Promise => { + const { intervalMs = 50, timeoutMs = 2000 } = options; + + const startTime = Date.now(); + + return new Promise((resolve, reject) => { + const intervalId = setInterval(() => { + try { + assertionFn(); + clearInterval(intervalId); + resolve(); + } catch { + if (Date.now() - startTime >= timeoutMs) { + clearInterval(intervalId); + reject(new Error(`waitFor: timeout reached after ${timeoutMs}ms`)); + } + } + }, intervalMs); + }); +}; diff --git a/packages/assets-controllers/src/utils/create-batch-handler.test.ts b/packages/assets-controllers/src/utils/create-batch-handler.test.ts new file mode 100644 index 00000000000..76b284e8d26 --- /dev/null +++ b/packages/assets-controllers/src/utils/create-batch-handler.test.ts @@ -0,0 +1,121 @@ +import { createBatchedHandler } from './create-batch-handler'; + +const TEST_BATCH_MS = 50; + +describe('createBatchedHandler', () => { + beforeEach(() => { + jest.useFakeTimers(); + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + const advanceAndFlush = async (): Promise => { + await jest.advanceTimersByTimeAsync(TEST_BATCH_MS); + }; + + const createNumberHandler = ( + onFlush = jest.fn().mockResolvedValue(undefined), + ): { capture: (n: number) => void; onFlush: jest.Mock } => { + const capture = createBatchedHandler( + (buffer) => buffer.reduce((sum, item) => sum + item, 0), + TEST_BATCH_MS, + onFlush, + ); + return { capture, onFlush }; + }; + + function createObjectHandler( + onFlush = jest.fn().mockResolvedValue(undefined), + ): { + capture: (item: { ids: number[] }) => void; + onFlush: jest.Mock; + } { + const capture = createBatchedHandler<{ ids: number[] }>( + (buffer) => ({ ids: buffer.flatMap((item) => item.ids) }), + TEST_BATCH_MS, + onFlush, + ); + return { capture, onFlush }; + } + + describe('buffering and aggregation', () => { + it.each([ + { + name: 'sums numbers and flushes once after debounce', + arrangeAct: (): ReturnType & { + act: () => void; + } => { + const ctx = createNumberHandler(); + return { + ...ctx, + act: (): void => { + ctx.capture(1); + ctx.capture(2); + ctx.capture(3); + }, + }; + }, + expectedCalls: 1, + expectedArg: 6, + }, + { + name: 'merges object arrays with custom aggregator', + arrangeAct: (): ReturnType & { + act: () => void; + } => { + const ctx = createObjectHandler(); + return { + ...ctx, + act: (): void => { + ctx.capture({ ids: [1] }); + ctx.capture({ ids: [2, 3] }); + }, + }; + }, + expectedCalls: 1, + expectedArg: { ids: [1, 2, 3] }, + }, + ])('$name', async ({ arrangeAct, expectedCalls, expectedArg }) => { + const { onFlush, act } = arrangeAct(); + expect(onFlush).not.toHaveBeenCalled(); + + act(); + expect(onFlush).not.toHaveBeenCalled(); + + await advanceAndFlush(); + + expect(onFlush).toHaveBeenCalledTimes(expectedCalls); + expect(onFlush).toHaveBeenCalledWith(expectedArg); + }); + }); + + describe('lifecycle and edge cases', () => { + it('does not call onFlush when capture was never invoked', async () => { + const onFlush = jest.fn().mockResolvedValue(undefined); + createBatchedHandler( + (b) => b.reduce((a, item) => a + item, 0), + TEST_BATCH_MS, + onFlush, + ); + + await advanceAndFlush(); + + expect(onFlush).not.toHaveBeenCalled(); + }); + + it('resets buffer after flush and can capture again', async () => { + const { capture, onFlush } = createNumberHandler(); + + capture(1); + await advanceAndFlush(); + expect(onFlush).toHaveBeenCalledWith(1); + + capture(2); + await advanceAndFlush(); + expect(onFlush).toHaveBeenCalledTimes(2); + expect(onFlush).toHaveBeenLastCalledWith(2); + }); + }); +}); diff --git a/packages/assets-controllers/src/utils/create-batch-handler.ts b/packages/assets-controllers/src/utils/create-batch-handler.ts new file mode 100644 index 00000000000..90b057642a2 --- /dev/null +++ b/packages/assets-controllers/src/utils/create-batch-handler.ts @@ -0,0 +1,39 @@ +import { debounce } from 'lodash'; + +/** + * Batched handler: buffers arguments, debounces flush, then runs an aggregator + * on the buffer and invokes onFlush with the result. Used to coalesce rapid + * updateBalances calls without dropping params. + * + * @param aggregatorFn - Reduces the buffered items into one. + * @param timeframeMs - Debounce wait before flushing. + * @param onFlush - Called with the aggregated result when flush runs. + * @returns Object with capture (push item and schedule flush) and cancel. + */ +export function createBatchedHandler( + aggregatorFn: (buffer: Item[]) => Item, + timeframeMs: number, + onFlush: (merged: Item) => void | Promise, +): (arg: Item) => void { + let eventBuffer: Item[] = []; + const flush = async (): Promise => { + if (eventBuffer.length === 0) { + return; + } + const merged = aggregatorFn(eventBuffer); + eventBuffer = []; + await onFlush(merged); + }; + const debouncedFlush = debounce(flush, timeframeMs, { + leading: false, + trailing: true, + }); + const capture = (arg: Item): void => { + eventBuffer.push(arg); + + // eslint-disable-next-line @typescript-eslint/no-floating-promises + debouncedFlush(); + }; + + return capture; +} From f163c70ab1c216529ad5bde058df47ba5f28feba Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Thu, 19 Mar 2026 12:19:42 +0000 Subject: [PATCH 2/5] feat(TokenBalancesController): implement batching for updateBalances calls Added functionality to the TokenBalancesController to batch rapid updateBalances requests, coalescing multiple calls into a single processed request. This enhancement reduces redundant balance fetches and improves performance. Updated changelog to reflect this new feature. --- packages/assets-controllers/CHANGELOG.md | 4 ++++ .../assets-controllers/src/TokenBalancesController.ts | 9 ++------- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/packages/assets-controllers/CHANGELOG.md b/packages/assets-controllers/CHANGELOG.md index fe89887921e..e7d3b595005 100644 --- a/packages/assets-controllers/CHANGELOG.md +++ b/packages/assets-controllers/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- `TokenBalancesController` batches rapid `updateBalances` calls: multiple requests within a short timeframe are coalesced and processed once, reducing redundant balance fetches ([#8246](https://github.com/MetaMask/core/pull/8246)) + ## [101.0.1] ### Changed diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts index faed304b688..640863c6cd8 100644 --- a/packages/assets-controllers/src/TokenBalancesController.ts +++ b/packages/assets-controllers/src/TokenBalancesController.ts @@ -419,12 +419,7 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ UPDATE_BALANCES_BATCH_MS, (merged: UpdateBalancesOptions): Promise => this.#executeUpdateBalances(merged).catch((error) => { - // With batched updateBalances, errors occur in the debounced flush. - // Log as polling failure so callers (e.g. interval polling) see consistent error reporting. - const chainIds = merged.chainIds ?? []; - const chainsLabel = - chainIds.length > 0 ? chainIds.join(', ') : 'unknown chains'; - console.warn(`Polling failed for chains ${chainsLabel}:`, error); + console.warn('Batched updated balances failed:', error); }), ); messenger.registerMethodActionHandlers(this, MESSENGER_EXPOSED_METHODS); @@ -722,7 +717,7 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ chainIds: ChainIdHex[]; queryAllAccounts?: boolean; }): Promise { - await this.updateBalances({ chainIds, queryAllAccounts }); + await this.#executeUpdateBalances({ chainIds, queryAllAccounts }); } updateChainPollingConfigs( From 5ebae46db95f898003b434746580bc2b8205380e Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Thu, 19 Mar 2026 14:03:50 +0000 Subject: [PATCH 3/5] fix(TokenBalancesController): improve batching logic for updateBalances Refined the batching logic in TokenBalancesController to handle empty buffers more effectively. The updated implementation now initializes the merge with the first buffer element when available, ensuring that the batching process remains robust and efficient. Additionally, enhanced error handling in the waitFor utility to provide more informative error messages upon timeout, improving debugging capabilities. --- .../src/TokenBalancesController.ts | 12 ++++++++---- .../src/__fixtures__/test-utils.ts | 12 ++++++++++-- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts index 640863c6cd8..5398f225793 100644 --- a/packages/assets-controllers/src/TokenBalancesController.ts +++ b/packages/assets-controllers/src/TokenBalancesController.ts @@ -412,10 +412,14 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ this.#subscribeToControllers(); this.#batchedUpdateBalances = createBatchedHandler( (buffer) => - buffer.reduce( - (acc, opts) => mergeUpdateBalancesOptions(acc, opts), - {}, - ), + buffer.length === 0 + ? {} + : buffer + .slice(1) + .reduce( + (acc, opts) => mergeUpdateBalancesOptions(acc, opts), + buffer[0], + ), UPDATE_BALANCES_BATCH_MS, (merged: UpdateBalancesOptions): Promise => this.#executeUpdateBalances(merged).catch((error) => { diff --git a/packages/assets-controllers/src/__fixtures__/test-utils.ts b/packages/assets-controllers/src/__fixtures__/test-utils.ts index f97f266894f..e57363787e8 100644 --- a/packages/assets-controllers/src/__fixtures__/test-utils.ts +++ b/packages/assets-controllers/src/__fixtures__/test-utils.ts @@ -19,15 +19,23 @@ export const waitFor = async ( const startTime = Date.now(); return new Promise((resolve, reject) => { + let lastError: unknown; const intervalId = setInterval(() => { try { assertionFn(); clearInterval(intervalId); resolve(); - } catch { + } catch (error) { + lastError = error; if (Date.now() - startTime >= timeoutMs) { clearInterval(intervalId); - reject(new Error(`waitFor: timeout reached after ${timeoutMs}ms`)); + const assertionDetail = + lastError instanceof Error ? lastError.message : String(lastError); + reject( + new Error( + `waitFor: timeout reached after ${timeoutMs}ms. Last assertion error: ${assertionDetail}`, + ), + ); } } }, intervalMs); From fd98a48b7c54031156f27bb915d1b67b57261245 Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Thu, 19 Mar 2026 14:18:24 +0000 Subject: [PATCH 4/5] refactor(TokenBalancesController, create-batch-handler): update batching to return Promises Modified the batching logic in TokenBalancesController to ensure that the #batchedUpdateBalances method returns a Promise, allowing for better error handling and synchronization. Updated the createBatchedHandler utility to return Promises for capture calls, enabling callers to await the completion of batch flushes and handle errors appropriately. Enhanced tests to verify the new asynchronous behavior and error propagation. --- .../src/TokenBalancesController.ts | 10 +-- .../src/utils/create-batch-handler.test.ts | 68 +++++++++++++++---- .../src/utils/create-batch-handler.ts | 40 ++++++++--- 3 files changed, 91 insertions(+), 27 deletions(-) diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts index 5398f225793..5e1232a9165 100644 --- a/packages/assets-controllers/src/TokenBalancesController.ts +++ b/packages/assets-controllers/src/TokenBalancesController.ts @@ -347,7 +347,9 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ pendingChanges: new Map(), }; - readonly #batchedUpdateBalances: (options: UpdateBalancesOptions) => void; + readonly #batchedUpdateBalances: ( + options: UpdateBalancesOptions, + ) => Promise; constructor({ messenger, @@ -422,9 +424,7 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ ), UPDATE_BALANCES_BATCH_MS, (merged: UpdateBalancesOptions): Promise => - this.#executeUpdateBalances(merged).catch((error) => { - console.warn('Batched updated balances failed:', error); - }), + this.#executeUpdateBalances(merged), ); messenger.registerMethodActionHandlers(this, MESSENGER_EXPOSED_METHODS); } @@ -742,7 +742,7 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ if (!this.isActive) { return; } - this.#batchedUpdateBalances(options); + await this.#batchedUpdateBalances(options); } async #executeUpdateBalances({ diff --git a/packages/assets-controllers/src/utils/create-batch-handler.test.ts b/packages/assets-controllers/src/utils/create-batch-handler.test.ts index 76b284e8d26..aee154c6875 100644 --- a/packages/assets-controllers/src/utils/create-batch-handler.test.ts +++ b/packages/assets-controllers/src/utils/create-batch-handler.test.ts @@ -17,7 +17,7 @@ describe('createBatchedHandler', () => { const createNumberHandler = ( onFlush = jest.fn().mockResolvedValue(undefined), - ): { capture: (n: number) => void; onFlush: jest.Mock } => { + ): { capture: (n: number) => Promise; onFlush: jest.Mock } => { const capture = createBatchedHandler( (buffer) => buffer.reduce((sum, item) => sum + item, 0), TEST_BATCH_MS, @@ -29,7 +29,7 @@ describe('createBatchedHandler', () => { function createObjectHandler( onFlush = jest.fn().mockResolvedValue(undefined), ): { - capture: (item: { ids: number[] }) => void; + capture: (item: { ids: number[] }) => Promise; onFlush: jest.Mock; } { const capture = createBatchedHandler<{ ids: number[] }>( @@ -45,15 +45,17 @@ describe('createBatchedHandler', () => { { name: 'sums numbers and flushes once after debounce', arrangeAct: (): ReturnType & { - act: () => void; + act: () => Promise; } => { const ctx = createNumberHandler(); return { ...ctx, - act: (): void => { - ctx.capture(1); - ctx.capture(2); - ctx.capture(3); + act: async (): Promise => { + await Promise.all([ + ctx.capture(1), + ctx.capture(2), + ctx.capture(3), + ]); }, }; }, @@ -63,14 +65,16 @@ describe('createBatchedHandler', () => { { name: 'merges object arrays with custom aggregator', arrangeAct: (): ReturnType & { - act: () => void; + act: () => Promise; } => { const ctx = createObjectHandler(); return { ...ctx, - act: (): void => { - ctx.capture({ ids: [1] }); - ctx.capture({ ids: [2, 3] }); + act: async (): Promise => { + await Promise.all([ + ctx.capture({ ids: [1] }), + ctx.capture({ ids: [2, 3] }), + ]); }, }; }, @@ -81,10 +85,11 @@ describe('createBatchedHandler', () => { const { onFlush, act } = arrangeAct(); expect(onFlush).not.toHaveBeenCalled(); - act(); + const promiseResult = act(); expect(onFlush).not.toHaveBeenCalled(); await advanceAndFlush(); + await promiseResult; expect(onFlush).toHaveBeenCalledTimes(expectedCalls); expect(onFlush).toHaveBeenCalledWith(expectedArg); @@ -108,14 +113,49 @@ describe('createBatchedHandler', () => { it('resets buffer after flush and can capture again', async () => { const { capture, onFlush } = createNumberHandler(); - capture(1); + const promise1 = capture(1); await advanceAndFlush(); + await promise1; expect(onFlush).toHaveBeenCalledWith(1); - capture(2); + const promise2 = capture(2); await advanceAndFlush(); + await promise2; expect(onFlush).toHaveBeenCalledTimes(2); expect(onFlush).toHaveBeenLastCalledWith(2); }); + + it('returns a Promise that resolves when the batch flush completes', async () => { + const { capture, onFlush } = createNumberHandler(); + + const promise = capture(1); + expect(onFlush).not.toHaveBeenCalled(); + + await advanceAndFlush(); + await promise; + expect(onFlush).toHaveBeenCalledWith(1); + }); + + it('rejects all callers in the same batch when onFlush throws', async () => { + const error = new Error('flush failed'); + const onFlush = jest.fn().mockRejectedValue(error); + const capture = createBatchedHandler( + (buffer) => buffer.reduce((sum, item) => sum + item, 0), + TEST_BATCH_MS, + onFlush, + ); + + const p1 = capture(1); + const p2 = capture(2); + const settled = Promise.allSettled([p1, p2]); + + await advanceAndFlush(); + const [r1, r2] = await settled; + + expect(r1.status).toBe('rejected'); + expect((r1 as PromiseRejectedResult).reason).toBe(error); + expect(r2.status).toBe('rejected'); + expect((r2 as PromiseRejectedResult).reason).toBe(error); + }); }); }); diff --git a/packages/assets-controllers/src/utils/create-batch-handler.ts b/packages/assets-controllers/src/utils/create-batch-handler.ts index 90b057642a2..33ce1bb7b46 100644 --- a/packages/assets-controllers/src/utils/create-batch-handler.ts +++ b/packages/assets-controllers/src/utils/create-batch-handler.ts @@ -1,38 +1,62 @@ import { debounce } from 'lodash'; +type PendingSettler = { + resolve: () => void; + reject: (reason: unknown) => void; +}; + /** * Batched handler: buffers arguments, debounces flush, then runs an aggregator * on the buffer and invokes onFlush with the result. Used to coalesce rapid * updateBalances calls without dropping params. * + * Each call to the returned function returns a Promise that resolves when the + * flush that includes that call completes, or rejects if onFlush throws, so + * callers can await or use .catch() for error handling. + * * @param aggregatorFn - Reduces the buffered items into one. * @param timeframeMs - Debounce wait before flushing. * @param onFlush - Called with the aggregated result when flush runs. - * @returns Object with capture (push item and schedule flush) and cancel. + * @returns Function that accepts an item, schedules a batched flush, and returns a Promise that settles when that batch completes. */ export function createBatchedHandler( aggregatorFn: (buffer: Item[]) => Item, timeframeMs: number, onFlush: (merged: Item) => void | Promise, -): (arg: Item) => void { +): (arg: Item) => Promise { let eventBuffer: Item[] = []; + let pendingSettlers: PendingSettler[] = []; + const flush = async (): Promise => { if (eventBuffer.length === 0) { return; } - const merged = aggregatorFn(eventBuffer); + const buffer = eventBuffer; + const settlers = pendingSettlers; eventBuffer = []; - await onFlush(merged); + pendingSettlers = []; + + const merged = aggregatorFn(buffer); + try { + await onFlush(merged); + settlers.forEach((settler) => settler.resolve()); + } catch (error) { + settlers.forEach((settler) => settler.reject(error)); + } }; + const debouncedFlush = debounce(flush, timeframeMs, { leading: false, trailing: true, }); - const capture = (arg: Item): void => { - eventBuffer.push(arg); - // eslint-disable-next-line @typescript-eslint/no-floating-promises - debouncedFlush(); + const capture = (arg: Item): Promise => { + return new Promise((resolve, reject) => { + eventBuffer.push(arg); + pendingSettlers.push({ resolve, reject }); + // eslint-disable-next-line @typescript-eslint/no-floating-promises -- Rejections are forwarded to capture() callers via pendingSettlers. + debouncedFlush(); + }); }; return capture; From 2cb9ef9c5b7434333d3109ee13d92eba43966e58 Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Thu, 19 Mar 2026 14:53:32 +0000 Subject: [PATCH 5/5] refactor(TokenBalancesController, create-batch-handler): restore subscription logic and enhance error handling Reintroduced the subscription logic in the TokenBalancesController to ensure proper updates during balance changes. Additionally, improved the createBatchedHandler utility to handle errors more effectively by ensuring that all promises are settled correctly, enhancing the robustness of the batching mechanism. Updated tests to validate the new error handling behavior and ensure all scenarios are covered. --- .../src/TokenBalancesController.ts | 4 +- .../src/utils/create-batch-handler.test.ts | 42 +++++++++++++++---- .../src/utils/create-batch-handler.ts | 2 +- 3 files changed, 36 insertions(+), 12 deletions(-) diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts index 5e1232a9165..9a5562795ec 100644 --- a/packages/assets-controllers/src/TokenBalancesController.ts +++ b/packages/assets-controllers/src/TokenBalancesController.ts @@ -113,7 +113,6 @@ export function mergeUpdateBalancesOptions( a: UpdateBalancesOptions, b: UpdateBalancesOptions, ): UpdateBalancesOptions { - // We will take const chainIds = a.chainIds && b.chainIds && union(a.chainIds, b.chainIds); const tokenAddresses = a.tokenAddresses && @@ -411,7 +410,6 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ const { isUnlocked } = this.messenger.call('KeyringController:getState'); this.#isUnlocked = isUnlocked; - this.#subscribeToControllers(); this.#batchedUpdateBalances = createBatchedHandler( (buffer) => buffer.length === 0 @@ -426,6 +424,8 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ (merged: UpdateBalancesOptions): Promise => this.#executeUpdateBalances(merged), ); + + this.#subscribeToControllers(); messenger.registerMethodActionHandlers(this, MESSENGER_EXPOSED_METHODS); } diff --git a/packages/assets-controllers/src/utils/create-batch-handler.test.ts b/packages/assets-controllers/src/utils/create-batch-handler.test.ts index aee154c6875..d3ac6601bd2 100644 --- a/packages/assets-controllers/src/utils/create-batch-handler.test.ts +++ b/packages/assets-controllers/src/utils/create-batch-handler.test.ts @@ -136,6 +136,23 @@ describe('createBatchedHandler', () => { expect(onFlush).toHaveBeenCalledWith(1); }); + const actAssertRejected = async ( + capture: (n: number) => Promise, + expectedError: unknown, + ): Promise => { + const p1 = capture(1); + const p2 = capture(2); + const settled = Promise.allSettled([p1, p2]); + + await advanceAndFlush(); + const [r1, r2] = await settled; + + expect(r1.status).toBe('rejected'); + expect((r1 as PromiseRejectedResult).reason).toBe(expectedError); + expect(r2.status).toBe('rejected'); + expect((r2 as PromiseRejectedResult).reason).toBe(expectedError); + }; + it('rejects all callers in the same batch when onFlush throws', async () => { const error = new Error('flush failed'); const onFlush = jest.fn().mockRejectedValue(error); @@ -145,17 +162,24 @@ describe('createBatchedHandler', () => { onFlush, ); - const p1 = capture(1); - const p2 = capture(2); - const settled = Promise.allSettled([p1, p2]); + await actAssertRejected(capture, error); + expect(onFlush).toHaveBeenCalled(); + }); - await advanceAndFlush(); - const [r1, r2] = await settled; + it('rejects all callers in the same batch when aggregatorFn throws', async () => { + const error = new Error('aggregation failed'); + const aggregatorFn = jest.fn(() => { + throw error; + }); + const onFlush = jest.fn().mockResolvedValue(undefined); + const capture = createBatchedHandler( + aggregatorFn, + TEST_BATCH_MS, + onFlush, + ); - expect(r1.status).toBe('rejected'); - expect((r1 as PromiseRejectedResult).reason).toBe(error); - expect(r2.status).toBe('rejected'); - expect((r2 as PromiseRejectedResult).reason).toBe(error); + await actAssertRejected(capture, error); + expect(onFlush).not.toHaveBeenCalled(); }); }); }); diff --git a/packages/assets-controllers/src/utils/create-batch-handler.ts b/packages/assets-controllers/src/utils/create-batch-handler.ts index 33ce1bb7b46..4d202b9575d 100644 --- a/packages/assets-controllers/src/utils/create-batch-handler.ts +++ b/packages/assets-controllers/src/utils/create-batch-handler.ts @@ -36,8 +36,8 @@ export function createBatchedHandler( eventBuffer = []; pendingSettlers = []; - const merged = aggregatorFn(buffer); try { + const merged = aggregatorFn(buffer); await onFlush(merged); settlers.forEach((settler) => settler.resolve()); } catch (error) {