From 3174da02caabef1862ea364aa2a8ff2a86873cef Mon Sep 17 00:00:00 2001 From: Nicholas Smith Date: Wed, 30 Oct 2024 23:22:31 -0500 Subject: [PATCH] chore: add unit tests and clean up --- .../src/AccountTrackerController.test.ts | 325 +++++++++++++++++- .../src/AccountTrackerController.ts | 15 +- .../src/AssetsContractController.test.ts | 136 ++++++++ ...tractControllerWithNetworkClientId.test.ts | 75 ++++ 4 files changed, 544 insertions(+), 7 deletions(-) diff --git a/packages/assets-controllers/src/AccountTrackerController.test.ts b/packages/assets-controllers/src/AccountTrackerController.test.ts index bcd7f9f55c7..bd27eecb87e 100644 --- a/packages/assets-controllers/src/AccountTrackerController.test.ts +++ b/packages/assets-controllers/src/AccountTrackerController.test.ts @@ -248,6 +248,123 @@ describe('AccountTrackerController', () => { }, ); }); + + it('should update staked balance when includeStakedAssets is enabled', async () => { + mockedQuery + .mockReturnValueOnce(Promise.resolve('0x10')) + .mockReturnValueOnce(Promise.resolve('0x11')); + + await withController( + { + options: { + includeStakedAssets: true, + getStakedBalanceForChain: jest.fn().mockResolvedValue('0x1'), + }, + isMultiAccountBalancesEnabled: false, + selectedAccount: ACCOUNT_1, + listAccounts: [ACCOUNT_1, ACCOUNT_2], + }, + async ({ controller }) => { + await controller.refresh(); + + expect(controller.state).toStrictEqual({ + accounts: { + [CHECKSUM_ADDRESS_1]: { balance: '0x10', stakedBalance: '0x1' }, + [CHECKSUM_ADDRESS_2]: { balance: '0x0' }, + }, + accountsByChainId: { + '0x1': { + [CHECKSUM_ADDRESS_1]: { + balance: '0x10', + stakedBalance: '0x1', + }, + [CHECKSUM_ADDRESS_2]: { + balance: '0x0', + }, + }, + }, + }); + }, + ); + }); + + it('should not update staked balance when includeStakedAssets is disabled', async () => { + mockedQuery + .mockReturnValueOnce(Promise.resolve('0x13')) + .mockReturnValueOnce(Promise.resolve('0x14')); + + await withController( + { + options: { + includeStakedAssets: false, + getStakedBalanceForChain: jest.fn().mockResolvedValue('0x1'), + }, + isMultiAccountBalancesEnabled: false, + selectedAccount: ACCOUNT_1, + listAccounts: [ACCOUNT_1, ACCOUNT_2], + }, + async ({ controller }) => { + await controller.refresh(); + + expect(controller.state).toStrictEqual({ + accounts: { + [CHECKSUM_ADDRESS_1]: { balance: '0x13' }, + [CHECKSUM_ADDRESS_2]: { balance: '0x0' }, + }, + accountsByChainId: { + '0x1': { + [CHECKSUM_ADDRESS_1]: { + balance: '0x13', + }, + [CHECKSUM_ADDRESS_2]: { + balance: '0x0', + }, + }, + }, + }); + }, + ); + }); + + it('should update staked balance when includeStakedAssets and multi-account is enabled', async () => { + mockedQuery + .mockReturnValueOnce(Promise.resolve('0x11')) + .mockReturnValueOnce(Promise.resolve('0x12')); + + await withController( + { + options: { + includeStakedAssets: true, + getStakedBalanceForChain: jest.fn().mockResolvedValue('0x1'), + }, + isMultiAccountBalancesEnabled: true, + selectedAccount: ACCOUNT_1, + listAccounts: [ACCOUNT_1, ACCOUNT_2], + }, + async ({ controller }) => { + await controller.refresh(); + + expect(controller.state).toStrictEqual({ + accounts: { + [CHECKSUM_ADDRESS_1]: { balance: '0x11', stakedBalance: '0x1' }, + [CHECKSUM_ADDRESS_2]: { balance: '0x12', stakedBalance: '0x1' }, + }, + accountsByChainId: { + '0x1': { + [CHECKSUM_ADDRESS_1]: { + balance: '0x11', + stakedBalance: '0x1', + }, + [CHECKSUM_ADDRESS_2]: { + balance: '0x12', + stakedBalance: '0x1', + }, + }, + }, + }); + }, + ); + }); }); describe('with networkClientId', () => { @@ -438,6 +555,185 @@ describe('AccountTrackerController', () => { }, ); }); + + it('should update staked balance when includeStakedAssets is enabled', async () => { + const networkClientId = 'holesky'; + mockedQuery + .mockReturnValueOnce(Promise.resolve('0x10')) + .mockReturnValueOnce(Promise.resolve('0x11')); + + await withController( + { + options: { + includeStakedAssets: true, + getStakedBalanceForChain: jest.fn().mockResolvedValue('0x1'), + }, + isMultiAccountBalancesEnabled: false, + selectedAccount: ACCOUNT_1, + listAccounts: [ACCOUNT_1, ACCOUNT_2], + networkClientById: { + [networkClientId]: buildCustomNetworkClientConfiguration({ + chainId: '0x4268', + }), + }, + }, + async ({ controller }) => { + await controller.refresh(); + + expect(controller.state).toStrictEqual({ + accounts: { + [CHECKSUM_ADDRESS_1]: { balance: '0x10', stakedBalance: '0x1' }, + [CHECKSUM_ADDRESS_2]: { balance: '0x0' }, + }, + accountsByChainId: { + '0x1': { + [CHECKSUM_ADDRESS_1]: { + balance: '0x10', + stakedBalance: '0x1', + }, + [CHECKSUM_ADDRESS_2]: { + balance: '0x0', + }, + }, + }, + }); + }, + ); + }); + + it('should not update staked balance when includeStakedAssets is disabled', async () => { + const networkClientId = 'holesky'; + mockedQuery + .mockReturnValueOnce(Promise.resolve('0x13')) + .mockReturnValueOnce(Promise.resolve('0x14')); + + await withController( + { + options: { + includeStakedAssets: false, + getStakedBalanceForChain: jest.fn().mockResolvedValue('0x1'), + }, + isMultiAccountBalancesEnabled: false, + selectedAccount: ACCOUNT_1, + listAccounts: [ACCOUNT_1, ACCOUNT_2], + networkClientById: { + [networkClientId]: buildCustomNetworkClientConfiguration({ + chainId: '0x4268', + }), + }, + }, + async ({ controller }) => { + await controller.refresh(); + + expect(controller.state).toStrictEqual({ + accounts: { + [CHECKSUM_ADDRESS_1]: { balance: '0x13' }, + [CHECKSUM_ADDRESS_2]: { balance: '0x0' }, + }, + accountsByChainId: { + '0x1': { + [CHECKSUM_ADDRESS_1]: { + balance: '0x13', + }, + [CHECKSUM_ADDRESS_2]: { + balance: '0x0', + }, + }, + }, + }); + }, + ); + }); + + it('should update staked balance when includeStakedAssets and multi-account is enabled', async () => { + const networkClientId = 'holesky'; + mockedQuery + .mockReturnValueOnce(Promise.resolve('0x11')) + .mockReturnValueOnce(Promise.resolve('0x12')); + + await withController( + { + options: { + includeStakedAssets: true, + getStakedBalanceForChain: jest.fn().mockResolvedValue('0x1'), + }, + isMultiAccountBalancesEnabled: true, + selectedAccount: ACCOUNT_1, + listAccounts: [ACCOUNT_1, ACCOUNT_2], + networkClientById: { + [networkClientId]: buildCustomNetworkClientConfiguration({ + chainId: '0x4268', + }), + }, + }, + async ({ controller }) => { + await controller.refresh(); + + expect(controller.state).toStrictEqual({ + accounts: { + [CHECKSUM_ADDRESS_1]: { balance: '0x11', stakedBalance: '0x1' }, + [CHECKSUM_ADDRESS_2]: { balance: '0x12', stakedBalance: '0x1' }, + }, + accountsByChainId: { + '0x1': { + [CHECKSUM_ADDRESS_1]: { + balance: '0x11', + stakedBalance: '0x1', + }, + [CHECKSUM_ADDRESS_2]: { + balance: '0x12', + stakedBalance: '0x1', + }, + }, + }, + }); + }, + ); + }); + + it('should not update staked balance when includeStakedAssets and multi-account is enabled if network unsupported', async () => { + const networkClientId = 'polygon'; + mockedQuery + .mockReturnValueOnce(Promise.resolve('0x11')) + .mockReturnValueOnce(Promise.resolve('0x12')); + + await withController( + { + options: { + includeStakedAssets: true, + getStakedBalanceForChain: jest.fn().mockResolvedValue(undefined), + }, + isMultiAccountBalancesEnabled: true, + selectedAccount: ACCOUNT_1, + listAccounts: [ACCOUNT_1, ACCOUNT_2], + networkClientById: { + [networkClientId]: buildCustomNetworkClientConfiguration({ + chainId: '0x89', + }), + }, + }, + async ({ controller }) => { + await controller.refresh(); + + expect(controller.state).toStrictEqual({ + accounts: { + [CHECKSUM_ADDRESS_1]: { balance: '0x11' }, + [CHECKSUM_ADDRESS_2]: { balance: '0x12' }, + }, + accountsByChainId: { + '0x1': { + [CHECKSUM_ADDRESS_1]: { + balance: '0x11', + }, + [CHECKSUM_ADDRESS_2]: { + balance: '0x12', + }, + }, + }, + }); + }, + ); + }); }); }); @@ -462,6 +758,33 @@ describe('AccountTrackerController', () => { }, ); }); + + it('should sync staked balance with addresses', async () => { + await withController( + { + options: { + includeStakedAssets: true, + getStakedBalanceForChain: jest.fn().mockResolvedValue('0x1'), + }, + isMultiAccountBalancesEnabled: true, + selectedAccount: ACCOUNT_1, + listAccounts: [], + }, + async ({ controller }) => { + mockedQuery + .mockReturnValueOnce(Promise.resolve('0x10')) + .mockReturnValueOnce(Promise.resolve('0x20')); + const result = await controller.syncBalanceWithAddresses([ + ADDRESS_1, + ADDRESS_2, + ]); + expect(result[ADDRESS_1].balance).toBe('0x10'); + expect(result[ADDRESS_2].balance).toBe('0x20'); + expect(result[ADDRESS_1].stakedBalance).toBe('0x1'); + expect(result[ADDRESS_2].stakedBalance).toBe('0x1'); + }, + ); + }); }); it('should call refresh every interval on legacy polling', async () => { @@ -647,8 +970,8 @@ async function withController( const controller = new AccountTrackerController({ messenger: accountTrackerMessenger, - ...options, getStakedBalanceForChain: jest.fn(), + ...options, }); return await testFunction({ diff --git a/packages/assets-controllers/src/AccountTrackerController.ts b/packages/assets-controllers/src/AccountTrackerController.ts index e2762c3dbad..b6ad00396dd 100644 --- a/packages/assets-controllers/src/AccountTrackerController.ts +++ b/packages/assets-controllers/src/AccountTrackerController.ts @@ -26,7 +26,10 @@ import { type Hex, assert } from '@metamask/utils'; import { Mutex } from 'async-mutex'; import { cloneDeep } from 'lodash'; -import type { AssetsContractController } from './AssetsContractController'; +import type { + AssetsContractController, + StakedBalance, +} from './AssetsContractController'; /** * The name of the {@link AccountTrackerController}. @@ -428,19 +431,19 @@ export class AccountTrackerController extends StaticIntervalPollingController> { + ): Promise< + Record + > { const { ethQuery } = this.#getCorrectNetworkClient(networkClientId); return await Promise.all( addresses.map( - ( - address, - ): Promise<[string, string, string | undefined] | undefined> => { + (address): Promise<[string, string, StakedBalance] | undefined> => { return safelyExecuteWithTimeout(async () => { assert(ethQuery, 'Provider not set.'); const balance = await query(ethQuery, 'getBalance', [address]); - let stakedBalance: string | undefined; + let stakedBalance: StakedBalance; if (this.#includeStakedAssets) { stakedBalance = await this.#getStakedBalanceForChain( address, diff --git a/packages/assets-controllers/src/AssetsContractController.test.ts b/packages/assets-controllers/src/AssetsContractController.test.ts index bcfdb6ba141..73ab03a4a11 100644 --- a/packages/assets-controllers/src/AssetsContractController.test.ts +++ b/packages/assets-controllers/src/AssetsContractController.test.ts @@ -1274,4 +1274,140 @@ describe('AssetsContractController', () => { expect(uri.toLowerCase()).toStrictEqual(expectedUri); messenger.clearEventSubscriptions('NetworkController:networkDidChange'); }); + + it('should get the staked ethereum balance for an address', async () => { + const { assetsContract, messenger, provider, networkClientConfiguration } = + await setupAssetContractControllers(); + assetsContract.configure({ provider }); + + mockNetworkWithDefaultChainId({ + networkClientConfiguration, + mocks: [ + // getShares + { + request: { + method: 'eth_call', + params: [ + { + to: '0x4fef9d741011476750a243ac70b9789a63dd47df', + data: '0xf04da65b0000000000000000000000005a3ca5cd63807ce5e4d7841ab32ce6b6d9bbba2d', + }, + 'latest', + ], + }, + response: { + result: + '0x0000000000000000000000000000000000000000000000000de0b6b3a7640000', + }, + }, + // convertToAssets + { + request: { + method: 'eth_call', + params: [ + { + to: '0x4fef9d741011476750a243ac70b9789a63dd47df', + data: '0x07a2d13a0000000000000000000000000000000000000000000000000de0b6b3a7640000', + }, + 'latest', + ], + }, + response: { + result: + '0x0000000000000000000000000000000000000000000000001bc16d674ec80000', + }, + }, + ], + }); + + const balance = await assetsContract.getStakedBalanceForChain( + TEST_ACCOUNT_PUBLIC_ADDRESS, + ); + + // exchange rate shares = 1e18 + // exchange rate share to assets = 2e18 + // user shares = 1e18 + // user assets = 2e18 + + expect(balance).toBeDefined(); + expect(balance).toBe('0x1bc16d674ec80000'); + expect(BigNumber.from(balance).toString()).toBe((2e18).toString()); + + messenger.clearEventSubscriptions('NetworkController:networkDidChange'); + }); + + it('should return default of zero hex as staked ethereum balance if user has no shares', async () => { + const errorSpy = jest.spyOn(console, 'error'); + const { assetsContract, messenger, provider, networkClientConfiguration } = + await setupAssetContractControllers(); + assetsContract.configure({ provider }); + + mockNetworkWithDefaultChainId({ + networkClientConfiguration, + mocks: [ + // getShares + { + request: { + method: 'eth_call', + params: [ + { + to: '0x4fef9d741011476750a243ac70b9789a63dd47df', + data: '0xf04da65b0000000000000000000000005a3ca5cd63807ce5e4d7841ab32ce6b6d9bbba2d', + }, + 'latest', + ], + }, + response: { + result: + '0x0000000000000000000000000000000000000000000000000000000000000000', + }, + }, + ], + }); + + const balance = await assetsContract.getStakedBalanceForChain( + TEST_ACCOUNT_PUBLIC_ADDRESS, + ); + + expect(balance).toBeDefined(); + expect(balance).toBe('0x00'); + expect(BigNumber.from(balance).toString()).toBe('0'); + expect(errorSpy).toHaveBeenCalledTimes(0); + + errorSpy.mockRestore(); + messenger.clearEventSubscriptions('NetworkController:networkDidChange'); + }); + + it('should return default of zero hex as staked ethereum balance if there is any error thrown', async () => { + let error; + const errorSpy = jest + .spyOn(console, 'error') + .mockImplementationOnce((e) => { + error = e; + }); + const { assetsContract, messenger, provider } = + await setupAssetContractControllers(); + assetsContract.configure({ provider }); + + const balance = await assetsContract.getStakedBalanceForChain( + TEST_ACCOUNT_PUBLIC_ADDRESS, + ); + + expect(balance).toBeDefined(); + expect(balance).toBe('0x00'); + expect(BigNumber.from(balance).toString()).toBe('0'); + expect(errorSpy).toHaveBeenCalledTimes(1); + expect(errorSpy).toHaveBeenCalledWith(error); + + errorSpy.mockRestore(); + messenger.clearEventSubscriptions('NetworkController:networkDidChange'); + }); + + it('should throw missing provider error when getting staked ethereum balance and missing provider', async () => { + const { assetsContract, messenger } = await setupAssetContractControllers(); + await expect( + assetsContract.getStakedBalanceForChain(TEST_ACCOUNT_PUBLIC_ADDRESS), + ).rejects.toThrow(MISSING_PROVIDER_ERROR); + messenger.clearEventSubscriptions('NetworkController:networkDidChange'); + }); }); diff --git a/packages/assets-controllers/src/AssetsContractControllerWithNetworkClientId.test.ts b/packages/assets-controllers/src/AssetsContractControllerWithNetworkClientId.test.ts index 97760db63f5..4b40dccb99a 100644 --- a/packages/assets-controllers/src/AssetsContractControllerWithNetworkClientId.test.ts +++ b/packages/assets-controllers/src/AssetsContractControllerWithNetworkClientId.test.ts @@ -1,3 +1,4 @@ +import { BigNumber } from '@ethersproject/bignumber'; import { BUILT_IN_NETWORKS } from '@metamask/controller-utils'; import { NetworkClientType } from '@metamask/network-controller'; @@ -900,4 +901,78 @@ describe('AssetsContractController with NetworkClientId', () => { expect(uri.toLowerCase()).toStrictEqual(expectedUri); messenger.clearEventSubscriptions('NetworkController:networkDidChange'); }); + + it('should get the staked ethereum balance for an address', async () => { + const { assetsContract, messenger, provider, networkClientConfiguration } = + await setupAssetContractControllers(); + assetsContract.configure({ provider }); + + mockNetworkWithDefaultChainId({ + networkClientConfiguration, + mocks: [ + // getShares + { + request: { + method: 'eth_call', + params: [ + { + to: '0x4fef9d741011476750a243ac70b9789a63dd47df', + data: '0xf04da65b0000000000000000000000005a3ca5cd63807ce5e4d7841ab32ce6b6d9bbba2d', + }, + 'latest', + ], + }, + response: { + result: + '0x0000000000000000000000000000000000000000000000000de0b6b3a7640000', + }, + }, + // convertToAssets + { + request: { + method: 'eth_call', + params: [ + { + to: '0x4fef9d741011476750a243ac70b9789a63dd47df', + data: '0x07a2d13a0000000000000000000000000000000000000000000000000de0b6b3a7640000', + }, + 'latest', + ], + }, + response: { + result: + '0x0000000000000000000000000000000000000000000000001bc16d674ec80000', + }, + }, + ], + }); + + const balance = await assetsContract.getStakedBalanceForChain( + TEST_ACCOUNT_PUBLIC_ADDRESS, + 'mainnet', + ); + + // exchange rate shares = 1e18 + // exchange rate share to assets = 2e18 + // user shares = 1e18 + // user assets = 2e18 + + expect(balance).toBeDefined(); + expect(balance).toBe('0x1bc16d674ec80000'); + expect(BigNumber.from(balance).toString()).toBe((2e18).toString()); + + messenger.clearEventSubscriptions('NetworkController:networkDidChange'); + }); + + it('should default staked ethereum balance to undefined if network is not supported', async () => { + const { assetsContract, provider } = await setupAssetContractControllers(); + assetsContract.configure({ provider }); + + const balance = await assetsContract.getStakedBalanceForChain( + TEST_ACCOUNT_PUBLIC_ADDRESS, + 'sepolia', + ); + + expect(balance).toBeUndefined(); + }); });