From 530d0faa565262ff72ca66381edb0ee4f4d5f7c9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Leszczyk?= Date: Fri, 6 Dec 2024 13:33:09 +0100 Subject: [PATCH] feat: wallet balance rollup (#101) --- .../connections/extensionConnection/models.ts | 1 + .../extensionConnection/registry.ts | 5 + .../services/accounts/AccountsService.ts | 4 + .../handlers/avalanche_getAddressesInRange.ts | 45 +- .../accounts/utils/getAddressesInRange.ts | 21 + .../balances/BalanceAggregatorService.test.ts | 57 +++ .../balances/BalanceAggregatorService.ts | 8 +- .../getTotalBalanceForWallet.test.ts | 475 ++++++++++++++++++ .../getTotalBalanceForWallet.ts | 160 ++++++ .../calculateTotalBalanceForAccounts.test.ts | 97 ++++ .../calculateTotalBalanceForAccounts.ts | 15 + .../helpers/getAccountsWithActivity.test.ts | 90 ++++ .../helpers/getAccountsWithActivity.ts | 60 +++ .../helpers/getAllAddressesForAccounts.ts | 8 + .../helpers/getIncludedNetworks.test.ts | 63 +++ .../helpers/getIncludedNetworks.ts | 18 + .../getTotalBalanceForWallet/helpers/index.ts | 5 + .../helpers/isDone.test.ts | 10 + .../helpers/isDone.ts | 5 + .../helpers/processGlacierAddresses.test.ts | 105 ++++ .../helpers/processGlacierAddresses.ts | 37 ++ .../getTotalBalanceForWallet/index.ts | 3 + .../getTotalBalanceForWallet/models.ts | 22 + .../utils/calculateTotalBalance.test.ts | 74 --- .../balances/utils/calculateTotalBalance.ts | 17 - .../services/glacier/GlacierService.ts | 20 +- src/contexts/BalancesProvider.tsx | 46 +- src/localization/locales/en/translation.json | 1 + src/pages/Accounts/Accounts.tsx | 29 +- src/pages/Accounts/components/AccountItem.tsx | 2 +- .../components/AccountListImported.tsx | 11 +- .../components/AccountListPrimary.tsx | 8 +- .../Accounts/components/WalletContainer.tsx | 49 +- .../Accounts/components/WalletHeader.tsx | 14 +- .../Accounts/hooks/useWalletTotalBalance.ts | 19 + .../providers/WalletTotalBalanceProvider.tsx | 135 +++++ .../PchainActiveNetworkWidgetContent.tsx | 14 +- src/popup/AppRoutes.tsx | 12 +- src/utils/calculateTotalBalance.test.ts | 19 +- src/utils/calculateTotalBalance.ts | 8 +- src/utils/getAddressForChain.ts | 2 +- src/utils/getAllAddressesForAccount.ts | 2 +- src/utils/getDefaultChainIds.ts | 15 + src/utils/hasAccountBalances.ts | 2 +- 44 files changed, 1617 insertions(+), 196 deletions(-) create mode 100644 src/background/services/accounts/utils/getAddressesInRange.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/getTotalBalanceForWallet.test.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/getTotalBalanceForWallet.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/calculateTotalBalanceForAccounts.test.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/calculateTotalBalanceForAccounts.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getAccountsWithActivity.test.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getAccountsWithActivity.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getAllAddressesForAccounts.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getIncludedNetworks.test.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getIncludedNetworks.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/index.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/isDone.test.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/isDone.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/processGlacierAddresses.test.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/processGlacierAddresses.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/index.ts create mode 100644 src/background/services/balances/handlers/getTotalBalanceForWallet/models.ts delete mode 100644 src/background/services/balances/utils/calculateTotalBalance.test.ts delete mode 100644 src/background/services/balances/utils/calculateTotalBalance.ts create mode 100644 src/pages/Accounts/hooks/useWalletTotalBalance.ts create mode 100644 src/pages/Accounts/providers/WalletTotalBalanceProvider.tsx create mode 100644 src/utils/getDefaultChainIds.ts diff --git a/src/background/connections/extensionConnection/models.ts b/src/background/connections/extensionConnection/models.ts index 8104c38eb..94b73d7f8 100644 --- a/src/background/connections/extensionConnection/models.ts +++ b/src/background/connections/extensionConnection/models.ts @@ -25,6 +25,7 @@ export enum ExtensionRequest { BALANCES_GET = 'balances_get', BALANCES_START_POLLING = 'balances_start_polling', BALANCES_STOP_POLLING = 'balances_stop_polling', + BALANCES_GET_TOTAL_FOR_WALLET = 'balance_get_total_for_wallet', NETWORK_BALANCES_UPDATE = 'network_balances_update', NFT_BALANCES_GET = 'nft_balances_get', NFT_REFRESH_METADATA = 'nft_refresh_metadata', diff --git a/src/background/connections/extensionConnection/registry.ts b/src/background/connections/extensionConnection/registry.ts index 49f607ce1..f2e7e4867 100644 --- a/src/background/connections/extensionConnection/registry.ts +++ b/src/background/connections/extensionConnection/registry.ts @@ -129,6 +129,7 @@ import { StopBalancesPollingHandler } from '@src/background/services/balances/ha import { BalancesUpdatedEvents } from '@src/background/services/balances/events/balancesUpdatedEvent'; import { UnifiedBridgeTrackTransfer } from '@src/background/services/unifiedBridge/handlers/unifiedBridgeTrackTransfer'; import { UpdateActionTxDataHandler } from '@src/background/services/actions/handlers/updateTxData'; +import { GetTotalBalanceForWalletHandler } from '@src/background/services/balances/handlers/getTotalBalanceForWallet/getTotalBalanceForWallet'; /** * TODO: GENERATE THIS FILE AS PART OF THE BUILD PROCESS @@ -373,6 +374,10 @@ import { UpdateActionTxDataHandler } from '@src/background/services/actions/hand token: 'ExtensionRequestHandler', useToken: StopBalancesPollingHandler, }, + { + token: 'ExtensionRequestHandler', + useToken: GetTotalBalanceForWalletHandler, + }, ]) export class ExtensionRequestHandlerRegistry {} diff --git a/src/background/services/accounts/AccountsService.ts b/src/background/services/accounts/AccountsService.ts index 43bc8e0fb..919e11d82 100644 --- a/src/background/services/accounts/AccountsService.ts +++ b/src/background/services/accounts/AccountsService.ts @@ -333,6 +333,10 @@ export class AccountsService implements OnLock, OnUnlock { ); } + getPrimaryAccountsByWalletId(walletId: string) { + return this.accounts.primary[walletId] ?? []; + } + #buildAccount( accountData, importType: ImportType, diff --git a/src/background/services/accounts/handlers/avalanche_getAddressesInRange.ts b/src/background/services/accounts/handlers/avalanche_getAddressesInRange.ts index 2e85c3ba1..80a1c847f 100644 --- a/src/background/services/accounts/handlers/avalanche_getAddressesInRange.ts +++ b/src/background/services/accounts/handlers/avalanche_getAddressesInRange.ts @@ -2,7 +2,6 @@ import { ethErrors } from 'eth-rpc-errors'; import { injectable } from 'tsyringe'; import { DAppRequestHandler } from '@src/background/connections/dAppConnection/DAppRequestHandler'; import { DAppProviderRequest } from '@src/background/connections/dAppConnection/models'; -import { Avalanche } from '@avalabs/core-wallets-sdk'; import { SecretsService } from '../../secrets/SecretsService'; import { NetworkService } from '../../network/NetworkService'; import { canSkipApproval } from '@src/utils/canSkipApproval'; @@ -21,6 +20,7 @@ type Params = [ internalLimit: number ]; import { AccountsService } from '../AccountsService'; +import { getAddressesInRange } from '../utils/getAddressesInRange'; const EXPOSED_DOMAINS = [ 'develop.avacloud-app.pages.dev', @@ -71,38 +71,23 @@ export class AvalancheGetAddressesInRangeHandler extends DAppRequestHandler< if (secrets?.xpubXP) { if (externalLimit > 0) { - for ( - let index = externalStart; - index < externalStart + externalLimit; - index++ - ) { - addresses.external.push( - Avalanche.getAddressFromXpub( - secrets.xpubXP, - index, - provXP, - 'X' - ).split('-')[1] as string // since addresses are the same for X/P we return them without the chain alias prefix (e.g.: fuji1jsduya7thx2ayrawf9dnw7v9jz7vc6xjycra2m) - ); - } + addresses.external = getAddressesInRange( + secrets.xpubXP, + provXP, + false, + externalStart, + externalLimit + ); } if (internalLimit > 0) { - for ( - let index = internalStart; - index < internalStart + internalLimit; - index++ - ) { - addresses.internal.push( - Avalanche.getAddressFromXpub( - secrets.xpubXP, - index, - provXP, - 'X', - true - ).split('-')[1] as string // only X has "internal" (change) addresses, but we remove the chain alias here as well to make it consistent with the external address list - ); - } + addresses.internal = getAddressesInRange( + secrets.xpubXP, + provXP, + true, + internalStart, + internalLimit + ); } } diff --git a/src/background/services/accounts/utils/getAddressesInRange.ts b/src/background/services/accounts/utils/getAddressesInRange.ts new file mode 100644 index 000000000..2ff423ef3 --- /dev/null +++ b/src/background/services/accounts/utils/getAddressesInRange.ts @@ -0,0 +1,21 @@ +import { Avalanche } from '@avalabs/core-wallets-sdk'; + +export function getAddressesInRange( + xpubXP: string, + providerXP: Avalanche.JsonRpcProvider, + internal = false, + start = 0, + limit = 64 +) { + const addresses: string[] = []; + + for (let i = start; i < start + limit; i++) { + addresses.push( + Avalanche.getAddressFromXpub(xpubXP, i, providerXP, 'P', internal).split( + '-' + )[1] as string + ); + } + + return addresses; +} diff --git a/src/background/services/balances/BalanceAggregatorService.test.ts b/src/background/services/balances/BalanceAggregatorService.test.ts index 39295eadb..ec715bf1d 100644 --- a/src/background/services/balances/BalanceAggregatorService.test.ts +++ b/src/background/services/balances/BalanceAggregatorService.test.ts @@ -218,6 +218,63 @@ describe('src/background/services/balances/BalanceAggregatorService.ts', () => { ); }); + it('only updates the balances of the requested accounts', async () => { + // Mock the existing balances for other accounts + (balancesServiceMock.getBalancesForNetwork as jest.Mock).mockReset(); + + balancesServiceMock.getBalancesForNetwork + .mockResolvedValueOnce({ + [account2.addressC]: { + [networkToken1.symbol]: network1TokenBalance, + }, + }) + .mockResolvedValueOnce({ + [account1.addressC]: { + [networkToken1.symbol]: network1TokenBalance, + }, + }); + + // Get balances for the `account2` so they get cached + await service.getBalancesForNetworks( + [network1.chainId], + [account2], + [TokenType.NATIVE, TokenType.ERC20, TokenType.ERC721] + ); + + expect(balancesServiceMock.getBalancesForNetwork).toHaveBeenCalledTimes( + 1 + ); + + expect(service.balances).toEqual({ + [network1.chainId]: { + [account2.addressC]: { + [networkToken1.symbol]: network1TokenBalance, + }, + }, + }); + + // Now get the balances for the first account and verify the `account2` balances are kept in cache + await service.getBalancesForNetworks( + [network1.chainId], + [account1], + [TokenType.NATIVE, TokenType.ERC20, TokenType.ERC721] + ); + + expect(balancesServiceMock.getBalancesForNetwork).toHaveBeenCalledTimes( + 2 + ); + expect(service.balances).toEqual({ + [network1.chainId]: { + [account1.addressC]: { + [networkToken1.symbol]: network1TokenBalance, + }, + [account2.addressC]: { + [networkToken1.symbol]: network1TokenBalance, + }, + }, + }); + }); + it('can fetch the balance for multiple networks and one account', async () => { const balances = await service.getBalancesForNetworks( [network1.chainId, network2.chainId], diff --git a/src/background/services/balances/BalanceAggregatorService.ts b/src/background/services/balances/BalanceAggregatorService.ts index c6951cc2a..e2b87f1ec 100644 --- a/src/background/services/balances/BalanceAggregatorService.ts +++ b/src/background/services/balances/BalanceAggregatorService.ts @@ -6,7 +6,7 @@ import { BalancesService } from './BalancesService'; import { NetworkService } from '../network/NetworkService'; import { EventEmitter } from 'events'; import * as Sentry from '@sentry/browser'; -import { isEqual, pick } from 'lodash'; +import { isEqual, omit, pick } from 'lodash'; import { LockService } from '../lock/LockService'; import { StorageService } from '../storage/StorageService'; @@ -55,7 +55,8 @@ export class BalanceAggregatorService implements OnLock, OnUnlock { async getBalancesForNetworks( chainIds: number[], accounts: Account[], - tokenTypes: TokenType[] + tokenTypes: TokenType[], + cacheResponse = true ): Promise<{ tokens: Balances; nfts: Balances }> { const sentryTracker = Sentry.startTransaction({ name: 'BalanceAggregatorService: getBatchedUpdatedBalancesForNetworks', @@ -143,6 +144,7 @@ export class BalanceAggregatorService implements OnLock, OnUnlock { for (const [chainId, chainBalances] of freshData) { for (const [address, addressBalance] of Object.entries(chainBalances)) { aggregatedBalances[chainId] = { + ...omit(aggregatedBalances[chainId], address), // Keep cached balances for other accounts ...chainBalances, [address]: addressBalance, }; @@ -150,7 +152,7 @@ export class BalanceAggregatorService implements OnLock, OnUnlock { } } - if (hasChanges && !this.lockService.locked) { + if (cacheResponse && hasChanges && !this.lockService.locked) { this.#balances = aggregatedBalances; this.#nfts = aggregatedNfts; diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/getTotalBalanceForWallet.test.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/getTotalBalanceForWallet.test.ts new file mode 100644 index 000000000..170786798 --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/getTotalBalanceForWallet.test.ts @@ -0,0 +1,475 @@ +import { Network } from '@avalabs/glacier-sdk'; +import { ChainId } from '@avalabs/core-chains-sdk'; +import { TokenType, type TokenWithBalance } from '@avalabs/vm-module-types'; + +import { buildRpcCall } from '@src/tests/test-utils'; +import { ExtensionRequest } from '@src/background/connections/extensionConnection/models'; +import type { SecretsService } from '@src/background/services/secrets/SecretsService'; +import type { NetworkService } from '@src/background/services/network/NetworkService'; +import type { GlacierService } from '@src/background/services/glacier/GlacierService'; +import type { AccountsService } from '@src/background/services/accounts/AccountsService'; +import { + type Accounts, + AccountType, + PrimaryAccount, + ImportedAccount, +} from '@src/background/services/accounts/models'; + +import type { Balances } from '../../models'; +import type { BalanceAggregatorService } from '../../BalanceAggregatorService'; + +import { getAccountsWithActivity } from './helpers'; +import { IMPORTED_ACCOUNTS_WALLET_ID } from './models'; +import { GetTotalBalanceForWalletHandler } from './getTotalBalanceForWallet'; + +jest.mock('./helpers/getAccountsWithActivity'); + +describe('background/services/balances/handlers/getTotalBalanceForWallet.test.ts', () => { + const secretsService: jest.Mocked = { + getWalletAccountsSecretsById: jest.fn(), + } as any; + + const glacierService: jest.Mocked = { + getChainIdsForAddresses: jest.fn(), + } as any; + + const networkService: jest.Mocked = { + getAvalanceProviderXP: jest.fn(), + getFavoriteNetworks: jest.fn(), + isMainnet: jest.fn(), + activeNetworks: { + promisify: jest.fn(), + }, + } as any; + + const accountsService: jest.Mocked = { + getAccounts: jest.fn(), + } as any; + + const balanceAggregatorService: jest.Mocked = { + getBalancesForNetworks: jest.fn(), + } as any; + + const FAVORITE_NETWORKS = [ + ChainId.BITCOIN, + ChainId.BITCOIN_TESTNET, + ChainId.ETHEREUM_HOMESTEAD, + ChainId.ETHEREUM_TEST_SEPOLIA, + ]; + const PROVIDER_XP = {} as any; + + const MAINNETS = { + [ChainId.AVALANCHE_MAINNET_ID]: {}, + [ChainId.AVALANCHE_X]: {}, + [ChainId.AVALANCHE_P]: {}, + [ChainId.ETHEREUM_HOMESTEAD]: {}, + [ChainId.BITCOIN]: {}, + } as any; + + const TESTNETS = { + [ChainId.AVALANCHE_TESTNET_ID]: {}, + [ChainId.AVALANCHE_TEST_X]: {}, + [ChainId.AVALANCHE_TEST_P]: {}, + [ChainId.ETHEREUM_TEST_SEPOLIA]: {}, + [ChainId.BITCOIN_TESTNET]: {}, + } as any; + + const buildHandler = () => + new GetTotalBalanceForWalletHandler( + secretsService, + glacierService, + networkService, + accountsService, + balanceAggregatorService + ); + + const handleRequest = (walletId: string) => + buildHandler().handle( + buildRpcCall({ + id: '123', + method: ExtensionRequest.BALANCES_GET_TOTAL_FOR_WALLET, + params: { + walletId, + }, + }) + ); + + const mockEnv = (isMainnet = true) => { + networkService.isMainnet.mockReturnValue(isMainnet); + jest + .mocked(networkService.activeNetworks.promisify) + .mockResolvedValue(isMainnet ? MAINNETS : TESTNETS); + }; + + const mockAccounts = (accounts = ACCOUNTS) => { + accountsService.getAccounts.mockReturnValue(accounts); + }; + + const mockSecrets = (xpubXP?: string) => { + secretsService.getWalletAccountsSecretsById.mockResolvedValueOnce({ + xpubXP, + } as any); + }; + + const mockAccountsWithActivity = (addresses: string[]) => { + jest.mocked(getAccountsWithActivity).mockResolvedValue(addresses); + }; + + const buildAccount = ({ id, ...opts }) => + ({ + id, + name: `name-${id}`, + addressC: `${id}-addressC`, + addressPVM: `${id}-addressPVM`, + addressAVM: `${id}-addressAVM`, + addressBTC: `${id}-addressBTC`, + ...opts, + type: opts.type ?? AccountType.PRIMARY, + } as T); + + const ACCOUNT_IMPORTED_0 = buildAccount({ + id: 'imported-0', + type: AccountType.IMPORTED, + }); + const ACCOUNT_IMPORTED_1 = buildAccount({ + id: 'imported-1', + type: AccountType.IMPORTED, + }); + const ACCOUNT_SEED_0 = buildAccount({ + id: 'seedphrase-0', + index: 0, + walletId: 'seedphrase', + }); + const ACCOUNT_SEED_1 = buildAccount({ + id: 'seedphrase-1', + index: 1, + walletId: 'seedphrase', + }); + const ACCOUNT_LEDGER_0 = buildAccount({ + id: 'ledger-0', + index: 0, + walletId: 'ledger', + }); + const ACCOUNT_LEDGER_1 = buildAccount({ + id: 'ledger-1', + index: 1, + walletId: 'ledger', + }); + const ACCOUNT_SEEDLESS = buildAccount({ + id: 'seedless-0', + index: 0, + walletId: 'seedless', + }); + + const ACCOUNTS: Accounts = { + imported: { + [ACCOUNT_IMPORTED_0.id]: ACCOUNT_IMPORTED_0, + [ACCOUNT_IMPORTED_1.id]: ACCOUNT_IMPORTED_1, + }, + primary: { + seedphrase: [ACCOUNT_SEED_0, ACCOUNT_SEED_1], + ledger: [ACCOUNT_LEDGER_0, ACCOUNT_LEDGER_1], + seedless: [ACCOUNT_SEEDLESS], + }, + }; + + const buildBalance = (symbolOrAddress: string, value: number) => ({ + [symbolOrAddress]: { balanceInCurrency: value } as TokenWithBalance, + }); + + const mockBalances = ( + isMainnet = true, + secondCallBalances?: { + P: Balances[keyof Balances]; + X: Balances[keyof Balances]; + } + ) => { + balanceAggregatorService.getBalancesForNetworks.mockResolvedValue({ + nfts: {}, + tokens: { + [isMainnet + ? ChainId.AVALANCHE_MAINNET_ID + : ChainId.AVALANCHE_TESTNET_ID]: { + [ACCOUNT_SEED_0.addressC]: { + ...buildBalance('AVAX', 100), + ...buildBalance('BTC.b', 1000), + }, + [ACCOUNT_SEED_1.addressC]: { + ...buildBalance('AVAX', 10), + }, + [ACCOUNT_LEDGER_0.addressC]: { + ...buildBalance('AVAX', 20), + }, + [ACCOUNT_LEDGER_1.addressC]: { + ...buildBalance('AVAX', 120), + ...buildBalance('WETH.e', 1300), + }, + [ACCOUNT_IMPORTED_0.addressC]: { + ...buildBalance('AVAX', 50), + }, + [ACCOUNT_IMPORTED_1.addressC]: { + ...buildBalance('AVAX', 75), + }, + [ACCOUNT_SEEDLESS.addressC]: { + ...buildBalance('AVAX', 750), + ...buildBalance('JOE', 43000), + }, + }, + [isMainnet ? ChainId.BITCOIN : ChainId.BITCOIN_TESTNET]: { + [ACCOUNT_SEED_0.addressBTC]: { + ...buildBalance('BTC', 15000), + }, + }, + [isMainnet ? ChainId.AVALANCHE_P : ChainId.AVALANCHE_TEST_P]: { + [ACCOUNT_SEED_0.addressPVM as string]: { + ...buildBalance('AVAX', 350), + }, + ...secondCallBalances?.P, + }, + [isMainnet ? ChainId.AVALANCHE_X : ChainId.AVALANCHE_TEST_X]: { + [ACCOUNT_SEED_1.addressAVM as string]: { + ...buildBalance('AVAX', 650), + }, + ...secondCallBalances?.X, + }, + [isMainnet + ? ChainId.ETHEREUM_HOMESTEAD + : ChainId.ETHEREUM_TEST_SEPOLIA]: { + [ACCOUNT_SEED_1.addressC as string]: { + ...buildBalance('AVAX', 400), + }, + }, + }, + }); + }; + + beforeEach(() => { + jest.resetAllMocks(); + + networkService.getFavoriteNetworks.mockResolvedValue(FAVORITE_NETWORKS); + networkService.getAvalanceProviderXP.mockResolvedValue(PROVIDER_XP); + }); + + describe(`when passed walletId is "${IMPORTED_ACCOUNTS_WALLET_ID}"`, () => { + beforeEach(() => { + mockEnv(true); + mockAccounts(); + mockBalances(true); + }); + + it('does not look for underived addresses', async () => { + const response = await handleRequest(IMPORTED_ACCOUNTS_WALLET_ID); + expect(response.error).toBeUndefined(); + expect(getAccountsWithActivity).not.toHaveBeenCalled(); + }); + + it('only fetches balances for the imported accounts', async () => { + const response = await handleRequest(IMPORTED_ACCOUNTS_WALLET_ID); + expect(response.error).toBeUndefined(); + expect( + balanceAggregatorService.getBalancesForNetworks + ).toHaveBeenCalledTimes(1); + expect( + balanceAggregatorService.getBalancesForNetworks + ).toHaveBeenCalledWith( + expect.any(Array), + [ACCOUNT_IMPORTED_0, ACCOUNT_IMPORTED_1], + [TokenType.NATIVE, TokenType.ERC20] + ); + }); + + it('returns the correct total balance for imported accounts', async () => { + const response = await handleRequest(IMPORTED_ACCOUNTS_WALLET_ID); + expect(response.result).toEqual({ + hasBalanceOnUnderivedAccounts: false, + totalBalanceInCurrency: 125, + }); + }); + }); + + describe('when requested wallet does not include XP public key', () => { + beforeEach(() => { + mockEnv(true); + mockAccounts(); + mockBalances(true); + mockSecrets(undefined); // No xpubXP + }); + + it('does not look for underived addresses', async () => { + const response = await handleRequest('seedless'); + expect(response.error).toBeUndefined(); + expect(getAccountsWithActivity).not.toHaveBeenCalled(); + }); + + it('only fetches balances for already derived accounts', async () => { + const response = await handleRequest('seedless'); + expect(response.error).toBeUndefined(); + expect( + balanceAggregatorService.getBalancesForNetworks + ).toHaveBeenCalledTimes(1); + expect( + balanceAggregatorService.getBalancesForNetworks + ).toHaveBeenCalledWith( + expect.any(Array), + [ACCOUNT_SEEDLESS], + [TokenType.NATIVE, TokenType.ERC20] + ); + }); + + it('returns the correct total balance for already derived accounts', async () => { + const response = await handleRequest('seedless'); + expect(response.result).toEqual({ + hasBalanceOnUnderivedAccounts: false, + totalBalanceInCurrency: 43750, + }); + }); + }); + + describe('when requested wallet does include XP public key', () => { + beforeEach(() => { + mockEnv(true); + mockAccounts(); + mockBalances(true); + mockSecrets('xpubXP'); // We've got xpubXP + }); + + it('looks for XP-chain activity on underived addresses of the requested wallet', async () => { + const unresolvedAddresses = ['avaxUnresolvedAddress']; + mockAccountsWithActivity(unresolvedAddresses); + + const response = await handleRequest('seedphrase'); + expect(response.error).toBeUndefined(); + expect(getAccountsWithActivity).toHaveBeenCalledWith( + 'xpubXP', + PROVIDER_XP, + expect.any(Function) + ); + + // Let's also make sure the passed activity fetcher actually invokes the Glacier API: + const fetcher = jest.mocked(getAccountsWithActivity).mock.lastCall?.[2]; + expect(fetcher).toEqual(expect.any(Function)); + fetcher?.(unresolvedAddresses); + expect(glacierService.getChainIdsForAddresses).toHaveBeenCalledWith({ + addresses: unresolvedAddresses, + network: Network.MAINNET, + }); + }); + + it('fetches C-, X- and P-Chain balances along with favorite networks for already derived accounts within the wallet', async () => { + mockAccountsWithActivity([]); // No underived addresses with activity + + const response = await handleRequest('seedphrase'); + expect(response.error).toBeUndefined(); + expect( + balanceAggregatorService.getBalancesForNetworks + ).toHaveBeenCalledTimes(1); + expect( + balanceAggregatorService.getBalancesForNetworks + ).toHaveBeenCalledWith( + [ + ChainId.AVALANCHE_MAINNET_ID, + ChainId.AVALANCHE_P, + ChainId.AVALANCHE_X, + ChainId.BITCOIN, + ChainId.ETHEREUM_HOMESTEAD, + ], + [ACCOUNT_SEED_0, ACCOUNT_SEED_1], + [TokenType.NATIVE, TokenType.ERC20] + ); + + expect(response.result).toEqual({ + hasBalanceOnUnderivedAccounts: false, + totalBalanceInCurrency: 17510, + }); + }); + + it('works with testnets', async () => { + mockEnv(false); + mockBalances(false); + mockAccountsWithActivity([]); // No underived addresses with activity + + const response = await handleRequest('seedphrase'); + expect(response.error).toBeUndefined(); + expect( + balanceAggregatorService.getBalancesForNetworks + ).toHaveBeenCalledTimes(1); + expect( + balanceAggregatorService.getBalancesForNetworks + ).toHaveBeenCalledWith( + [ + ChainId.AVALANCHE_TESTNET_ID, + ChainId.AVALANCHE_TEST_P, + ChainId.AVALANCHE_TEST_X, + ChainId.BITCOIN_TESTNET, + ChainId.ETHEREUM_TEST_SEPOLIA, + ], + [ACCOUNT_SEED_0, ACCOUNT_SEED_1], + [TokenType.NATIVE, TokenType.ERC20] + ); + + expect(response.result).toEqual({ + hasBalanceOnUnderivedAccounts: false, + totalBalanceInCurrency: 17510, + }); + }); + + it('fetches XP balances for underived accounts with activity', async () => { + const xpAddress = 'ledger-2-address'; + const underivedAddresses = [xpAddress]; // One underived account with activity + mockAccountsWithActivity(underivedAddresses); + mockBalances(true, { + X: { + [`X-${xpAddress}`]: { + ...buildBalance('AVAX', 300), + }, + }, + P: { + [`P-${xpAddress}`]: { + ...buildBalance('AVAX', 450), + }, + }, + }); + + const response = await handleRequest('ledger'); + expect(response.error).toBeUndefined(); + + // Fetching balances of derived accounts + expect( + balanceAggregatorService.getBalancesForNetworks + ).toHaveBeenCalledTimes(2); + expect( + balanceAggregatorService.getBalancesForNetworks + ).toHaveBeenNthCalledWith( + 1, + [ + ChainId.AVALANCHE_MAINNET_ID, + ChainId.AVALANCHE_P, + ChainId.AVALANCHE_X, + ChainId.BITCOIN, + ChainId.ETHEREUM_HOMESTEAD, + ], + [ACCOUNT_LEDGER_0, ACCOUNT_LEDGER_1], + [TokenType.NATIVE, TokenType.ERC20] + ); + + // Fetching XP balances of underived accounts, without caching + expect( + balanceAggregatorService.getBalancesForNetworks + ).toHaveBeenCalledTimes(2); + expect( + balanceAggregatorService.getBalancesForNetworks + ).toHaveBeenNthCalledWith( + 2, + [ChainId.AVALANCHE_P, ChainId.AVALANCHE_X], + [{ addressPVM: `P-${xpAddress}`, addressAVM: `X-${xpAddress}` }], + [TokenType.NATIVE], + false + ); + + expect(response.result).toEqual({ + hasBalanceOnUnderivedAccounts: true, + totalBalanceInCurrency: 2190, // 750 on underived accounts + 1440 on those mocked by default (already derived) + }); + }); + }); +}); diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/getTotalBalanceForWallet.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/getTotalBalanceForWallet.ts new file mode 100644 index 000000000..49288cfd7 --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/getTotalBalanceForWallet.ts @@ -0,0 +1,160 @@ +import { injectable } from 'tsyringe'; +import { Network } from '@avalabs/glacier-sdk'; +import { TokenType } from '@avalabs/vm-module-types'; + +import { ExtensionRequest } from '@src/background/connections/extensionConnection/models'; +import { ExtensionRequestHandler } from '@src/background/connections/models'; + +import { SecretsService } from '../../../secrets/SecretsService'; +import { AccountsService } from '../../../accounts/AccountsService'; +import { GlacierService } from '../../../glacier/GlacierService'; +import { NetworkService } from '../../../network/NetworkService'; +import { BalanceAggregatorService } from '../../BalanceAggregatorService'; +import { Account } from '../../../accounts/models'; + +import { + GetTotalBalanceForWalletParams, + TotalBalanceForWallet, + isImportedAccountsRequest, +} from './models'; +import { + calculateTotalBalanceForAccounts, + getAccountsWithActivity, + getAllAddressesForAccounts, + getIncludedNetworks, +} from './helpers'; +import { getXPChainIds } from '@src/utils/getDefaultChainIds'; + +type HandlerType = ExtensionRequestHandler< + ExtensionRequest.BALANCES_GET_TOTAL_FOR_WALLET, + TotalBalanceForWallet, + GetTotalBalanceForWalletParams +>; + +@injectable() +export class GetTotalBalanceForWalletHandler implements HandlerType { + method = ExtensionRequest.BALANCES_GET_TOTAL_FOR_WALLET as const; + + constructor( + private secretService: SecretsService, + private glacierService: GlacierService, + private networkService: NetworkService, + private accountsService: AccountsService, + private balanceAggregatorService: BalanceAggregatorService + ) {} + + #getAddressesActivity = (addresses: string[]) => + this.glacierService.getChainIdsForAddresses({ + addresses, + network: this.networkService.isMainnet() ? Network.MAINNET : Network.FUJI, + }); + + async #findUnderivedAccounts(walletId: string, derivedAccounts: Account[]) { + const secrets = await this.secretService.getWalletAccountsSecretsById( + walletId + ); + + const derivedWalletAddresses = getAllAddressesForAccounts( + derivedAccounts ?? [] + ); + const derivedAddressesUnprefixed = derivedWalletAddresses.map((addr) => + addr.replace(/^[PXC]-/i, '') + ); + const underivedXPChainAddresses = secrets.xpubXP + ? ( + await getAccountsWithActivity( + secrets.xpubXP, + await this.networkService.getAvalanceProviderXP(), + this.#getAddressesActivity + ) + ).filter((address) => !derivedAddressesUnprefixed.includes(address)) + : []; + + return underivedXPChainAddresses.map>((address) => ({ + addressPVM: `P-${address}`, + addressAVM: `X-${address}`, + })); + } + + handle: HandlerType['handle'] = async ({ request }) => { + const { walletId } = request.params; + const requestsImportedAccounts = isImportedAccountsRequest(walletId); + + try { + const allAccounts = this.accountsService.getAccounts(); + const derivedAccounts = requestsImportedAccounts + ? Object.values(allAccounts.imported ?? {}) + : allAccounts.primary[walletId] ?? []; + + if (!derivedAccounts.length) { + return { + ...request, + result: { + totalBalanceInCurrency: 0, + hasBalanceOnUnderivedAccounts: false, + }, + }; + } + + const underivedAccounts = requestsImportedAccounts + ? [] + : await this.#findUnderivedAccounts(walletId, derivedAccounts); + + const networksIncludedInTotal = getIncludedNetworks( + this.networkService.isMainnet(), + await this.networkService.activeNetworks.promisify(), + await this.networkService.getFavoriteNetworks() + ); + + // Get balance for derived addresses + const { tokens: derivedAddressesBalances } = + await this.balanceAggregatorService.getBalancesForNetworks( + networksIncludedInTotal, + derivedAccounts, + [TokenType.NATIVE, TokenType.ERC20] + ); + + let totalBalanceInCurrency = calculateTotalBalanceForAccounts( + derivedAddressesBalances, + derivedAccounts, + networksIncludedInTotal + ); + let hasBalanceOnUnderivedAccounts = false; + + if (underivedAccounts.length > 0) { + // Get balance for underived addresses for X- and P-Chain. + // We DO NOT cache this response. When fetching balances for multiple X/P addresses at once, + // Glacier responds with all the balances aggregated into one record and the Avalanche Module + // returns it with the first address as the key. If cached, we'd save incorrect data. + const { tokens: underivedAddressesBalances } = + await this.balanceAggregatorService.getBalancesForNetworks( + getXPChainIds(this.networkService.isMainnet()), + underivedAccounts as Account[], + [TokenType.NATIVE], + false // Don't cache this + ); + + const underivedAccountsTotal = calculateTotalBalanceForAccounts( + underivedAddressesBalances, + underivedAccounts, + getXPChainIds(this.networkService.isMainnet()) + ); + totalBalanceInCurrency += underivedAccountsTotal; + hasBalanceOnUnderivedAccounts = underivedAccountsTotal > 0; + } + + return { + ...request, + result: { + totalBalanceInCurrency, + hasBalanceOnUnderivedAccounts, + }, + }; + } catch (e: any) { + return { + ...request, + error: e.toString(), + }; + } + }; +} diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/calculateTotalBalanceForAccounts.test.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/calculateTotalBalanceForAccounts.test.ts new file mode 100644 index 000000000..109ca6bfc --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/calculateTotalBalanceForAccounts.test.ts @@ -0,0 +1,97 @@ +import { Account } from '@src/background/services/accounts/models'; +import { calculateTotalBalance } from '@src/utils/calculateTotalBalance'; + +import { calculateTotalBalanceForAccounts } from './calculateTotalBalanceForAccounts'; + +jest.mock('@src/utils/calculateTotalBalance'); + +describe('src/background/services/balances/handlers/helpers/calculateTotalBalanceForAccounts', () => { + it('aggregates results of calculateTotalBalance() for provided accounts', () => { + jest + .mocked(calculateTotalBalance) + .mockReturnValueOnce({ + sum: 100, + priceChange: { + percentage: [0], + value: 0, + }, + }) + .mockReturnValueOnce({ + sum: 0, + priceChange: { + percentage: [0], + value: 0, + }, + }) + .mockReturnValueOnce({ + sum: null, + priceChange: { + percentage: [], + value: 0, + }, + }) + .mockReturnValueOnce({ + sum: 1500, + priceChange: { + percentage: [0], + value: 0, + }, + }); + + const accounts: Partial[] = [ + { + addressAVM: 'addressAVM', + addressPVM: 'addressPVM', + }, + { + addressPVM: 'addressPVM', + }, + { + addressC: 'addressC', + addressBTC: 'addressBTC', + }, + { + addressC: 'addressC', + addressAVM: 'addressAVM', + addressPVM: 'addressPVM', + }, + ]; + + const balances = {} as any; + const chainIds = []; + + const result = calculateTotalBalanceForAccounts( + balances, + accounts, + chainIds + ); + + expect(calculateTotalBalance).toHaveBeenCalledTimes(4); + expect(calculateTotalBalance).toHaveBeenNthCalledWith( + 1, + accounts[0], + chainIds, + balances + ); + expect(calculateTotalBalance).toHaveBeenNthCalledWith( + 2, + accounts[1], + chainIds, + balances + ); + expect(calculateTotalBalance).toHaveBeenNthCalledWith( + 3, + accounts[2], + chainIds, + balances + ); + expect(calculateTotalBalance).toHaveBeenNthCalledWith( + 4, + accounts[3], + chainIds, + balances + ); + + expect(result).toEqual(1600); + }); +}); diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/calculateTotalBalanceForAccounts.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/calculateTotalBalanceForAccounts.ts new file mode 100644 index 000000000..4b7bdc1a1 --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/calculateTotalBalanceForAccounts.ts @@ -0,0 +1,15 @@ +import { Account } from '@src/background/services/accounts/models'; +import { calculateTotalBalance } from '@src/utils/calculateTotalBalance'; + +import { Balances } from '../../../models'; + +export function calculateTotalBalanceForAccounts( + balances: Balances, + accounts: Partial[], + chainIds: number[] +): number { + return accounts.reduce((sum: number, account: Partial) => { + const accountBalance = calculateTotalBalance(account, chainIds, balances); + return sum + (accountBalance.sum ?? 0); + }, 0); +} diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getAccountsWithActivity.test.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getAccountsWithActivity.test.ts new file mode 100644 index 000000000..9a9431d9c --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getAccountsWithActivity.test.ts @@ -0,0 +1,90 @@ +import { Avalanche } from '@avalabs/core-wallets-sdk'; + +import { processGlacierAddresses } from './processGlacierAddresses'; +import { getAccountsWithActivity } from './getAccountsWithActivity'; + +jest.mock('@avalabs/core-wallets-sdk'); +jest.mock('./processGlacierAddresses'); + +describe('src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getAccountsWithActivity', () => { + const mockedActivityFetcher = jest.fn(); + const xpubXP = 'xpubXP'; + const providerXP = {} as any; + + beforeEach(() => { + jest.resetAllMocks(); + + jest + .spyOn(Avalanche, 'getAddressFromXpub') + .mockImplementation( + (_, index, __, prefix, internal) => + `${prefix}-address${index}/${internal}` + ); + + jest + .mocked(processGlacierAddresses) + // First batch for external addresses + .mockResolvedValueOnce({ + gap: 5, + result: ['ext-address-0', 'ext-address-15'], + }) + // First batch for internal addresses + .mockResolvedValueOnce({ + gap: 9, + result: ['int-address-0', 'int-address-11'], + }) + // Second batch for external addresses + .mockResolvedValueOnce({ + gap: 21, + result: [], + }) + // Second batch for internal addresses + .mockResolvedValue({ + gap: 21, + result: [], + }); + }); + + it(`iterates until it finds a set number of addresses with no activity`, async () => { + await getAccountsWithActivity(xpubXP, providerXP, mockedActivityFetcher); + + // We mocked only two batches (one of external, one of internal) to come back with active + // addresses, therefore it should stop fetching after seeing the next two batches come back empty. + expect(processGlacierAddresses).toHaveBeenCalledTimes(4); + expect(processGlacierAddresses).toHaveBeenNthCalledWith( + 1, + expect.any(Array), + mockedActivityFetcher, + 0 + ); + expect(processGlacierAddresses).toHaveBeenNthCalledWith( + 2, + expect.any(Array), + mockedActivityFetcher, + 0 + ); + expect(processGlacierAddresses).toHaveBeenNthCalledWith( + 3, + expect.any(Array), + mockedActivityFetcher, + 5 + ); + expect(processGlacierAddresses).toHaveBeenNthCalledWith( + 4, + expect.any(Array), + mockedActivityFetcher, + 9 + ); + }); + + it('returns the addresses with activity', async () => { + expect( + await getAccountsWithActivity(xpubXP, providerXP, mockedActivityFetcher) + ).toEqual([ + 'ext-address-0', + 'ext-address-15', + 'int-address-0', + 'int-address-11', + ]); + }); +}); diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getAccountsWithActivity.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getAccountsWithActivity.ts new file mode 100644 index 000000000..40679ba65 --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getAccountsWithActivity.ts @@ -0,0 +1,60 @@ +import { uniq } from 'lodash'; +import { Avalanche } from '@avalabs/core-wallets-sdk'; + +import { getAddressesInRange } from '@src/background/services/accounts/utils/getAddressesInRange'; + +import { + AddressActivityFetcher, + GLACIER_ADDRESS_FETCH_LIMIT, + ITERATION_LIMIT, +} from '../models'; + +import { processGlacierAddresses } from './processGlacierAddresses'; +import { isDone } from './isDone'; + +export async function getAccountsWithActivity( + xpubXP: string, + providerXP: Avalanche.JsonRpcProvider, + activityFetcher: AddressActivityFetcher +) { + let externalGap = 0; + let internalGap = 0; + let externalStart = 0; + let internalStart = 0; + let tooManyIterations = false; + let result: string[] = []; + let iteration = 0; + + while (!isDone(externalGap) && !isDone(internalGap) && !tooManyIterations) { + const external = getAddressesInRange( + xpubXP, + providerXP, + false, + externalStart + ); + const internal = getAddressesInRange( + xpubXP, + providerXP, + true, + internalStart + ); + const [externalResult, internalResult] = await Promise.all([ + processGlacierAddresses(external, activityFetcher, externalGap), + processGlacierAddresses(internal, activityFetcher, internalGap), + ]); + + result = [...result, ...externalResult.result]; + result = [...result, ...internalResult.result]; + + externalGap = externalResult.gap; + externalStart += GLACIER_ADDRESS_FETCH_LIMIT; + + internalGap = internalResult.gap; + internalStart += GLACIER_ADDRESS_FETCH_LIMIT; + + iteration += 1; + tooManyIterations = iteration >= ITERATION_LIMIT; + } + + return uniq(result); +} diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getAllAddressesForAccounts.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getAllAddressesForAccounts.ts new file mode 100644 index 000000000..180989dbe --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getAllAddressesForAccounts.ts @@ -0,0 +1,8 @@ +import { isString } from 'lodash'; + +import type { Account } from '@src/background/services/accounts/models'; +import getAllAddressesForAccount from '@src/utils/getAllAddressesForAccount'; + +export function getAllAddressesForAccounts(accounts: Account[]): string[] { + return accounts.flatMap(getAllAddressesForAccount).filter(isString); +} diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getIncludedNetworks.test.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getIncludedNetworks.test.ts new file mode 100644 index 000000000..c332680a5 --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getIncludedNetworks.test.ts @@ -0,0 +1,63 @@ +import { ChainId } from '@avalabs/core-chains-sdk'; + +import { ChainListWithCaipIds } from '@src/background/services/network/models'; + +import { getIncludedNetworks } from './getIncludedNetworks'; + +describe('src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getIncludedNetworks', () => { + const favoriteNetworks = [ChainId.BITCOIN, ChainId.ETHEREUM_TEST_SEPOLIA]; + + describe('for mainnet environment', () => { + const mainnets = { + [ChainId.AVALANCHE_MAINNET_ID]: {}, + [ChainId.AVALANCHE_X]: {}, + [ChainId.AVALANCHE_P]: {}, + [ChainId.ETHEREUM_HOMESTEAD]: {}, + [ChainId.BITCOIN]: {}, + } as unknown as ChainListWithCaipIds; + + it('always adds C-, X- and P-Chain mainnets', () => { + expect(getIncludedNetworks(true, mainnets, [])).toEqual([ + ChainId.AVALANCHE_MAINNET_ID, + ChainId.AVALANCHE_P, + ChainId.AVALANCHE_X, + ]); + }); + + it('adds favorite networks', () => { + expect(getIncludedNetworks(true, mainnets, favoriteNetworks)).toEqual([ + ChainId.AVALANCHE_MAINNET_ID, + ChainId.AVALANCHE_P, + ChainId.AVALANCHE_X, + ChainId.BITCOIN, + ]); + }); + }); + + describe('for mainnet environment', () => { + const testnets = { + [ChainId.AVALANCHE_TESTNET_ID]: {}, + [ChainId.AVALANCHE_TEST_X]: {}, + [ChainId.AVALANCHE_TEST_P]: {}, + [ChainId.ETHEREUM_TEST_SEPOLIA]: {}, + [ChainId.BITCOIN_TESTNET]: {}, + } as unknown as ChainListWithCaipIds; + + it('always adds testnet C-, X- and P-Chain', () => { + expect(getIncludedNetworks(false, testnets, [])).toEqual([ + ChainId.AVALANCHE_TESTNET_ID, + ChainId.AVALANCHE_TEST_P, + ChainId.AVALANCHE_TEST_X, + ]); + }); + + it('adds favorite networks', () => { + expect(getIncludedNetworks(false, testnets, favoriteNetworks)).toEqual([ + ChainId.AVALANCHE_TESTNET_ID, + ChainId.AVALANCHE_TEST_P, + ChainId.AVALANCHE_TEST_X, + ChainId.ETHEREUM_TEST_SEPOLIA, + ]); + }); + }); +}); diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getIncludedNetworks.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getIncludedNetworks.ts new file mode 100644 index 000000000..fd53907ef --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getIncludedNetworks.ts @@ -0,0 +1,18 @@ +import { uniq } from 'lodash'; + +import { ChainListWithCaipIds } from '@src/background/services/network/models'; +import { getDefaultChainIds } from '@src/utils/getDefaultChainIds'; + +export function getIncludedNetworks( + isMainnet: boolean, + currentChainList: ChainListWithCaipIds, + favoriteChainIds: number[] +) { + const currentEnvNetworks = Object.keys(currentChainList).map(Number); + + return uniq( + [...getDefaultChainIds(isMainnet), ...favoriteChainIds].filter((chainId) => + currentEnvNetworks.includes(chainId) + ) + ); +} diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/index.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/index.ts new file mode 100644 index 000000000..de16e0c64 --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/index.ts @@ -0,0 +1,5 @@ +export * from './calculateTotalBalanceForAccounts'; +export * from './getAccountsWithActivity'; +export * from './getAllAddressesForAccounts'; +export * from './getIncludedNetworks'; +export * from './isDone'; diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/isDone.test.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/isDone.test.ts new file mode 100644 index 000000000..c418e4ba6 --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/isDone.test.ts @@ -0,0 +1,10 @@ +import { ADDRESS_GAP_LIMIT } from '../models'; +import { isDone } from './isDone'; + +describe('src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/isDone', () => { + it(`returns false beyond ${ADDRESS_GAP_LIMIT}`, () => { + expect(isDone(ADDRESS_GAP_LIMIT - 1)).toBe(false); + expect(isDone(ADDRESS_GAP_LIMIT)).toBe(false); + expect(isDone(ADDRESS_GAP_LIMIT + 1)).toBe(true); + }); +}); diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/isDone.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/isDone.ts new file mode 100644 index 000000000..57f02341f --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/isDone.ts @@ -0,0 +1,5 @@ +import { ADDRESS_GAP_LIMIT } from '../models'; + +export function isDone(currentGap: number) { + return currentGap > ADDRESS_GAP_LIMIT; +} diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/processGlacierAddresses.test.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/processGlacierAddresses.test.ts new file mode 100644 index 000000000..a91166c04 --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/processGlacierAddresses.test.ts @@ -0,0 +1,105 @@ +import { startsWith } from 'lodash'; +import { BlockchainIds } from '@avalabs/glacier-sdk'; + +import { ADDRESS_GAP_LIMIT } from '../models'; +import { processGlacierAddresses } from './processGlacierAddresses'; + +describe('src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/getAccountsWithActivity', () => { + const allAddresses = [ + 'active0', + 'active1', + 'active2', + 'active3', + 'inactive4', + 'inactive5', + 'inactive6', + 'inactive7', + 'active8', + 'inactive9', + 'inactive10', + 'inactive11', + 'inactive12', + 'inactive13', + 'inactive14', + 'inactive15', + 'inactive16', + 'inactive17', + 'inactive18', + 'inactive19', + 'inactive20', + 'inactive21', + 'inactive22', + 'inactive23', + 'inactive24', + 'inactive25', + 'inactive26', + 'inactive27', + 'inactive28', + 'inactive29', + 'inactive30', + ]; + + it('returns early if gap is larger than we look for', async () => { + const activityFetcher = jest.fn(); + + expect( + await processGlacierAddresses( + allAddresses, + activityFetcher, + ADDRESS_GAP_LIMIT + 1 + ) + ).toEqual({ + gap: ADDRESS_GAP_LIMIT + 1, + result: [], + }); + + expect(activityFetcher).not.toHaveBeenCalled(); + }); + + it('fetches activity for given set of addresses', async () => { + const activityFetcher = jest.fn().mockResolvedValueOnce({ addresses: [] }); + + await processGlacierAddresses(allAddresses, activityFetcher, 0); + + expect(activityFetcher).toHaveBeenCalledWith(allAddresses); + }); + + it('sums the previous gap with the current gap', async () => { + const mockedActivityFetcher = jest.fn().mockResolvedValueOnce({ + addresses: [], + }); + + const addressesToProcess = allAddresses.slice(-10); // all inactive + const currentGap = 5; + + expect( + await processGlacierAddresses( + allAddresses.slice(-10), + mockedActivityFetcher, + currentGap + ) + ).toEqual({ + gap: currentGap + addressesToProcess.length, + result: [], + }); + }); + + it('returns all active addresses and the number of consecutive inactive addresses (gap)', async () => { + const activeAddresses = allAddresses.filter((addr) => + startsWith(addr, 'active') + ); + const mockedActivityFetcher = jest.fn().mockResolvedValueOnce({ + addresses: activeAddresses.map((address) => ({ + address, + blockchainIds: [BlockchainIds._11111111111111111111111111111111LPO_YY], + })), + }); + + expect( + await processGlacierAddresses(allAddresses, mockedActivityFetcher, 0) + ).toEqual({ + gap: 21, + result: activeAddresses, + }); + }); +}); diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/processGlacierAddresses.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/processGlacierAddresses.ts new file mode 100644 index 000000000..18cb1c7dc --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/helpers/processGlacierAddresses.ts @@ -0,0 +1,37 @@ +import { ChainAddressChainIdMap } from '@avalabs/glacier-sdk'; + +import { AddressActivityFetcher } from '../models'; +import { isDone } from './isDone'; + +export async function processGlacierAddresses( + addresses: string[], + fetchActivity: AddressActivityFetcher, + gap: number +) { + if (isDone(gap)) { + return { gap, result: [] }; + } else { + const { addresses: glacierAddresses } = await fetchActivity(addresses); + + const seenByGlacier: Record = + glacierAddresses.reduce( + (acc, addressInfo) => ({ + ...acc, + [addressInfo.address]: addressInfo, + }), + {} + ); + const result: string[] = []; + for (let i = 0; i < addresses.length && !isDone(gap); i++) { + const address = addresses[i]; + if (address && address in seenByGlacier) { + result.push(address); + gap = 0; + } else { + gap += 1; + } + } + + return { gap, result }; + } +} diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/index.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/index.ts new file mode 100644 index 000000000..d05050d4c --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/index.ts @@ -0,0 +1,3 @@ +import { GetTotalBalanceForWalletHandler } from './getTotalBalanceForWallet'; + +export { GetTotalBalanceForWalletHandler }; diff --git a/src/background/services/balances/handlers/getTotalBalanceForWallet/models.ts b/src/background/services/balances/handlers/getTotalBalanceForWallet/models.ts new file mode 100644 index 000000000..540b7bccb --- /dev/null +++ b/src/background/services/balances/handlers/getTotalBalanceForWallet/models.ts @@ -0,0 +1,22 @@ +import { ChainAddressChainIdMapListResponse } from '@avalabs/glacier-sdk'; + +export type GetTotalBalanceForWalletParams = { + walletId: string; +}; + +export type TotalBalanceForWallet = { + totalBalanceInCurrency?: number; + hasBalanceOnUnderivedAccounts: boolean; +}; + +export type AddressActivityFetcher = ( + addresses: string[] +) => Promise; + +export const ITERATION_LIMIT = 10; // Abitrary number to avoid an infinite loop. +export const ADDRESS_GAP_LIMIT = 20; +export const GLACIER_ADDRESS_FETCH_LIMIT = 64; // Requested addresses are encoded as query params, and Glacier enforces URI length limits +export const IMPORTED_ACCOUNTS_WALLET_ID = '__IMPORTED__'; + +export const isImportedAccountsRequest = (walletId: string) => + walletId === IMPORTED_ACCOUNTS_WALLET_ID; diff --git a/src/background/services/balances/utils/calculateTotalBalance.test.ts b/src/background/services/balances/utils/calculateTotalBalance.test.ts deleted file mode 100644 index 5f2f8b650..000000000 --- a/src/background/services/balances/utils/calculateTotalBalance.test.ts +++ /dev/null @@ -1,74 +0,0 @@ -import { PrimaryNetworkAssetType } from '@avalabs/glacier-sdk'; -import { calculateTotalBalance } from './calculateTotalBalance'; -import BN from 'bn.js'; - -describe('src/background/services/balances/utils/calculateTotalBalance.ts', () => { - it('should return the total balance based on PChainBalance', () => { - const utxoWithAmount1 = { - assetId: '1', - name: 'testToken', - symbol: 'AVAX', - denomination: 9, - type: PrimaryNetworkAssetType.SECP256K1, - amount: '1', - utxoCount: 1, - }; - const utxos = { - unlockedUnstaked: [utxoWithAmount1], - unlockedStaked: [utxoWithAmount1], - lockedPlatform: [utxoWithAmount1], - lockedStakeable: [utxoWithAmount1], - lockedStaked: [utxoWithAmount1], - pendingStaked: [utxoWithAmount1], - - atomicMemoryUnlocked: [ - { - ...utxoWithAmount1, - sharedWithChainId: '', - status: 'testing', - }, - ], - - atomicMemoryLocked: [ - { - ...utxoWithAmount1, - sharedWithChainId: '', - status: 'testing', - }, - ], - }; - - const result = calculateTotalBalance(utxos); - expect(result).toEqual(new BN(8)); - }); - it('should return the total balance based on XChainBalance', () => { - const utxoWithAmount1 = { - assetId: '1', - name: 'testToken', - symbol: 'AVAX', - denomination: 9, - type: PrimaryNetworkAssetType.SECP256K1, - amount: '1', - utxoCount: 1, - }; - const utxos = { - locked: [utxoWithAmount1], - - unlocked: [utxoWithAmount1], - atomicMemoryUnlocked: [ - { - ...utxoWithAmount1, - sharedWithChainId: '', - }, - ], - atomicMemoryLocked: [ - { - ...utxoWithAmount1, - sharedWithChainId: '', - }, - ], - }; - const result = calculateTotalBalance(utxos); - expect(result).toEqual(new BN(4)); - }); -}); diff --git a/src/background/services/balances/utils/calculateTotalBalance.ts b/src/background/services/balances/utils/calculateTotalBalance.ts deleted file mode 100644 index 55b7b4d38..000000000 --- a/src/background/services/balances/utils/calculateTotalBalance.ts +++ /dev/null @@ -1,17 +0,0 @@ -import { PChainBalance, XChainBalances } from '@avalabs/glacier-sdk'; -import BN from 'bn.js'; - -export function calculateTotalBalance( - uxtos: PChainBalance | XChainBalances -): BN { - const sum = Object.values(uxtos).reduce(function (totalAcc, utxoList) { - const typeSum = utxoList.reduce(function (typeAcc, utxo) { - const balanceToAdd = Number(utxo.amount); - return typeAcc + balanceToAdd; - }, 0); - - return totalAcc + typeSum; - }, 0); - - return new BN(sum); -} diff --git a/src/background/services/glacier/GlacierService.ts b/src/background/services/glacier/GlacierService.ts index 776164dcf..7ee7625ca 100644 --- a/src/background/services/glacier/GlacierService.ts +++ b/src/background/services/glacier/GlacierService.ts @@ -1,4 +1,9 @@ -import { Erc1155Token, Erc721Token, Glacier } from '@avalabs/glacier-sdk'; +import { + Erc1155Token, + Erc721Token, + Glacier, + Network, +} from '@avalabs/glacier-sdk'; import { singleton } from 'tsyringe'; import { wait } from '@avalabs/core-utils-sdk'; @@ -13,6 +18,19 @@ export class GlacierService { HEADERS, }); + async getChainIdsForAddresses({ + addresses, + network, + }: { + addresses: string[]; + network: Network; + }) { + return this.glacierSdkInstance.primaryNetwork.getChainIdsForAddresses({ + addresses: addresses.join(','), + network, + }); + } + async refreshNftMetadata(address: string, chainId: string, tokenId: string) { const requestTimestamp = Math.floor(Date.now() / 1000); const maxAttempts = 10; // Amount of fetches after which we give up. diff --git a/src/contexts/BalancesProvider.tsx b/src/contexts/BalancesProvider.tsx index a1aff4db2..f9f32a141 100644 --- a/src/contexts/BalancesProvider.tsx +++ b/src/contexts/BalancesProvider.tsx @@ -33,6 +33,7 @@ import { calculateTotalBalance } from '@src/utils/calculateTotalBalance'; import { NftTokenWithBalance, TokenType } from '@avalabs/vm-module-types'; import { Network } from '@src/background/services/network/models'; import { getAddressForChain } from '@src/utils/getAddressForChain'; +import { getDefaultChainIds } from '@src/utils/getDefaultChainIds'; export const IPFS_URL = 'https://ipfs.io'; @@ -167,8 +168,6 @@ export function BalancesProvider({ children }: { children: any }) { }); const [subscribers, setSubscribers] = useState({}); - const [isPolling, setIsPolling] = useState(false); - const polledChainIds = useMemo( () => favoriteNetworks.map(({ chainId }) => chainId), [favoriteNetworks] @@ -236,11 +235,11 @@ export function BalancesProvider({ children }: { children: any }) { return; } - if (isPolling) { - const tokenTypes = Object.entries(subscribers) - .filter(([, subscriberCount]) => subscriberCount > 0) - .map(([tokenType]) => tokenType as TokenType); + const tokenTypes = Object.entries(subscribers) + .filter(([, subscriberCount]) => subscriberCount > 0) + .map(([tokenType]) => tokenType as TokenType); + if (tokenTypes.length > 0) { request({ method: ExtensionRequest.BALANCES_START_POLLING, params: [activeAccount, polledChainIds, tokenTypes], @@ -250,6 +249,10 @@ export function BalancesProvider({ children }: { children: any }) { payload: balancesData, }); }); + } else { + request({ + method: ExtensionRequest.BALANCES_STOP_POLLING, + }); } return () => { @@ -257,19 +260,7 @@ export function BalancesProvider({ children }: { children: any }) { method: ExtensionRequest.BALANCES_STOP_POLLING, }); }; - }, [ - request, - isPolling, - activeAccount, - network?.chainId, - polledChainIds, - subscribers, - ]); - - useEffect(() => { - // Toggle balance polling based on the amount of dependent components. - setIsPolling(Object.values(subscribers).some((count) => count > 0)); - }, [subscribers]); + }, [request, activeAccount, network?.chainId, polledChainIds, subscribers]); const updateBalanceOnNetworks = useCallback( async (accounts: Account[], chainIds?: number[]) => { @@ -309,18 +300,27 @@ export function BalancesProvider({ children }: { children: any }) { const getTotalBalance = useCallback( (addressC: string) => { - if (balances.tokens) { + if (balances.tokens && network?.chainId) { return calculateTotalBalance( - network, getAccount(addressC), - favoriteNetworks.map(({ chainId }) => chainId), + [ + network.chainId, + ...getDefaultChainIds(!network?.isTestnet), + ...favoriteNetworks.map(({ chainId }) => chainId), + ], balances.tokens ); } return undefined; }, - [getAccount, favoriteNetworks, network, balances.tokens] + [ + getAccount, + favoriteNetworks, + network?.chainId, + network?.isTestnet, + balances.tokens, + ] ); const getTokenPrice = useCallback( diff --git a/src/localization/locales/en/translation.json b/src/localization/locales/en/translation.json index f29e2443e..623de1318 100644 --- a/src/localization/locales/en/translation.json +++ b/src/localization/locales/en/translation.json @@ -989,6 +989,7 @@ "View All Networks": "View All Networks", "View Balance": "View Balance", "View Details": "View Details", + "View P-Chain Details": "View P-Chain Details", "View QR code to scan with your authenticator app.": "View QR code to scan with your authenticator app.", "View Status": "View Status", "View in Explorer": "View in Explorer", diff --git a/src/pages/Accounts/Accounts.tsx b/src/pages/Accounts/Accounts.tsx index ee51e40b3..5b7118803 100644 --- a/src/pages/Accounts/Accounts.tsx +++ b/src/pages/Accounts/Accounts.tsx @@ -4,6 +4,7 @@ import { ChevronLeftIcon, Divider, IconButton, + LoadingDotsIcon, Scrollbars, Stack, TrashIcon, @@ -12,7 +13,6 @@ import { } from '@avalabs/core-k2-components'; import { t } from 'i18next'; import { useHistory } from 'react-router-dom'; -import { TokenType } from '@avalabs/vm-module-types'; import { useAccountsContext } from '@src/contexts/AccountsProvider'; import { useLedgerContext } from '@src/contexts/LedgerProvider'; @@ -20,7 +20,6 @@ import { useAnalyticsContext } from '@src/contexts/AnalyticsProvider'; import { LedgerApprovalDialog } from '@src/pages/SignTransaction/components/LedgerApprovalDialog'; import { AccountType } from '@src/background/services/accounts/models'; -import { useLiveBalance } from '@src/hooks/useLiveBalance'; import { useScopedToast } from '@src/hooks/useScopedToast'; import { NetworkSwitcher } from '@src/components/common/header/NetworkSwitcher'; import { Overlay } from '@src/components/common/Overlay'; @@ -35,8 +34,8 @@ import { AccountListPrimary } from './components/AccountListPrimary'; import { AccountListImported } from './components/AccountListImported'; import { AccountsActionButton } from './components/AccountsActionButton'; import { OverflowingTypography } from './components/OverflowingTypography'; - -const POLLED_BALANCES = [TokenType.NATIVE, TokenType.ERC20]; +import { useWalletTotalBalance } from './hooks/useWalletTotalBalance'; +import { useWalletTotalBalanceContext } from './providers/WalletTotalBalanceProvider'; export function Accounts() { const { @@ -55,6 +54,11 @@ export function Accounts() { const theme = useTheme(); const history = useHistory(); const { walletDetails } = useWalletContext(); + const { isLoading, totalBalanceInCurrency: activeWalletTotalBalance } = + useWalletTotalBalance( + isPrimaryAccount(active) ? active.walletId : undefined + ); + const { fetchBalanceForWallet } = useWalletTotalBalanceContext(); const canCreateAccount = active?.type === AccountType.PRIMARY; const { getTotalBalance } = useBalancesContext(); @@ -64,8 +68,6 @@ export function Accounts() { [active?.addressC, getTotalBalance] ); - useLiveBalance(POLLED_BALANCES); - const addAccountAndFocus = async () => { setAddAccountLoading(true); @@ -75,6 +77,11 @@ export function Accounts() { walletType: walletDetails?.type, }); await selectAccount(id); + + // Refresh total balance of the wallet after adding an account + if (walletDetails?.id) { + fetchBalanceForWallet(walletDetails.id); + } } catch (e) { toast.error(t('An error occurred, please try again later')); } @@ -158,8 +165,14 @@ export function Accounts() { fontSize={13} textAlign="end" color="text.secondary" + // Prevents UI from jumping due to LoadingDotsIcon being larger than they appear + sx={isLoading ? { height: 15, overflow: 'hidden' } : null} > - {/* TODO: total balance of the active wallet */} + {isLoading ? ( + + ) : typeof activeWalletTotalBalance === 'number' ? ( + currencyFormatter(activeWalletTotalBalance) + ) : null} )} @@ -202,7 +215,7 @@ export function Accounts() { - + {hasImportedAccounts && ( diff --git a/src/pages/Accounts/components/AccountItem.tsx b/src/pages/Accounts/components/AccountItem.tsx index 2d05d35d8..a6484011b 100644 --- a/src/pages/Accounts/components/AccountItem.tsx +++ b/src/pages/Accounts/components/AccountItem.tsx @@ -244,7 +244,7 @@ export const AccountItem = forwardRef( /> {address && ( - + { accounts: { active }, } = useAccountsContext(); const [isExpanded, setIsExpanded] = useState(true); + const { isLoading, hasErrorOccurred, totalBalanceInCurrency } = + useWalletTotalBalance(IMPORTED_ACCOUNTS_WALLET_ID); return ( - setIsExpanded((e) => !e)} diff --git a/src/pages/Accounts/components/AccountListPrimary.tsx b/src/pages/Accounts/components/AccountListPrimary.tsx index 3f5151556..534a05aac 100644 --- a/src/pages/Accounts/components/AccountListPrimary.tsx +++ b/src/pages/Accounts/components/AccountListPrimary.tsx @@ -9,20 +9,20 @@ import { useWalletContext } from '@src/contexts/WalletProvider'; import { WalletContainer } from './WalletContainer'; type AccountListProps = { - primaryAccount: Record; + primaryAccounts: Record; sx?: SxProps; }; export const AccountListPrimary = ({ - primaryAccount, + primaryAccounts, sx, }: AccountListProps) => { const { walletDetails: activeWalletDetails, wallets } = useWalletContext(); return ( - {Object.keys(primaryAccount).map((walletId) => { - const walletAccounts = primaryAccount[walletId]; + {Object.keys(primaryAccounts).map((walletId) => { + const walletAccounts = primaryAccounts[walletId]; const walletDetails = wallets.find((wallet) => wallet.id === walletId); if (!walletDetails) { diff --git a/src/pages/Accounts/components/WalletContainer.tsx b/src/pages/Accounts/components/WalletContainer.tsx index deac59d13..739ce5ea4 100644 --- a/src/pages/Accounts/components/WalletContainer.tsx +++ b/src/pages/Accounts/components/WalletContainer.tsx @@ -1,10 +1,19 @@ import { useState } from 'react'; -import { Collapse, Stack } from '@avalabs/core-k2-components'; +import { + Button, + Collapse, + Grow, + OutboundIcon, + Stack, +} from '@avalabs/core-k2-components'; +import { useTranslation } from 'react-i18next'; import { SecretType } from '@src/background/services/secrets/models'; import { WalletDetails } from '@src/background/services/wallet/models'; import { PrimaryAccount } from '@src/background/services/accounts/models'; +import { useNetworkContext } from '@src/contexts/NetworkProvider'; +import { useWalletTotalBalance } from '../hooks/useWalletTotalBalance'; import { SelectionMode } from '../providers/AccountManagerProvider'; import { AccountItem } from './AccountItem'; import WalletHeader from './WalletHeader'; @@ -19,7 +28,15 @@ export const WalletContainer = ({ isActive: boolean; accounts: PrimaryAccount[]; }) => { + const { t } = useTranslation(); + const { isDeveloperMode } = useNetworkContext(); const [isExpanded, setIsExpanded] = useState(true); + const { + isLoading, + hasErrorOccurred, + totalBalanceInCurrency, + hasBalanceOnUnderivedAccounts, + } = useWalletTotalBalance(walletDetails.id); return ( @@ -27,6 +44,9 @@ export const WalletContainer = ({ walletDetails={walletDetails} isActive={isActive} isExpanded={isExpanded} + isLoading={isLoading} + hasBalanceError={hasErrorOccurred} + totalBalance={totalBalanceInCurrency} toggle={() => setIsExpanded((e) => !e)} /> @@ -44,6 +64,33 @@ export const WalletContainer = ({ /> ))} + + + + + ); diff --git a/src/pages/Accounts/components/WalletHeader.tsx b/src/pages/Accounts/components/WalletHeader.tsx index 3fd5309dc..1cf32ec90 100644 --- a/src/pages/Accounts/components/WalletHeader.tsx +++ b/src/pages/Accounts/components/WalletHeader.tsx @@ -5,6 +5,7 @@ import { Grow, IconButton, LedgerIcon, + LoadingDotsIcon, PencilRoundIcon, Stack, Typography, @@ -17,6 +18,7 @@ import { WalletDetails } from '@src/background/services/wallet/models'; import { useAccountManager } from '../providers/AccountManagerProvider'; import { OverflowingTypography } from './OverflowingTypography'; import { useWalletRename } from '../hooks/useWalletRename'; +import { useSettingsContext } from '@src/contexts/SettingsProvider'; const commonTransitionProps = { timeout: 200, @@ -26,6 +28,9 @@ const commonTransitionProps = { type WalletHeaderProps = { isActive: boolean; isExpanded: boolean; + isLoading: boolean; + totalBalance?: number; + hasBalanceError: boolean; toggle: () => void; } & ( | { @@ -40,10 +45,13 @@ export default function WalletHeader({ name, isActive, isExpanded, + isLoading, + totalBalance, toggle, }: WalletHeaderProps) { const { t } = useTranslation(); const { isManageMode } = useAccountManager(); + const { currencyFormatter } = useSettingsContext(); const [isHovered, setIsHovered] = useState(false); const { prompt: promptRename, renderDialog: renameDialog } = @@ -109,7 +117,11 @@ export default function WalletHeader({ textAlign="end" color="text.secondary" > - {/* TODO: total balance of the entire wallet */} + {isLoading ? ( + + ) : typeof totalBalance === 'number' ? ( + currencyFormatter(totalBalance) + ) : null} { + const { walletBalances } = useWalletTotalBalanceContext(); + + return useMemo( + (): WalletTotalBalanceState => + (walletId && walletBalances[walletId]) || { + isLoading: false, + hasErrorOccurred: false, + }, + [walletBalances, walletId] + ); +}; diff --git a/src/pages/Accounts/providers/WalletTotalBalanceProvider.tsx b/src/pages/Accounts/providers/WalletTotalBalanceProvider.tsx new file mode 100644 index 000000000..d00af5610 --- /dev/null +++ b/src/pages/Accounts/providers/WalletTotalBalanceProvider.tsx @@ -0,0 +1,135 @@ +import { + createContext, + useCallback, + useContext, + useEffect, + useMemo, + useState, +} from 'react'; +import { isString } from 'lodash'; + +import { useAccountsContext } from '@src/contexts/AccountsProvider'; +import { + IMPORTED_ACCOUNTS_WALLET_ID, + TotalBalanceForWallet, +} from '@src/background/services/balances/handlers/getTotalBalanceForWallet/models'; +import { useWalletContext } from '@src/contexts/WalletProvider'; +import { useConnectionContext } from '@src/contexts/ConnectionProvider'; +import { GetTotalBalanceForWalletHandler } from '@src/background/services/balances/handlers/getTotalBalanceForWallet'; +import { ExtensionRequest } from '@src/background/connections/extensionConnection/models'; + +interface WalletTotalBalanceContextProps { + children?: React.ReactNode; +} + +export type WalletTotalBalanceState = Partial & { + isLoading: boolean; + hasErrorOccurred: boolean; +}; + +export const WalletTotalBalanceContext = createContext<{ + fetchBalanceForWallet(walletId: string): Promise; + walletBalances: Record; +}>({ + walletBalances: {}, + fetchBalanceForWallet: () => Promise.resolve(), +}); + +export const WalletTotalBalanceProvider = ({ + children, +}: WalletTotalBalanceContextProps) => { + const { + accounts: { imported }, + } = useAccountsContext(); + const { wallets } = useWalletContext(); + const { request } = useConnectionContext(); + + const hasImportedAccounts = useMemo( + () => Object.keys(imported).length > 0, + [imported] + ); + + const [walletBalances, setWalletBalances] = useState< + Record + >({}); + + const fetchBalanceForWallet = useCallback( + async (walletId: string) => { + setWalletBalances((prevState) => ({ + ...prevState, + [walletId]: { + ...prevState[walletId], + hasErrorOccurred: false, + isLoading: true, + }, + })); + + request({ + method: ExtensionRequest.BALANCES_GET_TOTAL_FOR_WALLET, + params: { + walletId, + }, + }) + .then((walletBalanceInfo) => { + setWalletBalances((prevState) => ({ + ...prevState, + [walletId]: { + ...walletBalanceInfo, + hasErrorOccurred: false, + isLoading: false, + }, + })); + }) + .catch((err) => { + console.log('Error while fetching total balance for wallet', err); + setWalletBalances((prevState) => ({ + ...prevState, + [walletId]: { + ...prevState[walletId], + hasErrorOccurred: true, + isLoading: false, + }, + })); + }); + }, + [request] + ); + useEffect(() => { + let isMounted = true; + + const fetchWalletBalancesSequentially = async (walletIds: string[]) => { + for (const walletId of walletIds) { + await fetchBalanceForWallet(walletId); + if (!isMounted) { + return; + } + } + }; + + const walletIds = [ + ...wallets.map(({ id }) => id), + hasImportedAccounts ? IMPORTED_ACCOUNTS_WALLET_ID : undefined, + ].filter(isString); + + fetchWalletBalancesSequentially(walletIds); + + return () => { + isMounted = false; + }; + }, [wallets, hasImportedAccounts, fetchBalanceForWallet]); + + return ( + + {children} + + ); +}; + +export function useWalletTotalBalanceContext() { + return useContext(WalletTotalBalanceContext); +} diff --git a/src/pages/Home/components/Portfolio/NetworkWidget/PchainActiveNetworkWidgetContent.tsx b/src/pages/Home/components/Portfolio/NetworkWidget/PchainActiveNetworkWidgetContent.tsx index 56d4e74d9..21cda2207 100644 --- a/src/pages/Home/components/Portfolio/NetworkWidget/PchainActiveNetworkWidgetContent.tsx +++ b/src/pages/Home/components/Portfolio/NetworkWidget/PchainActiveNetworkWidgetContent.tsx @@ -2,6 +2,7 @@ import { Stack, Typography } from '@avalabs/core-k2-components'; import { useTranslation } from 'react-i18next'; import { BalanceColumn } from '@src/components/common/BalanceColumn'; import { TokenWithBalancePVM } from '@avalabs/vm-module-types'; +import { TokenUnit } from '@avalabs/core-utils-sdk'; interface PchainActiveNetworkWidgetContentProps { balances?: TokenWithBalancePVM; @@ -23,13 +24,13 @@ export function PchainActiveNetworkWidgetContent({ pendingStaked: t('Pending Staked'), }; - if (!balances) { + if (!balances?.balancePerType) { return null; } return ( <> - {Object.keys(typeDisplayNames).map((type) => { - const show = balances[type] && balances[type] > 0; + {Object.entries(balances.balancePerType).map(([type, balance]) => { + const show = balance > 0; if (!show) { return null; @@ -59,7 +60,12 @@ export function PchainActiveNetworkWidgetContent({ data-testid="token-row-token-balance" variant="caption" > - {`${balances[type]} AVAX`} + {new TokenUnit( + balance, + balances.decimals, + balances.symbol + ).toDisplay()}{' '} + AVAX diff --git a/src/popup/AppRoutes.tsx b/src/popup/AppRoutes.tsx index d23e41d4a..e32e516ef 100644 --- a/src/popup/AppRoutes.tsx +++ b/src/popup/AppRoutes.tsx @@ -187,6 +187,14 @@ const AccountManagerProvider = lazy(() => { ); }); +const WalletTotalBalanceProvider = lazy(() => { + return import('../pages/Accounts/providers/WalletTotalBalanceProvider').then( + (m) => ({ + default: m.WalletTotalBalanceProvider, + }) + ); +}); + const AccountDetailsView = lazy(() => { return import('../pages/Accounts/AccountDetailsView').then((m) => ({ default: m.AccountDetailsView, @@ -303,7 +311,9 @@ export const AppRoutes = () => ( }> - + + + diff --git a/src/utils/calculateTotalBalance.test.ts b/src/utils/calculateTotalBalance.test.ts index cd54ebe18..0eae74001 100644 --- a/src/utils/calculateTotalBalance.test.ts +++ b/src/utils/calculateTotalBalance.test.ts @@ -1,9 +1,4 @@ -import { - ChainId, - Network, - NetworkToken, - NetworkVMType, -} from '@avalabs/core-chains-sdk'; +import { ChainId, NetworkToken } from '@avalabs/core-chains-sdk'; import { Account, AccountType } from '@src/background/services/accounts/models'; import { Balances } from '@src/background/services/balances/models'; import { calculateTotalBalance } from './calculateTotalBalance'; @@ -27,17 +22,6 @@ describe('utils/calculateTotalBalance', () => { logoUri: 'network.token.one.com', }; - const network1: Network = { - chainName: 'test network 1', - chainId: ChainId.AVALANCHE_MAINNET_ID, - vmName: NetworkVMType.EVM, - rpcUrl: 'test.one.com/rpc', - explorerUrl: 'https://explorer.url', - networkToken: networkToken1, - logoUri: 'test.one.com/logo', - primaryColor: 'pink', - }; - const network1TokenBalance: NetworkTokenWithBalance = { ...networkToken1, type: TokenType.NATIVE, @@ -68,7 +52,6 @@ describe('utils/calculateTotalBalance', () => { it('it should calculate the balance', () => { const balance = calculateTotalBalance( - network1, account1, [ChainId.AVALANCHE_MAINNET_ID, ChainId.DFK], balances diff --git a/src/utils/calculateTotalBalance.ts b/src/utils/calculateTotalBalance.ts index b634cdb0f..e72fe9bf4 100644 --- a/src/utils/calculateTotalBalance.ts +++ b/src/utils/calculateTotalBalance.ts @@ -1,4 +1,3 @@ -import { Network } from '@avalabs/core-chains-sdk'; import { Account } from '@src/background/services/accounts/models'; import { Balances, @@ -8,12 +7,11 @@ import { getAddressForChain } from '@src/utils/getAddressForChain'; import { hasAccountBalances } from './hasAccountBalances'; export function calculateTotalBalance( - network?: Network, - account?: Account, + account?: Partial, networkIds?: number[], balances?: Balances ) { - if (!account || !balances || !network) { + if (!account || !balances || !networkIds?.length) { return { sum: null, priceChange: { @@ -23,7 +21,7 @@ export function calculateTotalBalance( }; } - const chainIdsToSum = new Set([network.chainId, ...(networkIds ?? [])]); + const chainIdsToSum = new Set(networkIds); const hasBalances = hasAccountBalances( balances, diff --git a/src/utils/getAddressForChain.ts b/src/utils/getAddressForChain.ts index b8c21e65d..96720bee4 100644 --- a/src/utils/getAddressForChain.ts +++ b/src/utils/getAddressForChain.ts @@ -3,7 +3,7 @@ import { Account } from '@src/background/services/accounts/models'; import { isPchainNetworkId } from '@src/background/services/network/utils/isAvalanchePchainNetwork'; import { isXchainNetworkId } from '@src/background/services/network/utils/isAvalancheXchainNetwork'; -export function getAddressForChain(chainId: number, account: Account) { +export function getAddressForChain(chainId: number, account: Partial) { return isBitcoinChainId(chainId) ? account.addressBTC : isPchainNetworkId(chainId) diff --git a/src/utils/getAllAddressesForAccount.ts b/src/utils/getAllAddressesForAccount.ts index 2f7dd4417..8e15d8a55 100644 --- a/src/utils/getAllAddressesForAccount.ts +++ b/src/utils/getAllAddressesForAccount.ts @@ -1,6 +1,6 @@ import { Account } from '@src/background/services/accounts/models'; -export default function getAllAddressesForAccount(acc: Account) { +export default function getAllAddressesForAccount(acc: Partial) { return [ acc.addressC, acc.addressBTC, diff --git a/src/utils/getDefaultChainIds.ts b/src/utils/getDefaultChainIds.ts new file mode 100644 index 000000000..489bede2a --- /dev/null +++ b/src/utils/getDefaultChainIds.ts @@ -0,0 +1,15 @@ +import { ChainId } from '@avalabs/core-chains-sdk'; + +export function getXPChainIds(isMainnet: boolean) { + const xChainId = isMainnet ? ChainId.AVALANCHE_X : ChainId.AVALANCHE_TEST_X; + const pChainId = isMainnet ? ChainId.AVALANCHE_P : ChainId.AVALANCHE_TEST_P; + + return [pChainId, xChainId]; +} + +export function getDefaultChainIds(isMainnet: boolean) { + return [ + isMainnet ? ChainId.AVALANCHE_MAINNET_ID : ChainId.AVALANCHE_TESTNET_ID, + ...getXPChainIds(isMainnet), + ]; +} diff --git a/src/utils/hasAccountBalances.ts b/src/utils/hasAccountBalances.ts index 3f94d1c77..9c03282dc 100644 --- a/src/utils/hasAccountBalances.ts +++ b/src/utils/hasAccountBalances.ts @@ -4,7 +4,7 @@ import getAllAddressesForAccount from './getAllAddressesForAccount'; export function hasAccountBalances( balances: Balances, - account: Account, + account: Partial, networkIds: number[] ) { const accountAddresses = getAllAddressesForAccount(account);