[Counterfactual] Display option to refresh balances

src/components/balances/AssetsTable/index.tsx

@@ -1,3 +1,4 @@
+import CheckBalance from '@/features/counterfactual/CheckBalance'
 import { type ReactElement, useMemo, useContext } from 'react'
 import { Button, Tooltip, Typography, SvgIcon, IconButton, Box, Checkbox, Skeleton } from '@mui/material'
 import type { TokenInfo } from '@safe-global/safe-gateway-typescript-sdk'
@@ -138,7 +139,7 @@ const AssetsTable = ({
     [hiddenAssets, balances.items, showHiddenAssets],
   )
 
-  const hasNoAssets = balances.items.length === 1 && balances.items[0].balance === '0'
+  const hasNoAssets = !loading && balances.items.length === 1 && balances.items[0].balance === '0'
 
   const selectedAssetCount = visibleAssets?.filter((item) => isAssetSelected(item.tokenInfo.address)).length || 0
 
@@ -257,6 +258,8 @@ const AssetsTable = ({
           <EnhancedTable rows={rows} headCells={headCells} />
         </div>
       )}
+
+      <CheckBalance />
     </>
   )
 }

src/components/common/CooldownButton/index.test.tsx

@@ -0,0 +1,77 @@
+import { render, waitFor } from '@/tests/test-utils'
+import CooldownButton from './index'
+
+describe('CooldownButton', () => {
+  beforeAll(() => {
+    jest.useFakeTimers()
+  })
+
+  afterAll(() => {
+    jest.useRealTimers()
+  })
+  it('should be disabled initially if startDisabled is set and become enabled after <cooldown> seconds', async () => {
+    const onClickEvent = jest.fn()
+    const result = render(
+      <CooldownButton cooldown={30} onClick={onClickEvent} startDisabled>
+        Try again
+      </CooldownButton>,
+    )
+
+    expect(result.getByRole('button')).toBeDisabled()
+    expect(result.getByText('Try again in 30s')).toBeVisible()
+
+    jest.advanceTimersByTime(10_000)
+
+    await waitFor(() => {
+      expect(result.getByRole('button')).toBeDisabled()
+      expect(result.getByText('Try again in 20s')).toBeVisible()
+    })
+
+    jest.advanceTimersByTime(5_000)
+
+    await waitFor(() => {
+      expect(result.getByRole('button')).toBeDisabled()
+      expect(result.getByText('Try again in 15s')).toBeVisible()
+    })
+
+    jest.advanceTimersByTime(15_000)
+
+    await waitFor(() => {
+      expect(result.getByRole('button')).toBeEnabled()
+    })
+    result.getByRole('button').click()
+
+    expect(onClickEvent).toHaveBeenCalledTimes(1)
+    await waitFor(() => {
+      expect(result.getByRole('button')).toBeDisabled()
+    })
+  })
+
+  it('should be enabled initially if startDisabled is not set and become disabled after click', async () => {
+    const onClickEvent = jest.fn()
+    const result = render(
+      <CooldownButton cooldown={30} onClick={onClickEvent}>
+        Try again
+      </CooldownButton>,
+    )
+
+    expect(result.getByRole('button')).toBeEnabled()
+    result.getByRole('button').click()
+
+    expect(onClickEvent).toHaveBeenCalledTimes(1)
+
+    await waitFor(() => {
+      expect(result.getByRole('button')).toBeDisabled()
+      expect(result.getByText('Try again in 30s')).toBeVisible()
+    })
+
+    jest.advanceTimersByTime(30_000)
+
+    await waitFor(() => {
+      expect(result.getByRole('button')).toBeEnabled()
+    })
+    result.getByRole('button').click()
+
+    expect(onClickEvent).toHaveBeenCalledTimes(2)
+  })
+})

src/components/common/CooldownButton/index.tsx

