From 4394679c0a6e25cb6312a92dc7c1663335edaf42 Mon Sep 17 00:00:00 2001 From: Brian Bergeron Date: Tue, 22 Oct 2024 16:06:49 -0700 Subject: [PATCH] merge main, tests --- packages/assets-controllers/package.json | 1 + .../src/TokenBalancesController.test.ts | 406 ++++-------------- .../src/TokenBalancesController.ts | 23 +- yarn.lock | 1 + 4 files changed, 104 insertions(+), 327 deletions(-) diff --git a/packages/assets-controllers/package.json b/packages/assets-controllers/package.json index e69368eccd..a70e500b3b 100644 --- a/packages/assets-controllers/package.json +++ b/packages/assets-controllers/package.json @@ -48,6 +48,7 @@ }, "dependencies": { "@ethereumjs/util": "^8.1.0", + "@ethersproject/abi": "^5.7.0", "@ethersproject/address": "^5.7.0", "@ethersproject/bignumber": "^5.7.0", "@ethersproject/contracts": "^5.7.0", diff --git a/packages/assets-controllers/src/TokenBalancesController.test.ts b/packages/assets-controllers/src/TokenBalancesController.test.ts index 751bfa4ad4..a2aee0b41c 100644 --- a/packages/assets-controllers/src/TokenBalancesController.test.ts +++ b/packages/assets-controllers/src/TokenBalancesController.test.ts @@ -1,26 +1,16 @@ import { ControllerMessenger } from '@metamask/base-controller'; -import { NetworkType, toHex } from '@metamask/controller-utils'; -import type { InternalAccount } from '@metamask/keyring-api'; import BN from 'bn.js'; import { useFakeTimers } from 'sinon'; -import { advanceTime, flushPromises } from '../../../tests/helpers'; -import { createMockInternalAccount } from '../../accounts-controller/src/tests/mocks'; +import { advanceTime } from '../../../tests/helpers'; import * as multicall from './multicall'; -import { multicallOrFallback } from './multicall'; import type { AllowedActions, AllowedEvents, TokenBalancesControllerActions, TokenBalancesControllerEvents, - TokenBalancesControllerMessenger, } from './TokenBalancesController'; import { TokenBalancesController } from './TokenBalancesController'; -import type { Token } from './TokenRatesController'; -import { - getDefaultTokensState, - type TokensControllerState, -} from './TokensController'; const setupController = ({ config, @@ -41,17 +31,14 @@ const setupController = ({ controllerMessenger.registerActionHandler( 'NetworkController:getNetworkClientById', jest.fn().mockImplementation((networkClientId) => { - const chainId = - networkClientId === 'mainnet' - ? '0x1' - : networkClientId === 'sepolia' - ? '0xaa36a7' - : undefined; - - if (!chainId) { + let chainId; + if (networkClientId === 'mainnet') { + chainId = '0x1'; + } else if (networkClientId === 'sepolia') { + chainId = '0xaa36a7'; + } else { throw new Error('unknown networkClientId'); } - return { configuration: { chainId }, provider: jest.fn(), @@ -87,7 +74,10 @@ describe('TokenBalancesController', () => { const interval = 10; const controller = setupController({ config: { interval } }); - controller.startPollingByNetworkClientId('mainnet'); + controller.startPolling({ + networkClientId: 'mainnet', + tokensPerAccount: {}, + }); await advanceTime({ clock, duration: 1 }); expect(pollSpy).toHaveBeenCalled(); @@ -97,334 +87,112 @@ describe('TokenBalancesController', () => { expect(pollSpy).toHaveBeenCalledTimes(2); }); - it('should update balances', async () => { + it('should update balances on success', async () => { const controller = setupController({}); + expect(controller.state.tokenBalances).toStrictEqual({}); jest.spyOn(multicall, 'multicallOrFallback').mockResolvedValue([ { success: true, - value: new BN(1), + value: new BN(2), }, ]); const accountAddresss = '0x0000000000000000000000000000000000000000'; const tokenAddress = '0x0000000000000000000000000000000000000001'; - await controller._executePoll('mainnet', { - [accountAddresss]: [tokenAddress], + await controller._executePoll({ + networkClientId: 'mainnet', + tokensPerAccount: { + [accountAddresss]: [tokenAddress], + }, }); expect(controller.state.tokenBalances).toStrictEqual({ [accountAddresss]: { '0x1': { - [tokenAddress]: '0x1', + [tokenAddress]: '0x2', }, }, }); }); - // TODO: Consider more tests - - // it('should not update balances if disabled', async () => { - // const address = '0x86fa049857e0209aa7d9e616f7eb3b3b78ecfdb0'; - // const { controller } = setupController({ - // config: { - // disabled: true, - // tokens: [{ address, decimals: 18, symbol: 'EOS', aggregators: [] }], - // interval: 10, - // }, - // mock: { - // selectedAccount: createMockInternalAccount({ address: '0x1234' }), - // getBalanceOf: new BN(1), - // }, - // }); - - // await controller.updateBalances(); - - // expect(controller.state.contractBalances).toStrictEqual({}); - // }); - - // it('should update balances if controller is manually enabled', async () => { - // const address = '0x86fa049857e0209aa7d9e616f7eb3b3b78ecfdb0'; - // const { controller } = setupController({ - // config: { - // disabled: true, - // tokens: [{ address, decimals: 18, symbol: 'EOS', aggregators: [] }], - // interval: 10, - // }, - // mock: { - // selectedAccount: createMockInternalAccount({ address: '0x1234' }), - // getBalanceOf: new BN(1), - // }, - // }); - - // await controller.updateBalances(); - - // expect(controller.state.contractBalances).toStrictEqual({}); - - // controller.enable(); - // await controller.updateBalances(); - - // expect(controller.state.contractBalances).toStrictEqual({ - // '0x86fa049857e0209aa7d9e616f7eb3b3b78ecfdb0': toHex(new BN(1)), - // }); - // }); - - // it('should not update balances if controller is manually disabled', async () => { - // const address = '0x86fa049857e0209aa7d9e616f7eb3b3b78ecfdb0'; - // const { controller } = setupController({ - // config: { - // disabled: false, - // tokens: [{ address, decimals: 18, symbol: 'EOS', aggregators: [] }], - // interval: 10, - // }, - // mock: { - // selectedAccount: createMockInternalAccount({ address: '0x1234' }), - // getBalanceOf: new BN(1), - // }, - // }); - - // await controller.updateBalances(); - - // expect(controller.state.contractBalances).toStrictEqual({ - // '0x86fa049857e0209aa7d9e616f7eb3b3b78ecfdb0': toHex(new BN(1)), - // }); - - // controller.disable(); - // await controller.updateBalances(); - - // expect(controller.state.contractBalances).toStrictEqual({ - // '0x86fa049857e0209aa7d9e616f7eb3b3b78ecfdb0': toHex(new BN(1)), - // }); - // }); - - // it('should update balances if tokens change and controller is manually enabled', async () => { - // const address = '0x86fa049857e0209aa7d9e616f7eb3b3b78ecfdb0'; - // const { controller, triggerTokensStateChange } = setupController({ - // config: { - // disabled: true, - // tokens: [{ address, decimals: 18, symbol: 'EOS', aggregators: [] }], - // interval: 10, - // }, - // mock: { - // selectedAccount: createMockInternalAccount({ address: '0x1234' }), - // getBalanceOf: new BN(1), - // }, - // }); - - // await controller.updateBalances(); - - // expect(controller.state.contractBalances).toStrictEqual({}); - - // controller.enable(); - // await triggerTokensStateChange({ - // ...getDefaultTokensState(), - // tokens: [ - // { - // address: '0x00', - // symbol: 'FOO', - // decimals: 18, - // }, - // ], - // }); - - // expect(controller.state.contractBalances).toStrictEqual({ - // '0x00': toHex(new BN(1)), - // }); - // }); - - // it('should not update balances if tokens change and controller is manually disabled', async () => { - // const address = '0x86fa049857e0209aa7d9e616f7eb3b3b78ecfdb0'; - // const { controller, triggerTokensStateChange } = setupController({ - // config: { - // disabled: false, - // tokens: [{ address, decimals: 18, symbol: 'EOS', aggregators: [] }], - // interval: 10, - // }, - // mock: { - // selectedAccount: createMockInternalAccount({ address: '0x1234' }), - // getBalanceOf: new BN(1), - // }, - // }); - - // await controller.updateBalances(); - - // expect(controller.state.contractBalances).toStrictEqual({ - // '0x86fa049857e0209aa7d9e616f7eb3b3b78ecfdb0': toHex(new BN(1)), - // }); - - // controller.disable(); - // await triggerTokensStateChange({ - // ...getDefaultTokensState(), - // tokens: [ - // { - // address: '0x00', - // symbol: 'FOO', - // decimals: 18, - // }, - // ], - // }); - - // expect(controller.state.contractBalances).toStrictEqual({ - // '0x86fa049857e0209aa7d9e616f7eb3b3b78ecfdb0': toHex(new BN(1)), - // }); - // }); - - // it('should clear previous interval', async () => { - // const { controller } = setupController({ - // config: { - // interval: 1337, - // }, - // mock: { - // selectedAccount: createMockInternalAccount({ address: '0x1234' }), - // getBalanceOf: new BN(1), - // }, - // }); - - // const mockClearTimeout = jest.spyOn(global, 'clearTimeout'); - - // await controller.poll(1338); - - // jest.advanceTimersByTime(1339); - - // expect(mockClearTimeout).toHaveBeenCalled(); - // }); - - // it('should update all balances', async () => { - // const selectedAddress = '0x0000000000000000000000000000000000000001'; - // const address = '0x86fa049857e0209aa7d9e616f7eb3b3b78ecfdb0'; - // const tokens: Token[] = [ - // { - // address, - // decimals: 18, - // symbol: 'EOS', - // aggregators: [], - // }, - // ]; - // const { controller } = setupController({ - // config: { - // interval: 1337, - // tokens, - // }, - // mock: { - // selectedAccount: createMockInternalAccount({ - // address: selectedAddress, - // }), - // getBalanceOf: new BN(1), - // }, - // }); - - // expect(controller.state.contractBalances).toStrictEqual({}); - - // await controller.updateBalances(); - - // expect(tokens[0].hasBalanceError).toBe(false); - // expect(Object.keys(controller.state.contractBalances)).toContain(address); - // expect(controller.state.contractBalances[address]).not.toBe(toHex(0)); - // }); - - // it('should handle `getERC20BalanceOf` error case', async () => { - // const errorMsg = 'Failed to get balance'; - // const address = '0x86fa049857e0209aa7d9e616f7eb3b3b78ecfdb0'; - // const tokens: Token[] = [ - // { - // address, - // decimals: 18, - // symbol: 'EOS', - // aggregators: [], - // }, - // ]; - - // const { controller, mockGetERC20BalanceOf } = setupController({ - // config: { - // interval: 1337, - // tokens, - // }, - // mock: { - // selectedAccount: createMockInternalAccount({ - // address, - // }), - // }, - // }); - - // // @ts-expect-error Testing error case - // mockGetERC20BalanceOf.mockReturnValueOnce(new Error(errorMsg)); - - // expect(controller.state.contractBalances).toStrictEqual({}); - - // await controller.updateBalances(); + it('should update balances when they change', async () => { + const controller = setupController({}); + expect(controller.state.tokenBalances).toStrictEqual({}); - // expect(tokens[0].hasBalanceError).toBe(true); - // expect(controller.state.contractBalances[address]).toBe(toHex(0)); + const accountAddresss = '0x0000000000000000000000000000000000000000'; + const tokenAddress = '0x0000000000000000000000000000000000000001'; - // mockGetERC20BalanceOf.mockReturnValueOnce(new BN(1)); - // await controller.updateBalances(); + for (let i = 0; i < 10; i++) { + jest.spyOn(multicall, 'multicallOrFallback').mockResolvedValue([ + { + success: true, + value: new BN(i), + }, + ]); - // expect(tokens[0].hasBalanceError).toBe(false); - // expect(Object.keys(controller.state.contractBalances)).toContain(address); - // expect(controller.state.contractBalances[address]).not.toBe(0); - // }); + await controller._executePoll({ + networkClientId: 'mainnet', + tokensPerAccount: { + [accountAddresss]: [tokenAddress], + }, + }); - // it('should update balances when tokens change', async () => { - // const { controller, triggerTokensStateChange } = setupController({ - // config: { - // interval: 1337, - // }, - // mock: { - // selectedAccount: createMockInternalAccount({ - // address: '0x1234', - // }), - // getBalanceOf: new BN(1), - // }, - // }); + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddresss]: { + '0x1': { + [tokenAddress]: `0x${i}`, + }, + }, + }); + } + }); - // const updateBalancesSpy = jest.spyOn(controller, 'updateBalances'); + it('should not update balances on failure', async () => { + const controller = setupController({}); - // await triggerTokensStateChange({ - // ...getDefaultTokensState(), - // tokens: [ - // { - // address: '0x00', - // symbol: 'FOO', - // decimals: 18, - // }, - // ], - // }); + const accountAddresss = '0x0000000000000000000000000000000000000000'; + const tokenAddress = '0x0000000000000000000000000000000000000001'; - // expect(updateBalancesSpy).toHaveBeenCalled(); - // }); + // Initial successfull call + jest.spyOn(multicall, 'multicallOrFallback').mockResolvedValue([ + { + success: true, + value: new BN(2), + }, + ]); - // it('should update token balances when detected tokens are added', async () => { - // const { controller, triggerTokensStateChange } = setupController({ - // config: { - // interval: 1337, - // }, - // mock: { - // selectedAccount: createMockInternalAccount({ - // address: '0x1234', - // }), - // getBalanceOf: new BN(1), - // }, - // }); + await controller._executePoll({ + networkClientId: 'mainnet', + tokensPerAccount: { + [accountAddresss]: [tokenAddress], + }, + }); - // expect(controller.state.contractBalances).toStrictEqual({}); + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddresss]: { + '0x1': { + [tokenAddress]: '0x2', + }, + }, + }); - // await triggerTokensStateChange({ - // ...getDefaultTokensState(), - // detectedTokens: [ - // { - // address: '0x02', - // decimals: 18, - // image: undefined, - // symbol: 'bar', - // isERC721: false, - // }, - // ], - // tokens: [], - // }); + // Failed call + jest.spyOn(multicall, 'multicallOrFallback').mockResolvedValue([ + { + success: false, + value: '', + }, + ]); - // expect(controller.state.contractBalances).toStrictEqual({ - // '0x02': toHex(new BN(1)), - // }); - // }); + // State should not change + expect(controller.state.tokenBalances).toStrictEqual({ + [accountAddresss]: { + '0x1': { + [tokenAddress]: '0x2', + }, + }, + }); + }); }); diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts index b18492aa72..5fba7d61bc 100644 --- a/packages/assets-controllers/src/TokenBalancesController.ts +++ b/packages/assets-controllers/src/TokenBalancesController.ts @@ -91,11 +91,17 @@ export function getDefaultTokenBalancesState(): TokenBalancesControllerState { }; } +/** The input to start polling for the {@link TokenBalancesController} */ +export type TokenBalancesPollingInput = { + networkClientId: NetworkClientId; + tokensPerAccount: Record; +}; + /** * Controller that passively polls on a set interval token balances * for tokens stored in the TokensController */ -export class TokenBalancesController extends StaticIntervalPollingController< +export class TokenBalancesController extends StaticIntervalPollingController()< typeof controllerName, TokenBalancesControllerState, TokenBalancesControllerMessenger @@ -128,13 +134,14 @@ export class TokenBalancesController extends StaticIntervalPollingController< /** * Polls for erc20 token balances. - * @param networkClientId - The network client id to poll with. - * @param options - A mapping from account addresses to token addresses to poll. + * @param input - The input for the poll. + * @param input.networkClientId - The network client id to poll with. + * @param input.tokensPerAccount - A mapping from account addresses to token addresses to poll. */ - async _executePoll( - networkClientId: NetworkClientId, - options: Record, - ): Promise { + async _executePoll({ + networkClientId, + tokensPerAccount, + }: TokenBalancesPollingInput): Promise { const networkClient = this.messagingSystem.call( `NetworkController:getNetworkClientById`, networkClientId, @@ -143,7 +150,7 @@ export class TokenBalancesController extends StaticIntervalPollingController< const { chainId } = networkClient.configuration; const provider = new Web3Provider(networkClient.provider); - const accountTokenPairs = Object.entries(options).flatMap( + const accountTokenPairs = Object.entries(tokensPerAccount).flatMap( ([accountAddress, tokenAddresses]) => tokenAddresses.map((tokenAddress) => ({ accountAddress: accountAddress as Hex, diff --git a/yarn.lock b/yarn.lock index 6106761bbb..f9100b9f0f 100644 --- a/yarn.lock +++ b/yarn.lock @@ -2128,6 +2128,7 @@ __metadata: resolution: "@metamask/assets-controllers@workspace:packages/assets-controllers" dependencies: "@ethereumjs/util": "npm:^8.1.0" + "@ethersproject/abi": "npm:^5.7.0" "@ethersproject/address": "npm:^5.7.0" "@ethersproject/bignumber": "npm:^5.7.0" "@ethersproject/contracts": "npm:^5.7.0"