fix: balance caching

ava-labs · Dec 6, 2024 · 0284bac · 0284bac
1 parent 3bc00fb
commit 0284bac
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 2 deletions.
diff --git a/src/background/services/balances/BalanceAggregatorService.test.ts b/src/background/services/balances/BalanceAggregatorService.test.ts
@@ -218,6 +218,63 @@ describe('src/background/services/balances/BalanceAggregatorService.ts', () => {
       );
     });
 
+    it('only updates the balances of the requested accounts', async () => {
+      // Mock the existing balances for other accounts
+      (balancesServiceMock.getBalancesForNetwork as jest.Mock).mockReset();
+
+      balancesServiceMock.getBalancesForNetwork
+        .mockResolvedValueOnce({
+          [account2.addressC]: {
+            [networkToken1.symbol]: network1TokenBalance,
+          },
+        })
+        .mockResolvedValueOnce({
+          [account1.addressC]: {
+            [networkToken1.symbol]: network1TokenBalance,
+          },
+        });
+
+      // Get balances for the `account2` so they get cached
+      await service.getBalancesForNetworks(
+        [network1.chainId],
+        [account2],
+        [TokenType.NATIVE, TokenType.ERC20, TokenType.ERC721]
+      );
+
+      expect(balancesServiceMock.getBalancesForNetwork).toHaveBeenCalledTimes(
+        1
+      );
+
+      expect(service.balances).toEqual({
+        [network1.chainId]: {
+          [account2.addressC]: {
+            [networkToken1.symbol]: network1TokenBalance,
+          },
+        },
+      });
+
+      // Now get the balances for the first account and verify the `account2` balances are kept in cache
+      await service.getBalancesForNetworks(
+        [network1.chainId],
+        [account1],
+        [TokenType.NATIVE, TokenType.ERC20, TokenType.ERC721]
+      );
+
+      expect(balancesServiceMock.getBalancesForNetwork).toHaveBeenCalledTimes(
+        2
+      );
+      expect(service.balances).toEqual({
+        [network1.chainId]: {
+          [account1.addressC]: {
+            [networkToken1.symbol]: network1TokenBalance,
+          },
+          [account2.addressC]: {
+            [networkToken1.symbol]: network1TokenBalance,
+          },
+        },
+      });
+    });
+
     it('can fetch the balance for multiple networks and one account', async () => {
       const balances = await service.getBalancesForNetworks(
         [network1.chainId, network2.chainId],

diff --git a/src/background/services/balances/BalanceAggregatorService.ts b/src/background/services/balances/BalanceAggregatorService.ts
@@ -6,7 +6,7 @@ import { BalancesService } from './BalancesService';
 import { NetworkService } from '../network/NetworkService';
 import { EventEmitter } from 'events';
 import * as Sentry from '@sentry/browser';
-import { isEqual, pick } from 'lodash';
+import { isEqual, omit, pick } from 'lodash';
 
 import { LockService } from '../lock/LockService';
 import { StorageService } from '../storage/StorageService';
@@ -144,7 +144,7 @@ export class BalanceAggregatorService implements OnLock, OnUnlock {
       for (const [chainId, chainBalances] of freshData) {
         for (const [address, addressBalance] of Object.entries(chainBalances)) {
           aggregatedBalances[chainId] = {
-            ...aggregatedBalances[chainId],
+            ...omit(aggregatedBalances[chainId], address), // Keep cached balances for other accounts
             ...chainBalances,
             [address]: addressBalance,
           };