@@ -0,0 +1,48 @@
+import { Button } from '@mui/material'
+import { useState, useCallback, useEffect, type ReactNode } from 'react'
+
+// TODO: Extract into a hook so it can be reused for links and not just buttons
+const CooldownButton = ({
+  onClick,
+  cooldown,
+  startDisabled = false,
+  children,
+}: {
+  onClick: () => void
+  startDisabled?: boolean
+  cooldown: number // Cooldown in seconds
+  children: ReactNode
+}) => {
+  const [remainingSeconds, setRemainingSeconds] = useState(startDisabled ? cooldown : 0)
+  const [lastSendTime, setLastSendTime] = useState(startDisabled ? Date.now() : 0)
+
+  const adjustSeconds = useCallback(() => {
+    const remainingCoolDownSeconds = Math.max(0, cooldown * 1000 - (Date.now() - lastSendTime)) / 1000
+    setRemainingSeconds(remainingCoolDownSeconds)
+  }, [cooldown, lastSendTime])
+
+  useEffect(() => {
+    // Counter for progress
+    const interval = setInterval(adjustSeconds, 1000)
+    return () => clearInterval(interval)
+  }, [adjustSeconds])
+
+  const handleClick = () => {
+    setLastSendTime(Date.now())
+    setRemainingSeconds(cooldown)
+    onClick()
+  }
+
+  const isDisabled = remainingSeconds > 0
+
+  return (
+    <Button onClick={handleClick} variant="contained" size="small" disabled={isDisabled}>
+      <span>
+        {children}
+        {remainingSeconds > 0 && ` in ${Math.floor(remainingSeconds)}s`}
+      </span>
+    </Button>
+  )
+}
+
+export default CooldownButton

src/features/counterfactual/CheckBalance.tsx

@@ -0,0 +1,43 @@
+import CooldownButton from '@/components/common/CooldownButton'
+import ExternalLink from '@/components/common/ExternalLink'
+import { getCounterfactualBalance } from '@/features/counterfactual/utils'
+import { useCurrentChain } from '@/hooks/useChains'
+import useSafeInfo from '@/hooks/useSafeInfo'
+import { useWeb3 } from '@/hooks/wallets/web3'
+import { getBlockExplorerLink } from '@/utils/chains'
+import { Box, Typography } from '@mui/material'
+
+const CheckBalance = () => {
+  const { safe, safeAddress } = useSafeInfo()
+  const chain = useCurrentChain()
+  const provider = useWeb3()
+
+  if (safe.deployed) return null
+
+  const handleBalanceRefresh = async () => {
+    void getCounterfactualBalance(safeAddress, provider, chain, true)
+  }
+
+  const blockExplorerLink = chain ? getBlockExplorerLink(chain, safeAddress) : undefined
+
+  return (
+    <Box textAlign="center" p={3}>
+      <Typography variant="h4" fontWeight="bold" mb={2}>
+        Don&apos;t see your tokens?
+      </Typography>
+      <Box display="flex" flexDirection="column" gap={1} alignItems="center">
+        <CooldownButton cooldown={15} onClick={handleBalanceRefresh}>
+          Refresh balance
+        </CooldownButton>
+        <Typography>or</Typography>
+        {blockExplorerLink && (
+          <Typography>
+            check on <ExternalLink href={blockExplorerLink.href}>Block Explorer</ExternalLink>
+          </Typography>
+        )}
+      </Box>
+    </Box>
+  )
+}
+
+export default CheckBalance

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

@@ -1,4 +1,9 @@
-import { getCounterfactualBalance, getUndeployedSafeInfo } from '@/features/counterfactual/utils'
+import {
+  getCounterfactualBalance,
+  getNativeBalance,
+  getUndeployedSafeInfo,
+  setNativeBalance,
+} from '@/features/counterfactual/utils'
 import * as web3 from '@/hooks/wallets/web3'
 import { chainBuilder } from '@/tests/builders/chains'
 import { faker } from '@faker-js/faker'
