fix: Replace useCounterfactualBalance hook with get function and writ…

…e tests

src/features/counterfactual/__tests__/utils.test.ts

@@ -1,6 +1,10 @@
-import { getUndeployedSafeInfo } from '@/features/counterfactual/utils'
+import { getCounterfactualBalance, getUndeployedSafeInfo } from '@/features/counterfactual/utils'
+import { chainBuilder } from '@/tests/builders/chains'
 import { faker } from '@faker-js/faker'
 import type { PredictedSafeProps } from '@safe-global/protocol-kit'
+import { ZERO_ADDRESS } from '@safe-global/protocol-kit/dist/src/utils/constants'
+import { TokenType } from '@safe-global/safe-gateway-typescript-sdk'
+import { BrowserProvider, type Eip1193Provider } from 'ethers'
 
 describe('Counterfactual utils', () => {
   describe('getUndeployedSafeInfo', () => {
@@ -25,4 +29,55 @@ describe('Counterfactual utils', () => {
       expect(result.owners[0].value).toEqual(undeployedSafe.safeAccountConfig.owners[0])
     })
   })
+
+  describe('getCounterfactualBalance', () => {
+    beforeEach(() => {
+      jest.clearAllMocks()
+    })
+
+    it('should return undefined if there is no provider', () => {
+      const mockSafeAddress = faker.finance.ethereumAddress()
+      const mockChain = chainBuilder().build()
+      const result = getCounterfactualBalance(mockSafeAddress, undefined, mockChain)
+
+      expect(result).resolves.toBeUndefined()
+    })
+
+    it('should return undefined if there is no chain info', () => {
+      const mockSafeAddress = faker.finance.ethereumAddress()
+      const mockProvider = new BrowserProvider(jest.fn() as unknown as Eip1193Provider)
+      mockProvider.getBalance = jest.fn(() => Promise.resolve(1n))
+
+      const result = getCounterfactualBalance(mockSafeAddress, mockProvider, undefined)
+
+      expect(result).resolves.toBeUndefined()
+    })
+
+    it('should return the native balance', () => {
+      const mockSafeAddress = faker.finance.ethereumAddress()
+      const mockProvider = new BrowserProvider(jest.fn() as unknown as Eip1193Provider)
+      const mockChain = chainBuilder().build()
+      const mockBalance = 1000000n
+
+      mockProvider.getBalance = jest.fn(() => Promise.resolve(mockBalance))
+
+      const result = getCounterfactualBalance(mockSafeAddress, mockProvider, mockChain)
+
+      expect(result).resolves.toEqual({
+        fiatTotal: '0',
+        items: [
+          {
+            tokenInfo: {
+              type: TokenType.NATIVE_TOKEN,
+              address: ZERO_ADDRESS,
+              ...mockChain.nativeCurrency,
+            },
+            balance: mockBalance.toString(),
+            fiatBalance: '0',
+            fiatConversion: '0',
+          },
+        ],
+      })
+    })
+  })
 })

src/features/counterfactual/utils.ts

@@ -1,7 +1,14 @@
 import { LATEST_SAFE_VERSION } from '@/config/constants'
 import { defaultSafeInfo } from '@/store/safeInfoSlice'
 import type { PredictedSafeProps } from '@safe-global/protocol-kit'
-import { ImplementationVersionState } from '@safe-global/safe-gateway-typescript-sdk'
+import { ZERO_ADDRESS } from '@safe-global/protocol-kit/dist/src/utils/constants'
+import {
+  type ChainInfo,
+  ImplementationVersionState,
+  type SafeBalanceResponse,
+  TokenType,
+} from '@safe-global/safe-gateway-typescript-sdk'
+import type { BrowserProvider } from 'ethers'
 
 export const getUndeployedSafeInfo = (undeployedSafe: PredictedSafeProps, address: string, chainId: string) => {
   return Promise.resolve({
@@ -17,3 +24,25 @@ export const getUndeployedSafeInfo = (undeployedSafe: PredictedSafeProps, addres
     deployed: false,
   })
 }
+
+export const getCounterfactualBalance = async (safeAddress: string, provider?: BrowserProvider, chain?: ChainInfo) => {
+  const balance = await provider?.getBalance(safeAddress)
+
+  if (!balance || !chain) return
+
+  return <SafeBalanceResponse>{
+    fiatTotal: '0',
+    items: [
+      {
+        tokenInfo: {
+          type: TokenType.NATIVE_TOKEN,
+          address: ZERO_ADDRESS,
+          ...chain?.nativeCurrency,
+        },
+        balance: balance.toString(),
+        fiatBalance: '0',
+        fiatConversion: '0',
+      },
+    ],
+  }
+}

src/hooks/loadables/useLoadBalances.ts

@@ -1,4 +1,5 @@
-import useCounterfactualBalance from '@/features/counterfactual/hooks/useCounterfactualBalance'
+import { getCounterfactualBalance } from '@/features/counterfactual/utils'
+import { useWeb3 } from '@/hooks/wallets/web3'
 import { useEffect, useMemo } from 'react'
 import { getBalances, type SafeBalanceResponse } from '@safe-global/safe-gateway-typescript-sdk'
 import { useAppSelector } from '@/store'
@@ -28,23 +29,25 @@ export const useLoadBalances = (): AsyncResult<SafeBalanceResponse> => {
   const currency = useAppSelector(selectCurrency)
   const isTrustedTokenList = useTokenListSetting()
   const { safe, safeAddress } = useSafeInfo()
-  const [balance] = useCounterfactualBalance(safeAddress, pollCount)
+  const web3 = useWeb3()
+  const chain = useCurrentChain()
   const chainId = safe.chainId
 
   // Re-fetch assets when the entire SafeInfo updates
-  const [data, error, loading] = useAsync<SafeBalanceResponse>(
+  const [data, error, loading] = useAsync<SafeBalanceResponse | undefined>(
     () => {
       if (!chainId || !safeAddress || isTrustedTokenList === undefined) return
+
       if (!safe.deployed) {
-        return balance ? Promise.resolve(balance) : undefined
+        return getCounterfactualBalance(safeAddress, web3, chain)
       }
 
       return getBalances(chainId, safeAddress, currency, {
         trusted: isTrustedTokenList,
       })
     },
     // eslint-disable-next-line react-hooks/exhaustive-deps
-    [safeAddress, chainId, currency, isTrustedTokenList, pollCount, safe.deployed],
+    [safeAddress, chainId, currency, isTrustedTokenList, pollCount, safe.deployed, web3, chain],
     false, // don't clear data between polls
   )