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.test.ts b/packages/assets-controllers/src/TokenBalancesController.test.ts index c4ce3292b55..1d0d9a4f7ac 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 () => { @@ -5122,20 +5477,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 () => { @@ -5165,10 +5519,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 () => { @@ -5245,23 +5598,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 () => { @@ -5323,18 +5674,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 () => { @@ -5390,18 +5740,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 () => { @@ -5451,19 +5800,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 () => { @@ -5521,17 +5869,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(); }); @@ -5607,26 +5954,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', () => { @@ -5663,6 +6009,7 @@ describe('TokenBalancesController', () => { mockAPIAccountsAPIMultichainAccountBalancesCamelCase(accountAddress); const account = createMockInternalAccount({ address: accountAddress }); + const tokenAddress = '0x2791bca1f2de4661ed88a30c99a7a9449aa84174'; // USDC from mock response const { controller } = setupController({ config: { @@ -5670,6 +6017,17 @@ describe('TokenBalancesController', () => { allowExternalServices: () => true, }, listAccounts: [account], + tokens: { + allTokens: { + [chainId]: { + [accountAddress]: [ + { address: tokenAddress, symbol: 'USDC', decimals: 6 }, + ], + }, + }, + allDetectedTokens: {}, + allIgnoredTokens: {}, + }, }); return { @@ -5686,10 +6044,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); }); @@ -5850,6 +6210,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: { @@ -6039,6 +6406,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({ @@ -6061,119 +6435,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(); + } }); }); @@ -6264,12 +6648,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 @@ -6286,9 +6682,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(); }); @@ -6466,15 +6862,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 }); @@ -6513,15 +6908,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 () => { @@ -6673,14 +7068,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(); @@ -6730,15 +7124,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(); }); @@ -6808,18 +7203,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]?.[ @@ -6872,11 +7268,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( @@ -6936,18 +7333,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 () => { @@ -6995,18 +7393,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 c0e55cbd690..9a5562795ec 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,39 @@ 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 { + 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 +346,10 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ pendingChanges: new Map(), }; + readonly #batchedUpdateBalances: ( + options: UpdateBalancesOptions, + ) => Promise; + constructor({ messenger, interval = DEFAULT_INTERVAL_MS, @@ -372,6 +410,21 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ const { isUnlocked } = this.messenger.call('KeyringController:getState'); this.#isUnlocked = isUnlocked; + this.#batchedUpdateBalances = createBatchedHandler( + (buffer) => + buffer.length === 0 + ? {} + : buffer + .slice(1) + .reduce( + (acc, opts) => mergeUpdateBalancesOptions(acc, opts), + buffer[0], + ), + UPDATE_BALANCES_BATCH_MS, + (merged: UpdateBalancesOptions): Promise => + this.#executeUpdateBalances(merged), + ); + this.#subscribeToControllers(); messenger.registerMethodActionHandlers(this, MESSENGER_EXPOSED_METHODS); } @@ -668,7 +721,7 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ chainIds: ChainIdHex[]; queryAllAccounts?: boolean; }): Promise { - await this.updateBalances({ chainIds, queryAllAccounts }); + await this.#executeUpdateBalances({ chainIds, queryAllAccounts }); } updateChainPollingConfigs( @@ -685,15 +738,18 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ } } - async updateBalances({ + async updateBalances(options: UpdateBalancesOptions = {}): Promise { + if (!this.isActive) { + return; + } + await 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..e57363787e8 --- /dev/null +++ b/packages/assets-controllers/src/__fixtures__/test-utils.ts @@ -0,0 +1,43 @@ +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) => { + let lastError: unknown; + const intervalId = setInterval(() => { + try { + assertionFn(); + clearInterval(intervalId); + resolve(); + } catch (error) { + lastError = error; + if (Date.now() - startTime >= timeoutMs) { + clearInterval(intervalId); + const assertionDetail = + lastError instanceof Error ? lastError.message : String(lastError); + reject( + new Error( + `waitFor: timeout reached after ${timeoutMs}ms. Last assertion error: ${assertionDetail}`, + ), + ); + } + } + }, 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..d3ac6601bd2 --- /dev/null +++ b/packages/assets-controllers/src/utils/create-batch-handler.test.ts @@ -0,0 +1,185 @@ +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) => Promise; 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[] }) => Promise; + 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: () => Promise; + } => { + const ctx = createNumberHandler(); + return { + ...ctx, + act: async (): Promise => { + await Promise.all([ + ctx.capture(1), + ctx.capture(2), + ctx.capture(3), + ]); + }, + }; + }, + expectedCalls: 1, + expectedArg: 6, + }, + { + name: 'merges object arrays with custom aggregator', + arrangeAct: (): ReturnType & { + act: () => Promise; + } => { + const ctx = createObjectHandler(); + return { + ...ctx, + act: async (): Promise => { + await Promise.all([ + 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(); + + const promiseResult = act(); + expect(onFlush).not.toHaveBeenCalled(); + + await advanceAndFlush(); + await promiseResult; + + 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(); + + const promise1 = capture(1); + await advanceAndFlush(); + await promise1; + expect(onFlush).toHaveBeenCalledWith(1); + + 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); + }); + + 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); + const capture = createBatchedHandler( + (buffer) => buffer.reduce((sum, item) => sum + item, 0), + TEST_BATCH_MS, + onFlush, + ); + + await actAssertRejected(capture, error); + expect(onFlush).toHaveBeenCalled(); + }); + + 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, + ); + + 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 new file mode 100644 index 00000000000..4d202b9575d --- /dev/null +++ b/packages/assets-controllers/src/utils/create-batch-handler.ts @@ -0,0 +1,63 @@ +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 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) => Promise { + let eventBuffer: Item[] = []; + let pendingSettlers: PendingSettler[] = []; + + const flush = async (): Promise => { + if (eventBuffer.length === 0) { + return; + } + const buffer = eventBuffer; + const settlers = pendingSettlers; + eventBuffer = []; + pendingSettlers = []; + + try { + const merged = aggregatorFn(buffer); + 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): 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; +}