@@ -32,6 +37,9 @@ describe('Counterfactual utils', () => {
   })
 
   describe('getCounterfactualBalance', () => {
+    const mockSafeAddress = faker.finance.ethereumAddress()
+    const mockChain = chainBuilder().build()
+
     beforeEach(() => {
       jest.clearAllMocks()
     })
@@ -41,13 +49,13 @@ describe('Counterfactual utils', () => {
       const mockReadOnlyProvider = {
         getBalance: jest.fn(() => Promise.resolve(mockBalance)),
       } as unknown as JsonRpcProvider
-      jest.spyOn(web3, 'getWeb3ReadOnly').mockImplementation(() => mockReadOnlyProvider)
+      jest.spyOn(web3, 'getWeb3ReadOnly').mockImplementationOnce(() => mockReadOnlyProvider)
 
-      const mockSafeAddress = faker.finance.ethereumAddress()
-      const mockChain = chainBuilder().build()
       const result = await getCounterfactualBalance(mockSafeAddress, undefined, mockChain)
+      const nativeBalanceCache = getNativeBalance()
 
-      expect(mockReadOnlyProvider.getBalance).toHaveBeenCalled()
+      expect(mockReadOnlyProvider.getBalance).toHaveBeenCalledTimes(1)
+      expect(nativeBalanceCache).toEqual(mockBalance)
       expect(result).toEqual({
         fiatTotal: '0',
         items: [
@@ -63,10 +71,14 @@ describe('Counterfactual utils', () => {
           },
         ],
       })
+
+      // Should use the cache now
+      const newResult = await getCounterfactualBalance(mockSafeAddress, undefined, mockChain)
+      expect(mockReadOnlyProvider.getBalance).toHaveBeenCalledTimes(1)
+      expect(newResult?.items[0].balance).toEqual('123')
     })
 
     it('should return undefined if there is no chain info', async () => {
-      const mockSafeAddress = faker.finance.ethereumAddress()
       const mockProvider = { getBalance: jest.fn(() => Promise.resolve(1n)) } as unknown as BrowserProvider
 
       const result = await getCounterfactualBalance(mockSafeAddress, mockProvider, undefined)
@@ -75,10 +87,8 @@ describe('Counterfactual utils', () => {
     })
 
     it('should return the native balance', async () => {
-      const mockSafeAddress = faker.finance.ethereumAddress()
       const mockBalance = 1000000n
       const mockProvider = { getBalance: jest.fn(() => Promise.resolve(mockBalance)) } as unknown as BrowserProvider
-      const mockChain = chainBuilder().build()
 
       const result = await getCounterfactualBalance(mockSafeAddress, mockProvider, mockChain)
 
@@ -99,5 +109,40 @@ describe('Counterfactual utils', () => {
         ],
       })
     })
+
+    it('should not use the cache if the ignoreCache flag is passed', async () => {
+      const mockBalance = 123n
+      const mockReadOnlyProvider = {
+        getBalance: jest.fn(() => Promise.resolve(mockBalance)),
+      } as unknown as JsonRpcProvider
+      jest.spyOn(web3, 'getWeb3ReadOnly').mockImplementationOnce(() => mockReadOnlyProvider)
+
+      // Set local cache
+      const mockCacheBalance = 10n
+      setNativeBalance(mockCacheBalance)
+      const nativeBalanceCache = getNativeBalance()
+      expect(nativeBalanceCache).toEqual(mockCacheBalance)
+
+      // Call function and ignore cache
+      const result = await getCounterfactualBalance(mockSafeAddress, undefined, mockChain, true)
+
+      expect(mockReadOnlyProvider.getBalance).toHaveBeenCalled()
+      expect(result?.items[0].balance).not.toEqual(mockCacheBalance)
+      expect(result).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

@@ -120,9 +120,14 @@ export const deploySafeAndExecuteTx = async (
   return dispatchTxExecutionAndDeploySafe(safeTx, txOptions, onboard, chainId, onSuccess)
 }
 
-const { getStore: getNativeBalance, setStore: setNativeBalance } = new ExternalStore<bigint | undefined>()
+export const { getStore: getNativeBalance, setStore: setNativeBalance } = new ExternalStore<bigint>(0n)
 
-export const getCounterfactualBalance = async (safeAddress: string, provider?: BrowserProvider, chain?: ChainInfo) => {
+export const getCounterfactualBalance = async (
+  safeAddress: string,
+  provider?: BrowserProvider,
+  chain?: ChainInfo,
+  ignoreCache?: boolean,
+) => {
   let balance: bigint | undefined
 
   if (!chain) return undefined
@@ -133,7 +138,8 @@ export const getCounterfactualBalance = async (safeAddress: string, provider?: B
     balance = await provider.getBalance(safeAddress)
   } else {
     const cachedBalance = getNativeBalance()
-    balance = cachedBalance !== undefined ? cachedBalance : await getWeb3ReadOnly()?.getBalance(safeAddress)
+    const useCache = cachedBalance && cachedBalance > 0n && !ignoreCache
+    balance = useCache ? cachedBalance : (await getWeb3ReadOnly()?.getBalance(safeAddress)) || 0n
     setNativeBalance(balance)
   }