From 5ad1783e2808b886fac272cd26a07287dbe7adc4 Mon Sep 17 00:00:00 2001
From: Micaela Estabillo <micaela.estabillo@consensys.net>
Date: Tue, 1 Oct 2024 17:18:46 -0700
Subject: [PATCH] chore: fetch latest src token balance

---
 ui/hooks/bridge/useLatestBalance.test.ts      | 88 +++++++++++++++++++
 ui/hooks/bridge/useLatestBalance.ts           | 58 ++++++++++++
 .../prepare-bridge-page.test.tsx.snap         |  8 +-
 .../bridge/prepare/bridge-input-group.tsx     |  8 +-
 .../prepare/prepare-bridge-page.test.tsx      | 11 +++
 5 files changed, 165 insertions(+), 8 deletions(-)
 create mode 100644 ui/hooks/bridge/useLatestBalance.test.ts
 create mode 100644 ui/hooks/bridge/useLatestBalance.ts

diff --git a/ui/hooks/bridge/useLatestBalance.test.ts b/ui/hooks/bridge/useLatestBalance.test.ts
new file mode 100644
index 000000000000..4965b5372f91
--- /dev/null
+++ b/ui/hooks/bridge/useLatestBalance.test.ts
@@ -0,0 +1,88 @@
+import { BigNumber } from 'ethers';
+import { renderHookWithProvider } from '../../../test/lib/render-helpers';
+import { CHAIN_IDS } from '../../../shared/constants/network';
+import { createBridgeMockStore } from '../../../test/jest/mock-store';
+import { zeroAddress } from '../../__mocks__/ethereumjs-util';
+import { createTestProviderTools } from '../../../test/stub/provider';
+import * as tokenutil from '../../../shared/lib/token-util';
+import useLatestBalance from './useLatestBalance';
+
+const mockGetBalance = jest.fn();
+jest.mock('@ethersproject/providers', () => {
+  return {
+    Web3Provider: jest.fn().mockImplementation(() => {
+      return {
+        getBalance: mockGetBalance,
+      };
+    }),
+  };
+});
+
+const mockFetchTokenBalance = jest.spyOn(tokenutil, 'fetchTokenBalance');
+jest.mock('../../../shared/lib/token-util', () => ({
+  ...jest.requireActual('../../../shared/lib/token-util'),
+  fetchTokenBalance: jest.fn(),
+}));
+
+const renderUseLatestBalance = (
+  token: { address: string; decimals?: number | string },
+  chainId: string,
+  mockStoreState: object,
+) =>
+  renderHookWithProvider(
+    () => useLatestBalance(token as any, chainId as any),
+    mockStoreState,
+  );
+
+describe('useLatestBalance', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    const { provider } = createTestProviderTools({
+      networkId: 'Ethereum',
+      chainId: CHAIN_IDS.MAINNET,
+    });
+
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    global.ethereumProvider = provider as any;
+  });
+
+  it('returns formattedBalance for native asset in current chain', async () => {
+    mockGetBalance.mockResolvedValue(BigNumber.from('1000000000000000000'));
+
+    const { result, waitForNextUpdate } = renderUseLatestBalance(
+      { address: zeroAddress() },
+      CHAIN_IDS.MAINNET,
+      createBridgeMockStore(),
+    );
+
+    await waitForNextUpdate();
+    expect(result.current.formattedBalance).toStrictEqual('1');
+
+    expect(mockGetBalance).toHaveBeenCalledTimes(1);
+    expect(mockGetBalance).toHaveBeenCalledWith(
+      '0x0dcd5d886577d5081b0c52e242ef29e70be3e7bc',
+    );
+    expect(mockFetchTokenBalance).toHaveBeenCalledTimes(0);
+  });
+
+  it('returns formattedBalance for ERC20 asset in current chain', async () => {
+    mockFetchTokenBalance.mockResolvedValueOnce(BigNumber.from('15390000'));
+
+    const { result, waitForNextUpdate } = renderUseLatestBalance(
+      { address: '0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48', decimals: '6' },
+      CHAIN_IDS.MAINNET,
+      createBridgeMockStore(),
+    );
+
+    await waitForNextUpdate();
+    expect(result.current.formattedBalance).toStrictEqual('15.39');
+
+    expect(mockFetchTokenBalance).toHaveBeenCalledTimes(1);
+    expect(mockFetchTokenBalance).toHaveBeenCalledWith(
+      '0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48',
+      '0x0dcd5d886577d5081b0c52e242ef29e70be3e7bc',
+      global.ethereumProvider,
+    );
+    expect(mockGetBalance).toHaveBeenCalledTimes(0);
+  });
+});
diff --git a/ui/hooks/bridge/useLatestBalance.ts b/ui/hooks/bridge/useLatestBalance.ts
new file mode 100644
index 000000000000..e54529b8cb20
--- /dev/null
+++ b/ui/hooks/bridge/useLatestBalance.ts
@@ -0,0 +1,58 @@
+import { useSelector } from 'react-redux';
+import { zeroAddress } from 'ethereumjs-util';
+import { Web3Provider } from '@ethersproject/providers';
+import { Hex } from '@metamask/utils';
+import { BigNumber } from 'ethers';
+import { Numeric } from '../../../shared/modules/Numeric';
+import { DEFAULT_PRECISION } from '../useCurrencyDisplay';
+import { fetchTokenBalance } from '../../../shared/lib/token-util';
+import {
+  getCurrentChainId,
+  getSelectedInternalAccount,
+  SwapsEthToken,
+} from '../../selectors';
+import { SwapsTokenObject } from '../../../shared/constants/swaps';
+import { useAsyncResult } from '../useAsyncResult';
+
+/**
+ * Custom hook to fetch and format the latest balance of a given token or native asset.
+ *
+ * @param token - The token object for which the balance is to be fetched. Can be null.
+ * @param chainId - The chain ID to be used for fetching the balance. Optional.
+ * @returns An object containing the formatted balance as a string.
+ */
+const useLatestBalance = (
+  token: SwapsTokenObject | SwapsEthToken | null,
+  chainId?: Hex,
+) => {
+  const { address: selectedAddress } = useSelector(getSelectedInternalAccount);
+  const currentChainId = useSelector(getCurrentChainId);
+
+  const { value: latestBalance } = useAsyncResult<BigNumber>(async () => {
+    if (token && chainId && currentChainId === chainId) {
+      if (!token.address || token.address === zeroAddress()) {
+        const ethersProvider = new Web3Provider(global.ethereumProvider);
+        return await ethersProvider.getBalance(selectedAddress);
+      }
+      return await fetchTokenBalance(
+        token.address,
+        selectedAddress,
+        global.ethereumProvider,
+      );
+    }
+    // TODO implement fetching balance on non-active chain and update unit test
+    return undefined;
+  }, [token, selectedAddress, global.ethereumProvider]);
+
+  return {
+    formattedBalance:
+      token && latestBalance
+        ? Numeric.from(latestBalance.toString(), 10)
+            .shiftedBy(token?.decimals ? Number(token.decimals) : 18)
+            .round(DEFAULT_PRECISION)
+            .toString()
+        : undefined,
+  };
+};
+
+export default useLatestBalance;
diff --git a/ui/pages/bridge/prepare/__snapshots__/prepare-bridge-page.test.tsx.snap b/ui/pages/bridge/prepare/__snapshots__/prepare-bridge-page.test.tsx.snap
index 9343ecf99701..03af05d50d09 100644
--- a/ui/pages/bridge/prepare/__snapshots__/prepare-bridge-page.test.tsx.snap
+++ b/ui/pages/bridge/prepare/__snapshots__/prepare-bridge-page.test.tsx.snap
@@ -99,9 +99,7 @@ exports[`PrepareBridgePage should render the component 1`] = `
           <p
             class="mm-box mm-text mm-text--body-md mm-box--color-text-default"
           >
-            Balance
-            : 
-            "0"
+             
           </p>
           <div
             class="mm-box currency-display-component mm-box--display-flex mm-box--flex-wrap-wrap mm-box--align-items-center"
@@ -184,9 +182,7 @@ exports[`PrepareBridgePage should render the component 1`] = `
           <p
             class="mm-box mm-text mm-text--body-md mm-box--color-text-default"
           >
