feat: use evm and avalanche modules to load balances CP-9002 (#43)

__mocks__/@blockaid/client.ts

@@ -0,0 +1,8 @@
+/**
+ * @blockaid/client runs environment checks on import.
+ * This results in errors being thrown if it gets incluced anywhere in the dependency tree.
+ * fetch is not available with the `jest-environment-jsdom` environment.
+ * To avoid errors, make sure we mock this library everywhere.
+ */
+
+export default {};

jest.config.js

@@ -1,4 +1,4 @@
-/** @type {import('ts-jest/dist/types').InitialOptionsTsJest} */
+/** @type {import('ts-jest/dist/types').JestConfigWithTsJest} */
 
 module.exports = {
   clearMocks: true,

package.json

@@ -12,7 +12,7 @@
     "dev": "yarn run build:inpage && webpack -w --config webpack.dev.js",
     "lint": "eslint --fix -c ./.eslintrc.js \"src/**/*.ts*\"",
     "typecheck": "yarn tsc --skipLibCheck --noEmit",
-    "postinstall": "husky install",
+    "postinstall": "husky install && patch-package",
     "prettify": "prettier --write \"src/**/*.ts*\"",
     "setup": "yarn install && yarn allow-scripts",
     "test": "yarn jest",
@@ -24,7 +24,8 @@
   },
   "dependencies": {
     "@avalabs/avalanchejs": "4.0.5",
-    "@avalabs/bitcoin-module": "0.5.0",
+    "@avalabs/bitcoin-module": "0.6.0",
+    "@avalabs/avalanche-module": "0.6.0",
     "@avalabs/bridge-unified": "2.1.0",
     "@avalabs/core-bridge-sdk": "3.1.0-alpha.4",
     "@avalabs/core-chains-sdk": "3.1.0-alpha.4",
@@ -36,11 +37,11 @@
     "@avalabs/core-token-prices-sdk": "3.1.0-alpha.4",
     "@avalabs/core-utils-sdk": "3.1.0-alpha.4",
     "@avalabs/core-wallets-sdk": "3.1.0-alpha.4",
-    "@avalabs/evm-module": "0.5.0",
+    "@avalabs/evm-module": "0.6.0",
     "@avalabs/glacier-sdk": "3.1.0-alpha.4",
     "@avalabs/hw-app-avalanche": "0.14.1",
     "@avalabs/types": "3.1.0-alpha.3",
-    "@avalabs/vm-module-types": "0.5.0",
+    "@avalabs/vm-module-types": "0.6.0",
     "@blockaid/client": "0.10.0",
     "@coinbase/cbpay-js": "1.6.0",
     "@cubist-labs/cubesigner-sdk": "0.3.28",
@@ -169,6 +170,7 @@
     "jest-environment-jsdom": "29.3.1",
     "lint-staged": "12.3.1",
     "node-polyfill-webpack-plugin": "1.1.4",
+    "patch-package": "8.0.0",
     "path-browserify": "1.0.1",
     "postinstall-postinstall": "2.1.0",
     "prettier": "2.5.1",
@@ -179,7 +181,7 @@
     "semantic-release": "19.0.3",
     "shell-quote": "1.8.1",
     "style-loader": "3.3.1",
-    "ts-jest": "29.0.5",
+    "ts-jest": "29.2.5",
     "ts-loader": "9.2.6",
     "typescript": "5.2.2",
     "webextension-polyfill": "0.10.0",
@@ -197,7 +199,8 @@
     "i18next-scanner/vinyl-fs/glob-stream/glob-parent": "5.1.2",
     "got": "11.8.5",
     "@ledgerhq/hw-app-eth": "6.36.1",
-    "@ledgerhq/hw-transport": "6.30.6"
+    "@ledgerhq/hw-transport": "6.30.6",
+    "jest-runner": "29.7.0"
   },
   "prettier": {
     "tabWidth": 2,

patches/jest-runner+29.7.0.patch

@@ -0,0 +1,13 @@
+diff --git a/node_modules/jest-runner/build/index.js b/node_modules/jest-runner/build/index.js
+index 65c0ed1..3bb4938 100644
+--- a/node_modules/jest-runner/build/index.js
++++ b/node_modules/jest-runner/build/index.js
+@@ -124,7 +124,7 @@ class TestRunner extends _types.EmittingTestRunner {
+       enableWorkerThreads: this._globalConfig.workerThreads,
+       exposedMethods: ['worker'],
+       forkOptions: {
+-        serialization: 'json',
++        serialization: 'advanced',
+         stdio: 'pipe'
+       },
+       // The workerIdleMemoryLimit should've been converted to a number during

src/background/connections/extensionConnection/registry.ts

@@ -65,7 +65,6 @@ import { UpdateCollectiblesVisibilityHandler } from '../../services/settings/han
 import { NetworksUpdatedEvents } from '@src/background/services/network/events/networksUpdatedEvent';
 import { NetworkUpdatedEvents } from '@src/background/services/network/events/networkUpdatedEvent';
 import { UpdateBalancesForNetworkHandler } from '@src/background/services/balances/handlers/updateBalancesForNetwork';
-import { GetNftBalancesHandler } from '@src/background/services/balances/handlers/getNftBalances';
 import { RemoveLedgerTransportHandler } from '@src/background/services/ledger/handlers/removeLedgerTransport';
 import { GetLockStateHandler } from '@src/background/services/lock/handlers/getLockState';
 import { LockStateChangedEvents } from '@src/background/services/lock/events/lockStateChangedEvent';
@@ -159,7 +158,6 @@ import { UnifiedBridgeTrackTransfer } from '@src/background/services/unifiedBrid
     useToken: UpdateBalancesForNetworkHandler,
   },
   { token: 'ExtensionRequestHandler', useToken: GetBalancesHandler },
-  { token: 'ExtensionRequestHandler', useToken: GetNftBalancesHandler },
   { token: 'ExtensionRequestHandler', useToken: BridgeGetConfigHandler },
   { token: 'ExtensionRequestHandler', useToken: BridgeGetStateHandler },
   { token: 'ExtensionRequestHandler', useToken: BridgeSetIsDevEnvHandler },

src/background/services/balances/BalanceAggregatorService.test.ts

@@ -1,4 +1,3 @@
-import BN from 'bn.js';
 import * as Sentry from '@sentry/browser';
 import { Network, NetworkToken, NetworkVMType } from '@avalabs/core-chains-sdk';
 
@@ -7,13 +6,13 @@ import { SettingsService } from '../settings/SettingsService';
 import { Account, AccountType } from '../accounts/models';
 
 import { BalanceAggregatorService } from './BalanceAggregatorService';
+import { BALANCES_CACHE_KEY, BalanceServiceEvents } from './models';
+import { LockService } from '../lock/LockService';
 import {
-  BALANCES_CACHE_KEY,
-  BalanceServiceEvents,
   NetworkTokenWithBalance,
+  NftTokenWithBalance,
   TokenType,
-} from './models';
-import { LockService } from '../lock/LockService';
+} from '@avalabs/vm-module-types';
 
 jest.mock('@sentry/browser');
 jest.mock('../lock/LockService');
@@ -111,19 +110,39 @@ describe('src/background/services/balances/BalanceAggregatorService.ts', () => {
   const network1TokenBalance: NetworkTokenWithBalance = {
     ...networkToken1,
     type: TokenType.NATIVE,
-    balance: new BN(100),
+    balance: 100n,
+    balanceDisplayValue: '0.00001',
+    coingeckoId: '',
+  };
+
+  const network1NftTokenBalance: NftTokenWithBalance = {
+    type: TokenType.ERC721,
+    balance: 100n,
+    balanceDisplayValue: '0.00001',
+    address: '0x1',
+    logoUri: 'logouri',
+    logoSmall: '',
+    tokenId: '123',
+    tokenUri: 'tokenuri',
+    collectionName: 'unknown',
+    description: '',
+    name: 'name',
+    symbol: '',
   };
 
   const balanceForNetwork1 = {
     [account1.addressC]: {
       [networkToken1.symbol]: network1TokenBalance,
+      ['0x1-123']: network1NftTokenBalance,
     },
   };
 
   const network2TokenBalance: NetworkTokenWithBalance = {
     ...networkToken2,
     type: TokenType.NATIVE,
-    balance: new BN(100),
+    balance: 100n,
+    balanceDisplayValue: '0.00001',
+    coingeckoId: '',
   };
 
   const balanceForNetwork2 = {
@@ -201,8 +220,24 @@ describe('src/background/services/balances/BalanceAggregatorService.ts', () => {
       );
       expect(balancesServiceMock.getBalancesForNetwork).toBeCalledTimes(2);
       const expected = {
-        [network1.chainId]: balanceForNetwork1,
-        [network2.chainId]: balanceForNetwork2,
+        tokens: {
+          [network1.chainId]: {
+            [account1.addressC]: {
+              [networkToken1.symbol]: network1TokenBalance,
+            },
+          },
+          [network2.chainId]: balanceForNetwork2,
+        },
+        nfts: {
+          [network1.chainId]: {
+            [account1.addressC]: {
+              '0x1-123': network1NftTokenBalance,
+            },
+          },
+          [network2.chainId]: {
+            [account1.addressC]: {},
+          },
+        },
       };
       expect(balances).toEqual(expected);
     });
@@ -214,8 +249,20 @@ describe('src/background/services/balances/BalanceAggregatorService.ts', () => {
       );
 
       const expected = {
-        [network1.chainId]: balanceForTwoAccounts,
-        [network2.chainId]: balanceForTwoAccounts,
+        tokens: {
+          [network1.chainId]: balanceForTwoAccounts,
+          [network2.chainId]: balanceForTwoAccounts,
+        },
+        nfts: {
+          [network1.chainId]: {
+            [account1.addressC]: {},
+            [account2.addressC]: {},
+          },
+          [network2.chainId]: {
+            [account1.addressC]: {},
+            [account2.addressC]: {},
+          },
+        },
       };
 
       expect(result).toEqual(expected);
@@ -264,17 +311,40 @@ describe('src/background/services/balances/BalanceAggregatorService.ts', () => {
       );
 
       expect(balancesServiceMock.getBalancesForNetwork).toBeCalledTimes(2);
-      const expected = {
-        [network1.chainId]: balanceForNetwork1,
+      const expectedTokens = {
+        [network1.chainId]: {
+          [account1.addressC]: {
+            [networkToken1.symbol]: network1TokenBalance,
+          },
+        },
         [network2.chainId]: balanceForNetwork2,
       };
-      expect(service.balances).toEqual(expected);
+      expect(service.balances).toEqual(expectedTokens);
+      expect(service.nfts).toEqual({
+        [network1.chainId]: {
+          [account1.addressC]: {
+            '0x1-123': network1NftTokenBalance,
+          },
+        },
+        [network2.chainId]: {
+          [account1.addressC]: {},
+        },
+      });
 
       expect(eventListener).toHaveBeenCalledTimes(1);
       expect(eventListener).toHaveBeenCalledWith({
         balances: {
-          [network1.chainId]: balanceForNetwork1,
-          [network2.chainId]: balanceForNetwork2,
+          tokens: expectedTokens,
+          nfts: {
+            [network1.chainId]: {
+              [account1.addressC]: {
+                '0x1-123': network1NftTokenBalance,
+              },
+            },
+            [network2.chainId]: {
+              [account1.addressC]: {},
+            },
+          },
         },
         isBalancesCached: false,
       });

src/background/services/balances/BalanceAggregatorService.ts

@@ -21,17 +21,25 @@ import {
 import { resolve } from '@avalabs/core-utils-sdk';
 import { SettingsService } from '../settings/SettingsService';
 import { isFulfilled } from '@src/utils/typeUtils';
+import { NftTokenWithBalance } from '@avalabs/vm-module-types';
+import { groupTokensByType } from './utils/groupTokensByType';
+import { BalancesInfo } from './events/balancesUpdatedEvent';
 
 @singleton()
 export class BalanceAggregatorService implements OnLock, OnUnlock {
   #eventEmitter = new EventEmitter();
   #balances: Balances = {};
+  #nfts: Balances<NftTokenWithBalance> = {};
   #isBalancesCached = true;
 
   get balances() {
     return this.#balances;
   }
 
+  get nfts() {
+    return this.#nfts;
+  }
+
   get isBalancesCached() {
     return this.#isBalancesCached;
   }
@@ -47,7 +55,7 @@ export class BalanceAggregatorService implements OnLock, OnUnlock {
   async getBalancesForNetworks(
     chainIds: number[],
     accounts: Account[]
-  ): Promise<Balances> {
+  ): Promise<{ tokens: Balances; nfts: Balances<NftTokenWithBalance> }> {
     const sentryTracker = Sentry.startTransaction({
       name: 'BalanceAggregatorService: getBatchedUpdatedBalancesForNetworks',
     });
@@ -85,34 +93,53 @@ export class BalanceAggregatorService implements OnLock, OnUnlock {
       )
       .map(({ chainId }) => chainId);
 
-    const freshBalances = updatedNetworks.reduce<Balances>(
+    const freshBalances = updatedNetworks.reduce<{
+      nfts: Balances<NftTokenWithBalance>;
+      tokens: Balances;
+    }>(
       (balances, balanceOfNetwork) => {
         const { chainId, networkBalances } = balanceOfNetwork;
-        balances[chainId] = networkBalances;
+
+        const networkBalancesByType = groupTokensByType(networkBalances);
+        balances.tokens[chainId] = networkBalancesByType.tokens;
+        balances.nfts[chainId] = networkBalancesByType.nfts;
 
         return balances;
       },
-      {}
+      { tokens: {}, nfts: {} }
     );
 
-    const aggregatedBalances = merge({}, this.balances, freshBalances);
+    const aggregatedBalances = merge({}, this.balances, freshBalances.tokens);
+    // NFTs don't have balance = 0, if they are sent they should be removed
+    // from the list, hence deep merge doesn't work
+    const aggregatedNfts = {
+      ...this.nfts,
+      ...freshBalances.nfts,
+    };
     const hasChanges = networksWithChanges.length > 0;
 
     if (hasChanges && !this.lockService.locked) {
       this.#balances = aggregatedBalances;
+      this.#nfts = aggregatedNfts;
 
       this.#eventEmitter.emit(BalanceServiceEvents.UPDATED, {
-        balances: aggregatedBalances,
+        balances: {
+          tokens: aggregatedBalances,
+          nfts: aggregatedNfts,
+        },
         isBalancesCached: false,
-      });
+      } as BalancesInfo);
       await this.#updateCachedBalancesInfo();
     }
 
     this.#isBalancesCached = false;
 
     sentryTracker.finish();
 
-    return aggregatedBalances;
+    return {
+      tokens: aggregatedBalances,
+      nfts: aggregatedNfts,
+    };
   }
 
   getPriceChangesData = async () => {
@@ -204,9 +231,12 @@ export class BalanceAggregatorService implements OnLock, OnUnlock {
     this.#isBalancesCached = true;
 
     this.#eventEmitter.emit(BalanceServiceEvents.UPDATED, {
-      balances: this.#balances,
+      balances: {
+        tokens: this.#balances,
+        nfts: this.#nfts,
+      },
       isBalancesCached: true,
-    });
+    } as BalancesInfo);
   }
 
   addListener<T = unknown>(