fix: show balances for custom tokens too (#55)

ava-labs · Oct 10, 2024 · c9fce37 · c9fce37
1 parent 9cc2635
commit c9fce37
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 16 deletions.
diff --git 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: [],
       });
     });
   });

diff --git 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
@@ -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
@@ -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;