From fddcad46c05ec11334cd1e8cffcc7ddbaba987bd Mon Sep 17 00:00:00 2001
From: Usame Algan <usame@safe.global>
Date: Sat, 10 Feb 2024 10:46:13 +0100
Subject: [PATCH] fix: Adjust cache retrieval logic and write a test for it

---
 .../counterfactual/__tests__/utils.test.ts    | 61 ++++++++++++++++---
 src/features/counterfactual/utils.ts          |  6 +-
 2 files changed, 56 insertions(+), 11 deletions(-)

diff --git a/src/features/counterfactual/__tests__/utils.test.ts b/src/features/counterfactual/__tests__/utils.test.ts
index 589372a73e..b13a4ef13b 100644
--- a/src/features/counterfactual/__tests__/utils.test.ts
+++ b/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',
+          },
+        ],
+      })
+    })
   })
 })
diff --git a/src/features/counterfactual/utils.ts b/src/features/counterfactual/utils.ts
index df71f315b4..8c54a9a7aa 100644
--- a/src/features/counterfactual/utils.ts
+++ b/src/features/counterfactual/utils.ts
@@ -120,7 +120,7 @@ export const deploySafeAndExecuteTx = async (
   return dispatchTxExecutionAndDeploySafe(safeTx, txOptions, onboard, chainId, onSuccess)
 }
 
-export 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,
@@ -138,8 +138,8 @@ export const getCounterfactualBalance = async (
     balance = await provider.getBalance(safeAddress)
   } else {
     const cachedBalance = getNativeBalance()
-    balance =
-      cachedBalance !== undefined && !ignoreCache ? cachedBalance : await getWeb3ReadOnly()?.getBalance(safeAddress)
+    const useCache = cachedBalance && cachedBalance > 0n && !ignoreCache
+    balance = useCache ? cachedBalance : (await getWeb3ReadOnly()?.getBalance(safeAddress)) || 0n
     setNativeBalance(balance)
   }