From 2bb4ed7065e133a46c3687300bf5cd160d3e4cd7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Leszczyk?= Date: Wed, 2 Oct 2024 13:47:22 +0200 Subject: [PATCH 1/2] fix: show balances for custom tokens too --- .../services/balances/BalancesService.ts | 15 +++++++++++---- src/background/services/settings/models.ts | 4 +++- src/hooks/useTokensWithBalances.ts | 7 ++++--- 3 files changed, 18 insertions(+), 8 deletions(-) diff --git a/src/background/services/balances/BalancesService.ts b/src/background/services/balances/BalancesService.ts index 91d9748d..f616a71d 100644 --- a/src/background/services/balances/BalancesService.ts +++ b/src/background/services/balances/BalancesService.ts @@ -6,7 +6,11 @@ import { ModuleManager } from '@src/background/vmModules/ModuleManager'; import { SettingsService } from '../settings/SettingsService'; import { getPriceChangeValues } from './utils/getPriceChangeValues'; import * as Sentry from '@sentry/browser'; -import { NetworkVMType, TokenWithBalance } from '@avalabs/vm-module-types'; +import { + NetworkVMType, + TokenType, + TokenWithBalance, +} from '@avalabs/vm-module-types'; @singleton() export class BalancesService { @@ -30,9 +34,11 @@ export class BalancesService { ); const module = await this.moduleManager.loadModuleByNetwork(network); - const currency = ( - await this.settingsService.getSettings() - ).currency.toLowerCase(); + const settings = await this.settingsService.getSettings(); + const currency = settings.currency.toLowerCase(); + const customTokens = Object.values( + settings.customTokens[network.chainId] ?? {} + ).map((t) => ({ ...t, type: TokenType.ERC20 as const })); const rawBalances = await module.getBalances({ // TODO: Use public key and module.getAddress instead to make this more modular @@ -56,6 +62,7 @@ export class BalancesService { .filter((address): address is string => !!address), network, currency, + customTokens, }); // Apply price changes data, VM Modules don't do this yet diff --git a/src/background/services/settings/models.ts b/src/background/services/settings/models.ts index 012e1854..c0e40004 100644 --- a/src/background/services/settings/models.ts +++ b/src/background/services/settings/models.ts @@ -48,7 +48,9 @@ export type AddCustomTokenData = { }; type CustomTokens = { - [chain: string]: NetworkContractToken; + [chain: string]: { + [tokenAddress: string]: NetworkContractToken; + }; }; export interface SettingsState { diff --git a/src/hooks/useTokensWithBalances.ts b/src/hooks/useTokensWithBalances.ts index d0af068b..d713fabb 100644 --- a/src/hooks/useTokensWithBalances.ts +++ b/src/hooks/useTokensWithBalances.ts @@ -60,18 +60,19 @@ export const useTokensWithBalances = ( if (!network?.chainId) { return {}; } - const customTokenForActiveNetwork = customTokens[network?.chainId]; - if (!customTokenForActiveNetwork) { + const customTokensForActiveNetwork = customTokens[network?.chainId]; + if (!customTokensForActiveNetwork) { return {}; } - return Object.entries(customTokenForActiveNetwork).reduce<{ + return Object.entries(customTokensForActiveNetwork).reduce<{ [address: string]: TokenWithBalance; }>((acc, [address, tokenData]) => { acc[address] = { ...tokenData, type: TokenType.ERC20, balance: 0n, + balanceDisplayValue: '0', }; return acc; From 94a63d269357cd8365bb0b4c2146dbc5db5f8229 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Leszczyk?= Date: Wed, 9 Oct 2024 14:17:40 +0200 Subject: [PATCH 2/2] test: update BalancesService tests --- .../services/balances/BalancesService.test.ts | 45 +++++++++++++++---- 1 file changed, 37 insertions(+), 8 deletions(-) diff --git a/src/background/services/balances/BalancesService.test.ts b/src/background/services/balances/BalancesService.test.ts index 20fd8a0b..3545b492 100644 --- a/src/background/services/balances/BalancesService.test.ts +++ b/src/background/services/balances/BalancesService.test.ts @@ -5,6 +5,7 @@ import { GetBalancesResponse, TokenType } from '@avalabs/vm-module-types'; import { ModuleManager } from '@src/background/vmModules/ModuleManager'; import { NetworkVMType } from '@avalabs/core-chains-sdk'; import * as Sentry from '@sentry/browser'; +import { SettingsService } from '../settings/SettingsService'; jest.mock('@src/background/vmModules/ModuleManager'); jest.mock('@sentry/browser', () => { @@ -180,6 +181,9 @@ describe('src/background/services/balances/BalancesService.ts', () => { const moduleMock = { getBalances: jest.fn(), }; + const settingsService: SettingsService = { + getSettings: jest.fn(), + } as any; beforeEach(() => { jest.resetAllMocks(); @@ -191,12 +195,10 @@ describe('src/background/services/balances/BalancesService.ts', () => { .mocked(moduleManager) .loadModuleByNetwork.mockResolvedValue(moduleMock as any); - service = new BalancesService( - { - getSettings: jest.fn().mockResolvedValue({ currency: 'EUR' }), - } as any, - moduleManager - ); + jest + .mocked(settingsService.getSettings) + .mockResolvedValue({ currency: 'EUR', customTokens: {} } as any); + service = new BalancesService(settingsService, moduleManager); }); describe('getBalancesForNetwork', () => { @@ -223,6 +225,7 @@ describe('src/background/services/balances/BalancesService.ts', () => { addresses: ['addressPVM'], currency: 'eur', network, + customTokens: [], }); }); @@ -247,12 +250,35 @@ describe('src/background/services/balances/BalancesService.ts', () => { addresses: ['addressAVM'], currency: 'eur', network, + customTokens: [], }); }); it('should get balances for EVM', async () => { + const customToken = { + address: 'tokenAddress', + contractType: 'ERC-20', + decimals: 18, + name: 'Custom Test Token', + symbol: 'CTT', + } as const; + moduleMock.getBalances.mockResolvedValue(evmResult); - const network = { vmName: NetworkVMType.EVM, caipId: 'evm' } as any; + + jest.mocked(settingsService.getSettings).mockResolvedValueOnce({ + currency: 'USD', + customTokens: { + '1337': { + [customToken.address]: customToken, + }, + }, + } as any); + + const network = { + vmName: NetworkVMType.EVM, + caipId: 'eip155:1337', + chainId: 1337, + } as any; const result = await service.getBalancesForNetwork(network, [account]); expect(result).toEqual({ @@ -270,8 +296,9 @@ describe('src/background/services/balances/BalancesService.ts', () => { expect(moduleMock.getBalances).toHaveBeenCalledTimes(1); expect(moduleMock.getBalances).toHaveBeenCalledWith({ addresses: ['addressC'], - currency: 'eur', + currency: 'usd', network, + customTokens: [{ ...customToken, type: TokenType.ERC20 }], }); }); @@ -297,6 +324,7 @@ describe('src/background/services/balances/BalancesService.ts', () => { addresses: ['addressBTC'], currency: 'eur', network, + customTokens: [], }); }); @@ -324,6 +352,7 @@ describe('src/background/services/balances/BalancesService.ts', () => { addresses: ['addressC'], currency: 'eur', network, + customTokens: [], }); }); });