From f82d825c55f39446799bf54ab6d2259d6b93c990 Mon Sep 17 00:00:00 2001
From: Manuel Gellfart <manu@safe.global>
Date: Thu, 14 Dec 2023 12:10:49 +0100
Subject: [PATCH] fix: poll wallet balance (#3011)

---
 .../__tests__/useWalletBalance.test.ts        | 48 +++++++++++++++++++
 src/hooks/wallets/useWalletBalance.ts         | 26 +++++++---
 2 files changed, 68 insertions(+), 6 deletions(-)
 create mode 100644 src/hooks/wallets/__tests__/useWalletBalance.test.ts

diff --git a/src/hooks/wallets/__tests__/useWalletBalance.test.ts b/src/hooks/wallets/__tests__/useWalletBalance.test.ts
new file mode 100644
index 0000000000..2b547d8974
--- /dev/null
+++ b/src/hooks/wallets/__tests__/useWalletBalance.test.ts
@@ -0,0 +1,48 @@
+import { act, renderHook, waitFor } from '@testing-library/react'
+import useWalletBalance from '../useWalletBalance'
+import * as web3 from '@/hooks/wallets/web3'
+import * as useWallet from '../useWallet'
+
+import { BigNumber } from 'ethers'
+import type { JsonRpcProvider } from '@ethersproject/providers'
+import { connectedWalletBuilder } from '@/tests/builders/wallet'
+import { POLLING_INTERVAL } from '@/config/constants'
+
+describe('useWalletBalance', () => {
+  beforeAll(() => {
+    jest.useFakeTimers()
+  })
+
+  afterAll(() => {
+    jest.useRealTimers()
+  })
+  it('should poll balance after interval passed', async () => {
+    jest.spyOn(web3, 'useWeb3ReadOnly').mockReturnValue({
+      getBalance: () => Promise.resolve(BigNumber.from(420)),
+    } as unknown as JsonRpcProvider)
+    jest.spyOn(useWallet, 'default').mockReturnValue(connectedWalletBuilder().build())
+
+    const { result } = renderHook(useWalletBalance)
+
+    await waitFor(() => {
+      expect(result.current[0]).toEqual(BigNumber.from(420))
+    })
+
+    jest.spyOn(web3, 'useWeb3ReadOnly').mockReturnValue({
+      getBalance: () => Promise.resolve(BigNumber.from(69)),
+    } as unknown as JsonRpcProvider)
+
+    // We only poll after the interval passed
+    await waitFor(() => {
+      expect(result.current[0]).toEqual(BigNumber.from(420))
+    })
+
+    act(() => {
+      jest.advanceTimersByTime(POLLING_INTERVAL)
+    })
+
+    await waitFor(() => {
+      expect(result.current[0]).toEqual(BigNumber.from(69))
+    })
+  })
+})
diff --git a/src/hooks/wallets/useWalletBalance.ts b/src/hooks/wallets/useWalletBalance.ts
index 9f7551ab28..1224ac41dd 100644
--- a/src/hooks/wallets/useWalletBalance.ts
+++ b/src/hooks/wallets/useWalletBalance.ts
@@ -2,18 +2,32 @@ import useAsync, { type AsyncResult } from '../useAsync'
 import useWallet from './useWallet'
 import { useWeb3ReadOnly } from '@/hooks/wallets/web3'
 import { type BigNumber } from 'ethers'
+import useIntervalCounter from '../useIntervalCounter'
+import { POLLING_INTERVAL } from '@/config/constants'
+import { useEffect } from 'react'
 
 const useWalletBalance = (): AsyncResult<BigNumber | undefined> => {
   const web3ReadOnly = useWeb3ReadOnly()
   const wallet = useWallet()
+  const [pollingInterval, resetPolling] = useIntervalCounter(POLLING_INTERVAL)
 
-  return useAsync<BigNumber | undefined>(() => {
-    if (!wallet || !web3ReadOnly) {
-      return undefined
-    }
+  // Reset polling if new wallet connects
+  useEffect(() => {
+    resetPolling()
+  }, [resetPolling, wallet, web3ReadOnly])
 
-    return web3ReadOnly.getBalance(wallet.address, 'latest')
-  }, [wallet, web3ReadOnly])
+  return useAsync<BigNumber | undefined>(
+    () => {
+      if (!wallet || !web3ReadOnly) {
+        return undefined
+      }
+
+      return web3ReadOnly.getBalance(wallet.address, 'latest')
+    },
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+    [wallet, web3ReadOnly, pollingInterval],
+    false,
+  )
 }
 
 export default useWalletBalance