-            Balance
-            : 
-            "0"
+             
           </p>
           <div
             class="mm-box currency-display-component mm-box--display-flex mm-box--flex-wrap-wrap mm-box--align-items-center"
diff --git a/ui/pages/bridge/prepare/bridge-input-group.tsx b/ui/pages/bridge/prepare/bridge-input-group.tsx
index 32ee541870bd..173580a917d4 100644
--- a/ui/pages/bridge/prepare/bridge-input-group.tsx
+++ b/ui/pages/bridge/prepare/bridge-input-group.tsx
@@ -26,6 +26,7 @@ import {
   CHAIN_ID_TO_CURRENCY_SYMBOL_MAP,
   CHAIN_ID_TOKEN_IMAGE_MAP,
 } from '../../../../shared/constants/network';
+import useLatestBalance from '../../../hooks/bridge/useLatestBalance';
 
 const generateAssetFromToken = (
   chainId: Hex,
@@ -91,7 +92,10 @@ export const BridgeInputGroup = ({
     true,
   );
 
-  const latestBalance = '0';
+  const { formattedBalance } = useLatestBalance(
+    token,
+    networkProps?.network?.chainId,
+  );
 
   return (
     <Box className={className}>
@@ -132,7 +136,7 @@ export const BridgeInputGroup = ({
       </Box>
       <Box className="prepare-bridge-page__amounts-row">
         <Text>
-          {t('balance')}: {JSON.stringify(latestBalance)}
+          {formattedBalance ? `${t('balance')}: ${formattedBalance}` : ' '}
         </Text>
         <CurrencyDisplay
           currency="usd"
diff --git a/ui/pages/bridge/prepare/prepare-bridge-page.test.tsx b/ui/pages/bridge/prepare/prepare-bridge-page.test.tsx
index f9db9b1213a7..8b6099eff709 100644
--- a/ui/pages/bridge/prepare/prepare-bridge-page.test.tsx
+++ b/ui/pages/bridge/prepare/prepare-bridge-page.test.tsx
@@ -3,9 +3,20 @@ import { renderWithProvider } from '../../../../test/jest';
 import configureStore from '../../../store/store';
 import { createBridgeMockStore } from '../../../../test/jest/mock-store';
 import { CHAIN_IDS } from '../../../../shared/constants/network';
+import { createTestProviderTools } from '../../../../test/stub/provider';
 import PrepareBridgePage from './prepare-bridge-page';
 
 describe('PrepareBridgePage', () => {
+  beforeAll(() => {
+    const { provider } = createTestProviderTools({
+      networkId: 'Ethereum',
+      chainId: CHAIN_IDS.MAINNET,
+    });
+
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    global.ethereumProvider = provider as any;
+  });
+
   it('should render the component', () => {
     const mockStore = createBridgeMockStore(
       {