fix: should throw error if account api fail

Author: salimtb

packages/assets-controllers/src/TokenBalancesController.test.ts

@@ -1,5 +1,6 @@
 import { Messenger } from '@metamask/base-controller';
 import { toHex } from '@metamask/controller-utils';
+import * as controllerUtils from '@metamask/controller-utils';
 import type { InternalAccount } from '@metamask/keyring-internal-api';
 import type { NetworkState } from '@metamask/network-controller';
 import type { PreferencesState } from '@metamask/preferences-controller';
@@ -1794,27 +1795,26 @@ describe('TokenBalancesController', () => {
       // Spy on console.warn
       const consoleWarnSpy = jest.spyOn(console, 'warn').mockImplementation();
 
-      // Mock the RPC balance fetcher to throw an error
-      const mockMulticall = jest
-        .spyOn(multicall, 'getTokenBalancesForMultipleAddresses')
-        .mockRejectedValue(mockError);
+      const { controller } = setupController();
 
-      const { controller } = setupController({
-        config: { useAccountsAPI: false }, // Force use of RPC fetcher only
-        tokens: {
-          allTokens: {
-            [chainId]: {
-              '0x1234567890123456789012345678901234567890': [
-                {
-                  address: '0x6B175474E89094C44Da98b954EedeAC495271d0F',
-                  symbol: 'DAI',
-                  decimals: 18,
-                },
-              ],
-            },
-          },
-          allDetectedTokens: {},
-        },
+      // Mock safelyExecuteWithTimeout to simulate the scenario where the error
+      // bypasses it and reaches the catch block directly (line 289-292)
+      const safelyExecuteSpy = jest
+        .spyOn(controllerUtils, 'safelyExecuteWithTimeout')
+        .mockImplementation(async () => {
+          // Instead of swallowing the error, throw it to reach the catch block
+          throw mockError;
+        });
+
+      // Mock a fetcher that supports the chain
+      const mockFetcher = {
+        supports: jest.fn().mockReturnValue(true),
+        fetch: jest.fn(),
+      };
+
+      Object.defineProperty(controller, '#balanceFetchers', {
+        value: [mockFetcher],
+        writable: true,
       });
 
       await controller.updateBalances({ chainIds: [chainId] });
@@ -1825,7 +1825,7 @@ describe('TokenBalancesController', () => {
       );
 
       // Restore mocks
-      mockMulticall.mockRestore();
+      safelyExecuteSpy.mockRestore();
       consoleWarnSpy.mockRestore();
     });
 

packages/assets-controllers/src/TokenBalancesController.ts

@@ -8,7 +8,11 @@ import type {
   ControllerStateChangeEvent,
   RestrictedMessenger,
 } from '@metamask/base-controller';
-import { isValidHexAddress, toHex } from '@metamask/controller-utils';
+import {
+  isValidHexAddress,
+  safelyExecuteWithTimeout,
+  toHex,
+} from '@metamask/controller-utils';
 import type { KeyringControllerAccountRemovedEvent } from '@metamask/keyring-controller';
 import type {
   NetworkControllerGetNetworkClientByIdAction,
@@ -260,20 +264,20 @@ export class TokenBalancesController extends StaticIntervalPollingController<{
       if (!supportedChains.length) {
         continue;
       }
+
       try {
-        const balances = await Promise.race([
-          fetcher.fetch({
-            chainIds: supportedChains,
-            queryAllAccounts: this.#queryAllAccounts,
-            selectedAccount: selected as ChecksumAddress,
-            allAccounts,
-          }),
-          new Promise<never>((_resolve, reject) =>
-            setTimeout(() => {
-              reject(new Error('timeout'));
-            }, this.getIntervalLength()),
-          ),
-        ]);
+        const balances = await safelyExecuteWithTimeout(
+          async () => {
+            return await fetcher.fetch({
+              chainIds: supportedChains,
+              queryAllAccounts: this.#queryAllAccounts,
+              selectedAccount: selected as ChecksumAddress,
+              allAccounts,
+            });
+          },
+          false,
+          this.getIntervalLength(),
+        );
 
         if (balances && balances.length > 0) {
           aggregated.push(...balances);

packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.test.ts

@@ -456,23 +456,6 @@ describe('AccountsApiBalanceFetcher', () => {
       expect(result.length).toBeGreaterThan(3);
     });
 
-    it('should handle API errors gracefully', async () => {
-      mockSafelyExecute.mockResolvedValue(undefined);
-
-      const result = await balanceFetcher.fetch({
-        chainIds: [MOCK_CHAIN_ID],
-        queryAllAccounts: false,
-        selectedAccount: MOCK_ADDRESS_1 as ChecksumAddress,
-        allAccounts: MOCK_INTERNAL_ACCOUNTS,
-      });
-
-      // Should still have native token guarantee even with API error
-      expect(result).toHaveLength(1);
-      expect(result[0].token).toBe(ZERO_ADDRESS);
-      expect(result[0].success).toBe(true);
-      expect(result[0].value).toStrictEqual(new BN('0'));
-    });
-
     it('should handle missing account address in response', async () => {
       const responseWithMissingAccount: GetBalancesResponse = {
         count: 1,
@@ -1100,7 +1083,6 @@ describe('AccountsApiBalanceFetcher', () => {
     });
 
     it('should test checksum and toCaipAccount helper functions indirectly', async () => {
-      // This test covers lines 47 and 52 by calling methods that use these helpers
       mockToChecksumHexAddress.mockReturnValue('0xCHECKSUMMED');
       mockAccountAddressToCaipReference.mockReturnValue(
         'eip155:1:0xCHECKSUMMED',
@@ -1492,6 +1474,68 @@ describe('AccountsApiBalanceFetcher', () => {
     });
   });
 
+  describe('API error handling and recovery', () => {
+    beforeEach(() => {
+      balanceFetcher = new AccountsApiBalanceFetcher('extension');
+    });
+
+    it('should not throw error when API fails but staked balances succeed', async () => {
+      // Mock console.error to suppress error logging
+      const consoleSpy = jest.spyOn(console, 'error').mockImplementation();
+
+      // Setup successful staking contract
+      const mockShares = {
+        toString: () => '1000000000000000000',
+        gt: jest.fn().mockReturnValue(true),
+      };
+      const mockAssets = {
+        toString: () => '2000000000000000000',
+      };
+
+      const localMockContract = {
+        getShares: jest.fn().mockResolvedValue(mockShares),
+        convertToAssets: jest.fn().mockResolvedValue(mockAssets),
+      };
+
+      const mockContractConstructor = jest.requireMock(
+        '@ethersproject/contracts',
+      ).Contract;
+      mockContractConstructor.mockImplementation(() => localMockContract);
+
+      const mockProvider = { call: jest.fn() };
+      const mockGetProvider = jest.fn().mockReturnValue(mockProvider);
+
+      const fetcherWithProvider = new AccountsApiBalanceFetcher(
+        'extension',
+        mockGetProvider,
+      );
+
+      // Make API fail but staking succeed
+      mockFetchMultiChainBalancesV4.mockRejectedValue(new Error('API failure'));
+
+      try {
+        const result = await fetcherWithProvider.fetch({
+          chainIds: [MOCK_CHAIN_ID],
+          queryAllAccounts: false,
+          selectedAccount: MOCK_ADDRESS_1 as ChecksumAddress,
+          allAccounts: MOCK_INTERNAL_ACCOUNTS,
+        });
+
+        // Should have mixed results: failed API entries + successful staked balance
+        const successfulEntries = result.filter((r) => r.success);
+        const errorEntries = result.filter((r) => !r.success);
+
+        expect(successfulEntries.length).toBeGreaterThan(0); // Staked balance succeeded
+        expect(errorEntries.length).toBeGreaterThan(0); // API entries failed
+
+        // Should not throw since we have some successful results
+        expect(result.length).toBeGreaterThan(0);
+      } finally {
+        consoleSpy.mockRestore();
+      }
+    });
+  });
+
   describe('precision handling in balance conversion', () => {
     beforeEach(() => {
       balanceFetcher = new AccountsApiBalanceFetcher('extension');

packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.ts

@@ -258,10 +258,20 @@ export class AccountsApiBalanceFetcher implements BalanceFetcher {
       return [];
     }
 
-    const [balances, stakedBalances] = await Promise.all([
-      safelyExecute(() => this.#fetchBalances(caipAddrs)),
-      this.#fetchStakedBalances(caipAddrs),
-    ]);
+    // Don't use safelyExecute here - let real errors propagate
+    let balances;
+    let apiError = false;
+
+    try {
+      balances = await this.#fetchBalances(caipAddrs);
+    } catch (error) {
+      // Mark that we had an API error so we don't add fake zero balances
+      apiError = true;
+      console.error('Failed to fetch balances from API:', error);
+      balances = undefined;
+    }
+
+    const stakedBalances = await this.#fetchStakedBalances(caipAddrs);
 
     const results: ProcessedBalance[] = [];
 
@@ -331,28 +341,48 @@ export class AccountsApiBalanceFetcher implements BalanceFetcher {
       results.push(...apiBalances);
     }
 
-    // Ensure native token entries exist for all addresses/chains, even if not returned by API
-    addressChainMap.forEach((chains, address) => {
-      chains.forEach((chainId) => {
-        const key = `${address}-${chainId}`;
-        const existingBalance = nativeBalancesFromAPI.get(key);
+    // Only add zero native balance entries if API succeeded but didn't return balances
+    // Don't add fake zero balances if the API failed entirely
+    if (!apiError) {
+      addressChainMap.forEach((chains, address) => {
+        chains.forEach((chainId) => {
+          const key = `${address}-${chainId}`;
+          const existingBalance = nativeBalancesFromAPI.get(key);
 
-        if (!existingBalance) {
-          // Add zero native balance entry if API didn't return one
+          if (!existingBalance) {
+            // Add zero native balance entry if API succeeded but didn't return one
+            results.push({
+              success: true,
+              value: new BN('0'),
+              account: address as ChecksumAddress,
+              token: ZERO_ADDRESS,
+              chainId,
+            });
+          }
+        });
+      });
+    } else {
+      // If API failed, add error entries for all requested addresses/chains
+      addressChainMap.forEach((chains, address) => {
+        chains.forEach((chainId) => {
           results.push({
-            success: true,
-            value: new BN('0'),
+            success: false,
             account: address as ChecksumAddress,
             token: ZERO_ADDRESS,
             chainId,
           });
-        }
+        });
       });
-    });
+    }
 
     // Add staked balances
     results.push(...stakedBalances);
 
+    // If we had an API error and no successful results, throw the error
+    if (apiError && results.every((r) => !r.success)) {
+      throw new Error('Failed to fetch any balance data due to API error');
+    }
+
     return results;
   }
 }