[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'
@@ -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/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 | undefined>()
 
-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)
+    balance =
+      cachedBalance !== undefined && !ignoreCache ? cachedBalance : await getWeb3ReadOnly()?.getBalance(safeAddress)
     setNativeBalance(balance)
   }