refactor: use BitcoinModule to fetch BTC balances (#14)

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

@@ -12,6 +12,7 @@ import {
   TokenWithBalanceAVM,
   TokenWithBalancePVM,
 } from './models';
+import { decorateWithCaipId } from '@src/utils/caipConversion';
 
 describe('src/background/services/balances/BalancesService.ts', () => {
   const balancesServiceEVMMock = {
@@ -35,17 +36,17 @@ describe('src/background/services/balances/BalancesService.ts', () => {
     isNetworkSupported: jest.fn(),
   } as any;
 
-  const pchain = {
+  const pchain = decorateWithCaipId({
     ...AVALANCHE_XP_TEST_NETWORK,
     vmName: NetworkVMType.PVM,
-  };
+  });
 
-  const xchain = {
+  const xchain = decorateWithCaipId({
     ...AVALANCHE_XP_TEST_NETWORK,
     vmName: NetworkVMType.AVM,
-  };
+  });
 
-  const testChain = {
+  const testChain = decorateWithCaipId({
     chainName: 'testChain',
     chainId: 123,
     vmName: NetworkVMType.EVM,
@@ -59,7 +60,7 @@ describe('src/background/services/balances/BalancesService.ts', () => {
       logoUri: 'testChainTokenLogoUri.com',
     },
     logoUri: 'testChainLogoUri.com',
-  };
+  });
 
   const account = {
     id: '123',

src/background/services/balances/BalancesService.ts

@@ -17,7 +17,7 @@ import { isPchainNetwork } from '../network/utils/isAvalanchePchainNetwork';
 import { TokensPriceShortData } from '../tokens/models';
 import { isXchainNetwork } from '../network/utils/isAvalancheXchainNetwork';
 import { BalancesServiceAVM } from './BalancesServiceAVM';
-import { Network } from '../network/models';
+import { NetworkWithCaipId } from '../network/models';
 
 @singleton()
 export class BalancesService {
@@ -44,7 +44,7 @@ export class BalancesService {
   }
 
   async getBalancesForNetwork(
-    network: Network,
+    network: NetworkWithCaipId,
     accounts: Account[],
     priceChanges?: TokensPriceShortData
   ): Promise<Record<string, Record<string, TokenWithBalance>>> {

src/background/services/balances/BalancesServiceBTC.ts

@@ -1,142 +1,61 @@
-import { satoshiToBtc } from '@avalabs/core-bridge-sdk';
-import { Network, NetworkVMType } from '@avalabs/core-chains-sdk';
-import { balanceToDisplayValue, bigToBN } from '@avalabs/core-utils-sdk';
-import { TokenPricesService } from '@src/background/services/balances/TokenPricesService';
-import { NetworkService } from '@src/background/services/network/NetworkService';
+import { BitcoinModule } from '@avalabs/bitcoin-module';
+import { isString, mapValues } from 'lodash';
 import { singleton } from 'tsyringe';
 import { Account } from '../accounts/models';
 import { SettingsService } from '../settings/SettingsService';
-import { TokenType, TokenWithBalance } from './models';
 import * as Sentry from '@sentry/browser';
 import { TokensPriceShortData } from '../tokens/models';
 import { BitcoinProvider } from '@avalabs/core-wallets-sdk';
 import { getPriceChangeValues } from './utils/getPriceChangeValues';
+import ModuleManager from '@src/background/vmModules/ModuleManager';
+import { NetworkWithCaipId } from '../network/models';
+import { TokenWithBalanceBTC } from '@avalabs/vm-module-types';
 
 @singleton()
 export class BalancesServiceBTC {
-  constructor(
-    private networkService: NetworkService,
-    private tokenPricesService: TokenPricesService,
-    private settingsService: SettingsService
-  ) {}
+  constructor(private settingsService: SettingsService) {}
 
   getServiceForProvider(provider: any) {
     if (provider instanceof BitcoinProvider) return this;
   }
 
   async getBalances(
     accounts: Account[],
-    network: Network,
+    network: NetworkWithCaipId,
     priceChanges?: TokensPriceShortData
-  ): Promise<Record<string, Record<string, TokenWithBalance>>> {
+  ): Promise<Record<string, Record<string, TokenWithBalanceBTC>>> {
     const sentryTracker = Sentry.startTransaction({
       name: 'BalancesServiceBTC: getBalances',
     });
-    const provider = await this.networkService.getBitcoinProvider();
-    const selectedCurrency = (await this.settingsService.getSettings())
-      .currency;
-    const coingeckoTokenId = network.pricingProviders?.coingecko.nativeTokenId;
-    const tokenPrice = coingeckoTokenId
-      ? await this.tokenPricesService.getPriceByCoinId(
-          coingeckoTokenId,
-          selectedCurrency
-        )
-      : undefined;
-
-    const balances = (
-      await Promise.allSettled(
-        accounts.map(async (account) => {
-          // WalletConnect accounts may be imported without the BTC address.
-          // Any accounts like that are filtered out below.
-          if (!account.addressBTC) {
-            throw new Error('This account does not support BTC');
-          }
-
-          if (network.vmName !== NetworkVMType.BITCOIN) {
-            return {
-              address: account.addressBTC,
-              balances: [],
-            };
-          }
-
-          const denomination = network.networkToken.decimals;
-          const {
-            balance: balanceSatoshis,
-            utxos,
-            balanceUnconfirmed: balanceSatoshisUnconfirmed,
-          } = await provider.getUtxoBalance(account.addressBTC, false);
-
-          const balanceBig = satoshiToBtc(balanceSatoshis);
-          const balance = bigToBN(balanceBig, denomination);
-          const balanceUSD =
-            tokenPrice === undefined
-              ? undefined
-              : balanceBig.times(tokenPrice).toNumber();
-
-          const unconfirmedBalanceBig = satoshiToBtc(
-            balanceSatoshisUnconfirmed
-          );
-          const unconfirmedBalance = bigToBN(
-            unconfirmedBalanceBig,
-            denomination
-          );
-          const unconfirmedBalanceUSD =
-            tokenPrice === undefined
-              ? undefined
-              : unconfirmedBalanceBig.times(tokenPrice).toNumber();
-
-          const symbol = network.networkToken.symbol;
+    const currency = (
+      await this.settingsService.getSettings()
+    ).currency.toLowerCase();
+
+    const module = (await ModuleManager.loadModuleByNetwork(
+      network
+    )) as BitcoinModule;
+    const addresses = accounts
+      .map(({ addressBTC }) => addressBTC)
+      .filter(isString);
+
+    const rawBalances = await module.getBalances({
+      addresses,
+      currency,
+      network,
+    });
 
-          return {
-            address: account.addressBTC,
-            balances: {
-              [symbol]: {
-                ...network.networkToken,
-                type: TokenType.NATIVE,
-                balance,
-                balanceDisplayValue: balanceToDisplayValue(
-                  balance,
-                  denomination
-                ),
-                balanceUSD,
-                balanceUsdDisplayValue:
-                  tokenPrice === undefined
-                    ? undefined
-                    : balanceBig.mul(tokenPrice).toFixed(2),
-                priceUSD: tokenPrice,
-                utxos,
-                unconfirmedBalance,
-                unconfirmedBalanceDisplayValue: balanceToDisplayValue(
-                  unconfirmedBalance,
-                  denomination
-                ),
-                unconfirmedBalanceUsdDisplayValue:
-                  tokenPrice === undefined
-                    ? undefined
-                    : unconfirmedBalanceBig.mul(tokenPrice).toFixed(2),
-                unconfirmedBalanceUSD,
-                logoUri:
-                  'https://raw.githubusercontent.com/trustwallet/assets/master/blockchains/bitcoin/info/logo.png',
-                priceChanges: getPriceChangeValues(
-                  symbol,
-                  balanceUSD,
-                  priceChanges
-                ),
-              },
-            },
-          };
-        })
-      )
-    ).reduce((acc, result) => {
-      if (result.status === 'rejected') {
-        return acc;
-      }
+    // Apply price changes data, VM Modules don't do this yet
+    const balances = mapValues(rawBalances, (accountBalance) => {
+      return mapValues(accountBalance, (tokenBalance) => ({
+        ...tokenBalance,
+        priceChanges: getPriceChangeValues(
+          tokenBalance.symbol,
+          tokenBalance.balanceInCurrency,
+          priceChanges
+        ),
+      }));
+    });
 
-      return {
-        ...acc,
-        [result.value.address]: result.value.balances,
-      };
-    }, {});
     sentryTracker.finish();
     return balances;
   }

src/background/services/balances/models.ts

@@ -1,9 +1,9 @@
 import { NetworkContractToken, NetworkToken } from '@avalabs/core-chains-sdk';
+import { TokenWithBalanceBTC } from '@avalabs/vm-module-types';
 import {
   PChainBalance as GlacierPchainBalance,
   XChainBalances as GlacierXchainBalance,
 } from '@avalabs/glacier-sdk';
-import { BitcoinInputUTXOWithOptionalScript } from '@avalabs/core-wallets-sdk';
 import { EnsureDefined } from '@src/background/models';
 import BN from 'bn.js';
 
@@ -40,16 +40,6 @@ export enum TokenType {
   ERC1155 = 'ERC1155',
 }
 
-export interface TokenWithBalanceBTC extends NetworkTokenWithBalance {
-  logoUri: string;
-  utxos: BitcoinInputUTXOWithOptionalScript[];
-  utxosUnconfirmed?: BitcoinInputUTXOWithOptionalScript[];
-  unconfirmedBalance?: BN;
-  unconfirmedBalanceDisplayValue?: string;
-  unconfirmedBalanceUsdDisplayValue?: string;
-  unconfirmedBalanceUSD?: number;
-}
-
 export interface TokenWithBalancePVM extends NetworkTokenWithBalance {
   available?: BN;
   availableUSD?: number;
@@ -208,3 +198,32 @@ export const isAvaxWithUnavailableBalance = (
       ('locked' in token || 'unlockedUnstaked' in token) &&
       token.available?.toNumber() !== token.balance?.toNumber()
   );
+
+export const isNewTokenBalance = (
+  token?: TokenWithBalance
+): token is TokenWithBalanceBTC =>
+  !token ? false : typeof token.balance === 'bigint';
+
+export const getTokenPrice = (token?: TokenWithBalance | null) => {
+  if (!token) {
+    return;
+  }
+
+  return isNewTokenBalance(token) ? token.priceInCurrency : token.priceUSD;
+};
+
+export const getBalanceInCurrency = (token?: TokenWithBalance) => {
+  if (!token) {
+    return undefined;
+  }
+
+  return isNewTokenBalance(token) ? token.balanceInCurrency : token.balanceUSD;
+};
+
+export const getUnconfirmedBalanceInCurrency = (token?: TokenWithBalance) => {
+  if (!token || !hasUnconfirmedBTCBalance(token)) {
+    return undefined;
+  }
+
+  return token.unconfirmedBalanceInCurrency;
+};

src/background/services/balances/utils/isTokenWithBalanceAVM.test.ts

@@ -3,10 +3,13 @@ import { isTokenWithBalanceAVM } from './isTokenWithBalanceAVM';
 import {
   TokenType,
   TokenWithBalanceAVM,
-  TokenWithBalanceBTC,
   TokenWithBalanceEVM,
   TokenWithBalancePVM,
 } from '../models';
+import {
+  TokenWithBalanceBTC,
+  TokenType as VMModulesTokenType,
+} from '@avalabs/vm-module-types';
 
 describe('src/background/services/balances/utils/isTokenWithBalanceAVM.ts', () => {
   const tokenWithBalanceAVM: TokenWithBalanceAVM = {
@@ -42,10 +45,12 @@ describe('src/background/services/balances/utils/isTokenWithBalanceAVM.ts', () =
   };
 
   const tokenWithBalanceBTC: TokenWithBalanceBTC = {
-    type: TokenType.NATIVE,
+    type: VMModulesTokenType.NATIVE,
     name: 'Avalanche',
     symbol: 'AVAX',
-    balance: new BN(1),
+    coingeckoId: 'avax',
+    balance: 1n,
+    balanceDisplayValue: '0.00000001',
     decimals: 9,
     description: 'description',
     logoUri: 'logoUri',

src/background/services/balances/utils/isTokenWithBalancePVM.test.ts

@@ -1,8 +1,8 @@
 import BN from 'bn.js';
+import { TokenWithBalanceBTC } from '@avalabs/vm-module-types';
 import {
   TokenType,
   TokenWithBalanceAVM,
-  TokenWithBalanceBTC,
   TokenWithBalanceEVM,
   TokenWithBalancePVM,
 } from '../models';
@@ -45,7 +45,9 @@ describe('src/background/services/balances/utils/isTokenWithBalancePVM.ts', () =
     type: TokenType.NATIVE,
     name: 'Avalanche',
     symbol: 'AVAX',
-    balance: new BN(1),
+    balance: 1n,
+    balanceDisplayValue: '0.00000001',
+    coingeckoId: 'avax',
     decimals: 9,
     description: 'description',
     logoUri: 'logoUri',

src/background/services/bridge/BridgeService.ts

@@ -33,7 +33,7 @@ import {
 import Big from 'big.js';
 import { BalanceAggregatorService } from '../balances/BalanceAggregatorService';
 import { FeatureGates } from '../featureFlags/models';
-import { TokenWithBalanceBTC } from '../balances/models';
+import { TokenWithBalanceBTC } from '@avalabs/vm-module-types';
 
 @singleton()
 export class BridgeService implements OnLock, OnStorageReady {

src/background/services/bridge/handlers/avalanche_bridgeAsset.ts

@@ -35,13 +35,13 @@ import { FeatureFlagService } from '../../featureFlags/FeatureFlagService';
 import { FeatureGates } from '../../featureFlags/models';
 import { isWalletConnectAccount } from '../../accounts/utils/typeGuards';
 import { NetworkFeeService } from '../../networkFee/NetworkFeeService';
-import { TokenWithBalanceBTC } from '../../balances/models';
 import {
   buildBtcTx,
   getBtcInputUtxos,
   validateBtcSend,
 } from '@src/utils/send/btcSendUtils';
 import { resolve } from '@src/utils/promiseResolver';
+import { TokenWithBalanceBTC } from '@avalabs/vm-module-types';
 
 type BridgeActionParams = [
   currentBlockchain: Blockchain,

src/background/services/onboarding/utils/addXPChainsToFavoriteIfNeeded.ts

@@ -1,7 +1,11 @@
 import { BalanceAggregatorService } from '../../balances/BalanceAggregatorService';
 import { NetworkService } from '../../network/NetworkService';
 import { ChainId } from '@avalabs/core-chains-sdk';
-import { Balances } from '../../balances/models';
+import {
+  Balances,
+  TokenWithBalanceAVM,
+  TokenWithBalancePVM,
+} from '../../balances/models';
 import { PrimaryAccount } from '../../accounts/models';
 import BN from 'bn.js';
 import { Network } from '../../network/models';
@@ -96,7 +100,9 @@ function hasBalance(
       return false;
     }
 
-    const avaxBalance = balance['AVAX'];
+    const avaxBalance = balance['AVAX'] as
+      | TokenWithBalanceAVM
+      | TokenWithBalancePVM;
 
     return avaxBalance && avaxBalance.balance.gt(new BN(0));
   });