From a9267de6538b50551050c20d0144e75c39d683fd Mon Sep 17 00:00:00 2001
From: Art <artur@vegaprotocol.io>
Date: Wed, 3 May 2023 21:35:05 +0200
Subject: [PATCH] feat(trading): asset balances inside asset selector (deposit,
 withdraw form) (#3590)

---
 libs/assets/src/lib/asset-option.tsx          | 45 ++++++++++--
 libs/assets/src/lib/index.ts                  |  3 +-
 libs/assets/src/lib/use-balances-store.ts     | 69 +++++++++++++++++++
 libs/deposits/src/lib/asset-balance.tsx       | 42 +++++++++++
 libs/deposits/src/lib/deposit-form.tsx        | 13 +++-
 libs/deposits/src/lib/index.ts                |  5 +-
 libs/deposits/src/lib/use-submit-faucet.ts    |  4 +-
 libs/utils/src/lib/is-asset-erc20.ts          |  2 +-
 libs/withdraws/src/lib/asset-balance.tsx      | 30 ++++++++
 .../src/lib/withdraw-form-container.spec.tsx  |  6 ++
 libs/withdraws/src/lib/withdraw-form.spec.tsx |  5 ++
 libs/withdraws/src/lib/withdraw-form.tsx      |  7 +-
 .../src/lib/withdraw-manager.spec.tsx         |  5 ++
 13 files changed, 223 insertions(+), 13 deletions(-)
 create mode 100644 libs/assets/src/lib/use-balances-store.ts
 create mode 100644 libs/deposits/src/lib/asset-balance.tsx
 create mode 100644 libs/withdraws/src/lib/asset-balance.tsx

diff --git a/libs/assets/src/lib/asset-option.tsx b/libs/assets/src/lib/asset-option.tsx
index 0df06a490..2d67c79a8 100644
--- a/libs/assets/src/lib/asset-option.tsx
+++ b/libs/assets/src/lib/asset-option.tsx
@@ -1,13 +1,50 @@
 import { Option } from '@vegaprotocol/ui-toolkit';
 import type { AssetFieldsFragment } from './__generated__/Asset';
+import classNames from 'classnames';
+import { t } from '@vegaprotocol/i18n';
+import type { ReactNode } from 'react';
 
-export const AssetOption = ({ asset }: { asset: AssetFieldsFragment }) => {
+type AssetOptionProps = {
+  asset: AssetFieldsFragment;
+  balance?: ReactNode;
+};
+
+export const Balance = ({
+  balance,
+  symbol,
+}: {
+  balance?: string;
+  symbol: string;
+}) =>
+  balance ? (
+    <div className="mt-1 font-alpha">
+      {balance} {symbol}
+    </div>
+  ) : (
+    <div className="text-vega-orange-500">{t('Fetching balance…')}</div>
+  );
+
+export const AssetOption = ({ asset, balance }: AssetOptionProps) => {
   return (
     <Option key={asset.id} value={asset.id}>
       <div className="flex flex-col items-start">
-        <span>{asset.name}</span>
-        <div className="text-[10px] font-mono w-full text-left break-all">
-          <span className="text-gray-500">{asset.id} -</span> {asset.symbol}
+        <div className="flex flex-row align-baseline gap-2">
+          <span>{asset.name}</span>{' '}
+          <span
+            className={classNames(
+              'bg-vega-light-200 dark:bg-vega-dark-200',
+              'text-black dark:text-white text-xs',
+              'py-[2px] px-[4px] rounded'
+            )}
+          >
+            {asset.symbol}
+          </span>
+        </div>
+        {balance}
+        <div className="text-[12px] font-mono w-full text-left break-all">
+          <span className="text-vega-light-300 dark:text-vega-dark-300">
+            {asset.id}
+          </span>
         </div>
       </div>
     </Option>
diff --git a/libs/assets/src/lib/index.ts b/libs/assets/src/lib/index.ts
index de6516378..f85c8581f 100644
--- a/libs/assets/src/lib/index.ts
+++ b/libs/assets/src/lib/index.ts
@@ -1,8 +1,9 @@
 export * from './__generated__/Asset';
 export * from './__generated__/Assets';
 export * from './asset-data-provider';
-export * from './assets-data-provider';
 export * from './asset-details-dialog';
 export * from './asset-details-table';
 export * from './asset-option';
+export * from './assets-data-provider';
 export * from './constants';
+export * from './use-balances-store';
diff --git a/libs/assets/src/lib/use-balances-store.ts b/libs/assets/src/lib/use-balances-store.ts
new file mode 100644
index 000000000..d934aa830
--- /dev/null
+++ b/libs/assets/src/lib/use-balances-store.ts
@@ -0,0 +1,69 @@
+import type BigNumber from 'bignumber.js';
+import type { AssetFieldsFragment } from './__generated__/Asset';
+import { create } from 'zustand';
+import * as Sentry from '@sentry/react';
+import { immer } from 'zustand/middleware/immer';
+
+type AssetWithBalance = {
+  asset: AssetFieldsFragment;
+  balanceOnEth?: BigNumber;
+  balanceOnVega?: BigNumber;
+  updatedAt: number;
+};
+
+type SetBalanceArgs = Omit<AssetWithBalance, 'updatedAt'> & {
+  ethBalanceFetcher?: () => Promise<BigNumber | undefined>;
+};
+
+type BalancesStore = {
+  balances: (AssetWithBalance & { fetchFromEth?: () => void })[];
+  getBalance: (assetId: string) => AssetWithBalance | undefined;
+  setBalance: (args: SetBalanceArgs) => void;
+  refetch: (assetId: string) => void;
+};
+
+export const useBalancesStore = create(
+  immer<BalancesStore>((set, get) => ({
+    balances: [],
+    getBalance: (assetId) =>
+      get().balances.find(({ asset: a }) => a.id === assetId),
+    setBalance: ({ asset, balanceOnEth, balanceOnVega, ethBalanceFetcher }) => {
+      set((state) => {
+        const found = state.balances.find(({ asset: a }) => a.id === asset.id);
+        const fetchFromEth = ethBalanceFetcher
+          ? () => {
+              if (!ethBalanceFetcher) return;
+              ethBalanceFetcher()
+                .then((balance) => {
+                  if (balance) {
+                    get().setBalance({ asset, balanceOnEth: balance });
+                  }
+                })
+                .catch((err) => {
+                  Sentry.captureException(err);
+                });
+            }
+          : undefined;
+
+        if (found) {
+          if (balanceOnEth) found.balanceOnEth = balanceOnEth;
+          if (balanceOnVega) found.balanceOnVega = balanceOnVega;
+          if (fetchFromEth) found.fetchFromEth = fetchFromEth;
+          found.updatedAt = Date.now();
+        } else {
+          state.balances.push({
+            asset,
+            balanceOnEth,
+            balanceOnVega,
+            fetchFromEth,
+            updatedAt: Date.now(),
+          });
+        }
+      });
+    },
+    refetch: (assetId) => {
+      const found = get().balances.find((a) => a.asset.id === assetId);
+      found?.fetchFromEth?.();
+    },
+  }))
+);
diff --git a/libs/deposits/src/lib/asset-balance.tsx b/libs/deposits/src/lib/asset-balance.tsx
new file mode 100644
index 000000000..563c8db57
--- /dev/null
+++ b/libs/deposits/src/lib/asset-balance.tsx
@@ -0,0 +1,42 @@
+import { useCallback, useEffect } from 'react';
+import { useGetBalanceOfERC20Token } from './use-get-balance-of-erc20-token';
+import type { AssetFieldsFragment } from '@vegaprotocol/assets';
+import { useBalancesStore } from '@vegaprotocol/assets';
+import { Balance } from '@vegaprotocol/assets';
+import { isAssetTypeERC20 } from '@vegaprotocol/utils';
+import { useTokenContract } from '@vegaprotocol/web3';
+
+const REFETCH_DELAY = 5000;
+
+export const AssetBalance = ({ asset }: { asset: AssetFieldsFragment }) => {
+  const [setBalance, getBalance] = useBalancesStore((state) => [
+    state.setBalance,
+    state.getBalance,
+  ]);
+
+  const tokenContract = useTokenContract(
+    isAssetTypeERC20(asset) ? asset.source.contractAddress : undefined
+  );
+  const ethBalanceFetcher = useGetBalanceOfERC20Token(tokenContract, asset);
+
+  const fetchFromEth = useCallback(async () => {
+    const balance = await ethBalanceFetcher();
+    if (balance) {
+      setBalance({ asset, balanceOnEth: balance, ethBalanceFetcher });
+    }
+  }, [asset, ethBalanceFetcher, setBalance]);
+
+  useEffect(() => {
+    const balance = getBalance(asset.id);
+    if (!balance || Date.now() - balance.updatedAt > REFETCH_DELAY) {
+      fetchFromEth();
+    }
+  }, [asset.id, fetchFromEth, getBalance]);
+
+  return (
+    <Balance
+      balance={getBalance(asset.id)?.balanceOnEth?.toString()}
+      symbol={asset.symbol}
+    />
+  );
+};
diff --git a/libs/deposits/src/lib/deposit-form.tsx b/libs/deposits/src/lib/deposit-form.tsx
index 960107169..9ec2eb05f 100644
--- a/libs/deposits/src/lib/deposit-form.tsx
+++ b/libs/deposits/src/lib/deposit-form.tsx
@@ -41,6 +41,7 @@ import type { DepositBalances } from './use-deposit-balances';
 import { FaucetNotification } from './faucet-notification';
 import { ApproveNotification } from './approve-notification';
 import { usePersistentDeposit } from './use-persistent-deposit';
+import { AssetBalance } from './asset-balance';
 
 interface FormFields {
   asset: string;
@@ -253,9 +254,15 @@ export const DepositForm = ({
               value={selectedAsset?.id}
               hasError={Boolean(errors.asset?.message)}
             >
-              {assets.filter(isAssetTypeERC20).map((a) => (
-                <AssetOption asset={a} key={a.id} />
-              ))}
+              {assets
+                .filter((asset) => isAssetTypeERC20(asset))
+                .map((asset) => (
+                  <AssetOption
+                    asset={asset}
+                    key={asset.id}
+                    balance={<AssetBalance asset={asset} />}
+                  />
+                ))}
             </RichSelect>
           )}
         />
diff --git a/libs/deposits/src/lib/index.ts b/libs/deposits/src/lib/index.ts
index 32152393f..f381470c1 100644
--- a/libs/deposits/src/lib/index.ts
+++ b/libs/deposits/src/lib/index.ts
@@ -1,15 +1,16 @@
 export * from './__generated__/Deposit';
+export * from './asset-balance';
 export * from './deposit-container';
+export * from './deposit-dialog';
 export * from './deposit-form';
 export * from './deposit-limits';
 export * from './deposit-manager';
+export * from './deposits-provider';
 export * from './deposits-table';
 export * from './use-deposit-balances';
-export * from './deposits-provider';
 export * from './use-get-allowance';
 export * from './use-get-balance-of-erc20-token';
 export * from './use-get-deposit-maximum';
 export * from './use-get-deposited-amount';
 export * from './use-submit-approval';
 export * from './use-submit-faucet';
-export * from './deposit-dialog';
diff --git a/libs/deposits/src/lib/use-submit-faucet.ts b/libs/deposits/src/lib/use-submit-faucet.ts
index 05ef3e5b8..3d61596e5 100644
--- a/libs/deposits/src/lib/use-submit-faucet.ts
+++ b/libs/deposits/src/lib/use-submit-faucet.ts
@@ -5,6 +5,7 @@ import {
 } from '@vegaprotocol/web3';
 import { isAssetTypeERC20 } from '@vegaprotocol/utils';
 import type { Asset } from '@vegaprotocol/assets';
+import { useBalancesStore } from '@vegaprotocol/assets';
 import { useEffect, useState } from 'react';
 
 export const useSubmitFaucet = (
@@ -25,8 +26,9 @@ export const useSubmitFaucet = (
   useEffect(() => {
     if (tx?.status === EthTxStatus.Confirmed) {
       getBalances();
+      if (asset) useBalancesStore.getState().refetch(asset.id);
     }
-  }, [tx?.status, getBalances]);
+  }, [tx?.status, getBalances, asset]);
 
   return {
     id,
diff --git a/libs/utils/src/lib/is-asset-erc20.ts b/libs/utils/src/lib/is-asset-erc20.ts
index a7e1de722..4437a01e4 100644
--- a/libs/utils/src/lib/is-asset-erc20.ts
+++ b/libs/utils/src/lib/is-asset-erc20.ts
@@ -7,7 +7,7 @@ interface Asset {
       };
 }
 
-interface ERC20Asset {
+export interface ERC20Asset {
   __typename?: 'Asset';
   source: {
     __typename: 'ERC20';
diff --git a/libs/withdraws/src/lib/asset-balance.tsx b/libs/withdraws/src/lib/asset-balance.tsx
new file mode 100644
index 000000000..8c4cce702
--- /dev/null
+++ b/libs/withdraws/src/lib/asset-balance.tsx
@@ -0,0 +1,30 @@
+import { useAccountBalance } from '@vegaprotocol/accounts';
+import type { AssetFieldsFragment } from '@vegaprotocol/assets';
+import { useBalancesStore } from '@vegaprotocol/assets';
+import { Balance } from '@vegaprotocol/assets';
+import { addDecimal } from '@vegaprotocol/utils';
+import BigNumber from 'bignumber.js';
+import { useEffect } from 'react';
+
+export const AssetBalance = ({ asset }: { asset: AssetFieldsFragment }) => {
+  const [setBalance, getBalance] = useBalancesStore((state) => [
+    state.setBalance,
+    state.getBalance,
+  ]);
+  const { accountBalance, accountDecimals } = useAccountBalance(asset?.id);
+
+  useEffect(() => {
+    const balance =
+      accountBalance && accountDecimals
+        ? new BigNumber(addDecimal(accountBalance, accountDecimals))
+        : undefined;
+    setBalance({ asset, balanceOnVega: balance });
+  }, [accountBalance, accountDecimals, asset, setBalance]);
+
+  return (
+    <Balance
+      balance={getBalance(asset.id)?.balanceOnVega?.toString()}
+      symbol={asset.symbol}
+    />
+  );
+};
diff --git a/libs/withdraws/src/lib/withdraw-form-container.spec.tsx b/libs/withdraws/src/lib/withdraw-form-container.spec.tsx
index 5ca9483e9..a16393b9e 100644
--- a/libs/withdraws/src/lib/withdraw-form-container.spec.tsx
+++ b/libs/withdraws/src/lib/withdraw-form-container.spec.tsx
@@ -1,6 +1,7 @@
 import { render, screen, act, waitFor } from '@testing-library/react';
 import { MockedProvider } from '@apollo/client/testing';
 import type { Account } from '@vegaprotocol/accounts';
+import { useAccountBalance } from '@vegaprotocol/accounts';
 import { WithdrawFormContainer } from './withdraw-form-container';
 import * as Types from '@vegaprotocol/types';
 import { useWeb3React } from '@web3-react/core';
@@ -12,6 +13,7 @@ jest.mock('@vegaprotocol/react-helpers', () => ({
   }),
 }));
 jest.mock('@web3-react/core');
+jest.mock('@vegaprotocol/accounts');
 
 describe('WithdrawFormContainer', () => {
   const props = {
@@ -91,6 +93,10 @@ describe('WithdrawFormContainer', () => {
 
   beforeEach(() => {
     (useWeb3React as jest.Mock).mockReturnValue({ account: MOCK_ETH_ADDRESS });
+    (useAccountBalance as jest.Mock).mockReturnValue({
+      accountBalance: 0,
+      accountDecimals: null,
+    });
   });
   afterEach(() => {
     jest.resetAllMocks();
diff --git a/libs/withdraws/src/lib/withdraw-form.spec.tsx b/libs/withdraws/src/lib/withdraw-form.spec.tsx
index a0e895cbd..d67a7ef99 100644
--- a/libs/withdraws/src/lib/withdraw-form.spec.tsx
+++ b/libs/withdraws/src/lib/withdraw-form.spec.tsx
@@ -8,6 +8,11 @@ import type { Asset } from '@vegaprotocol/assets';
 
 jest.mock('@web3-react/core');
 
+jest.mock('@vegaprotocol/accounts', () => ({
+  ...jest.requireActual('@vegaprotocol/accounts'),
+  useAccountBalance: jest.fn(() => ({ accountBalance: 0, accountDecimals: 0 })),
+}));
+
 const MOCK_ETH_ADDRESS = '0x72c22822A19D20DE7e426fB84aa047399Ddd8853';
 
 let assets: Asset[];
diff --git a/libs/withdraws/src/lib/withdraw-form.tsx b/libs/withdraws/src/lib/withdraw-form.tsx
index 5fec610a2..add36e0d9 100644
--- a/libs/withdraws/src/lib/withdraw-form.tsx
+++ b/libs/withdraws/src/lib/withdraw-form.tsx
@@ -33,6 +33,7 @@ import {
   useWeb3ConnectStore,
   useWeb3Disconnect,
 } from '@vegaprotocol/web3';
+import { AssetBalance } from './asset-balance';
 
 interface FormFields {
   asset: string;
@@ -154,7 +155,11 @@ export const WithdrawForm = ({
         hasError={Boolean(errors.asset?.message)}
       >
         {assets.filter(isAssetTypeERC20).map((a) => (
-          <AssetOption key={a.id} asset={a} />
+          <AssetOption
+            key={a.id}
+            asset={a}
+            balance={<AssetBalance asset={a} />}
+          />
         ))}
       </RichSelect>
     );
diff --git a/libs/withdraws/src/lib/withdraw-manager.spec.tsx b/libs/withdraws/src/lib/withdraw-manager.spec.tsx
index 4790fbe5d..88234ef36 100644
--- a/libs/withdraws/src/lib/withdraw-manager.spec.tsx
+++ b/libs/withdraws/src/lib/withdraw-manager.spec.tsx
@@ -25,6 +25,11 @@ jest.mock('./use-withdraw-asset', () => ({
   useWithdrawAsset: () => withdrawAsset,
 }));
 
+jest.mock('@vegaprotocol/accounts', () => ({
+  ...jest.requireActual('@vegaprotocol/accounts'),
+  useAccountBalance: jest.fn(() => ({ accountBalance: 0, accountDecimals: 0 })),
+}));
+
 jest.mock('@vegaprotocol/web3', () => ({
   ...jest.requireActual('@vegaprotocol/web3'),
   useGetWithdrawThreshold: () => {