fix: update non-evm token balances on flip (#24329)

<!--
Please submit this PR as a draft initially.
Do not mark it as "Ready for review" until the template has been
completely filled out, and PR status checks have passed at least once.
-->

## **Description**

fixes an issue where flipping source/dest assets do not update balances
when user has zero balance on non-evm chains

<!--
Write a short description of the changes included in this pull request,
also include relevant motivation and context. Have in mind the following
questions:
1. What is the reason for the change?
2. What is the improvement/solution?
-->

## **Changelog**

<!--
If this PR is not End-User-Facing and should not show up in the
CHANGELOG, you can choose to either:
1. Write `CHANGELOG entry: null`
2. Label with `no-changelog`

If this PR is End-User-Facing, please write a short User-Facing
description in the past tense like:
`CHANGELOG entry: Added a new tab for users to see their NFTs`
`CHANGELOG entry: Fixed a bug that was causing some NFTs to flicker`

(This helps the Release Engineer do their job more quickly and
accurately)
-->

CHANGELOG entry: fixes an issue where flipping source/dest assets do not
update balances when user has zero balance on non-evm chains

## **Related issues**

Fixes: https://consensyssoftware.atlassian.net/browse/SWAPS-3696

## **Manual testing steps**

```gherkin
Feature: my feature name

  Scenario: user [verb for user action]
    Given [describe expected initial app state]

    When user [verb for user action]
    Then [describe expected outcome]
```

## **Screenshots/Recordings**

<!-- If applicable, add screenshots and/or recordings to visualize the
before and after of your change. -->

### **Before**

<!-- [screenshots/recordings] -->

### **After**

<!-- [screenshots/recordings] -->

## **Pre-merge author checklist**

- [ ] I’ve followed [MetaMask Contributor
Docs](https://github.com/MetaMask/contributor-docs) and [MetaMask Mobile
Coding
Standards](https://github.com/MetaMask/metamask-mobile/blob/main/.github/guidelines/CODING_GUIDELINES.md).
- [ ] I've completed the PR template to the best of my ability
- [ ] I’ve included tests if applicable
- [ ] I’ve documented my code using [JSDoc](https://jsdoc.app/) format
if applicable
- [ ] I’ve applied the right labels on the PR (see [labeling
guidelines](https://github.com/MetaMask/metamask-mobile/blob/main/.github/guidelines/LABELING_GUIDELINES.md)).
Not required for external contributors.

## **Pre-merge reviewer checklist**

- [ ] I've manually tested the PR (e.g. pull and build branch, run the
app, test code being changed).
- [ ] I confirm that this PR addresses all acceptance criteria described
in the ticket it closes and includes the necessary testing evidence such
as recordings and or screenshots.

<!-- CURSOR_SUMMARY -->
---

> [!NOTE]
> Improves `useLatestBalance` behavior for non‑EVM and token switching
scenarios.
> 
> - Resets balance state when `token.address` changes to avoid stale UI;
returns cached balance immediately while fetching new data
> - For non‑EVM chains, uses controller data and sets balance to `0`
when token not found or balance is empty/undefined
> - Expands effect deps to include previous/current token address
> - Adds comprehensive tests covering address flips, non‑EVM fallbacks,
zero/invalid balances, and race conditions
> 
> <sup>Written by [Cursor
Bugbot](https://cursor.com/dashboard?tab=bugbot) for commit
330d2db. This will update automatically
on new commits. Configure
[here](https://cursor.com/dashboard?tab=bugbot).</sup>
<!-- /CURSOR_SUMMARY -->

Author: GeorgeGkas

app/components/UI/Bridge/hooks/useLatestBalance/index.ts

@@ -142,11 +142,21 @@ export const useLatestBalance = (token: {
           displayBalance,
           atomicBalance: parseUnits(displayBalance, token.decimals),
         });
+      } else {
+        setBalance({
+          displayBalance: '0',
+          atomicBalance: parseUnits('0'),
+        });
       }
     }
   }, [token.address, token.decimals, chainId, selectedAddress, nonEvmTokens]);
 
   useEffect(() => {
+    // Reset balance state when token address changes to prevent stale balance display
+    if (previousToken?.address !== token.address) {
+      setBalance(undefined);
+    }
+
     // In case chainId is undefined, exit early to avoid
     // calling handleFetchEvmAtomicBalance which will trigger an invalid address error
     // when selectedAddress is a non-EVM chain.
@@ -161,7 +171,13 @@ export const useLatestBalance = (token: {
     if (isCaipChainId(chainId) && isNonEvmChainId(chainId)) {
       handleNonEvmAtomicBalance();
     }
-  }, [handleFetchEvmAtomicBalance, handleNonEvmAtomicBalance, chainId]);
+  }, [
+    handleFetchEvmAtomicBalance,
+    handleNonEvmAtomicBalance,
+    chainId,
+    previousToken?.address,
+    token.address,
+  ]);
 
   const cachedBalance = useMemo(() => {
     const displayBalance = token.balance;

app/components/UI/Bridge/hooks/useLatestBalance/useLatestBalance.test.tsx

@@ -9,7 +9,7 @@ import { useLatestBalance } from '.';
 import { getProviderByChainId } from '../../../../../util/notifications/methods/common';
 import { BigNumber, constants } from 'ethers';
 import { waitFor } from '@testing-library/react-native';
-import { Hex } from '@metamask/utils';
+import { Hex, type CaipAssetId } from '@metamask/utils';
 import { SolScope } from '@metamask/keyring-api';
 import { cloneDeep } from 'lodash';
 
@@ -366,12 +366,13 @@ describe('useLatestBalance', () => {
       { state: initialState },
     );
 
-    // Should not call EVM provider methods when chainId is CAIP (Solana)
-    expect(getProviderByChainId).not.toHaveBeenCalled();
-    expect(mockProvider.getBalance).not.toHaveBeenCalled();
-    expect(result.current).toEqual({
-      displayBalance: '50.0',
-      atomicBalance: BigNumber.from('50000000000000000000'),
+    // When EVM address is used with non-EVM chainId, it tries to find it in nonEvmTokens
+    // Since it's not found, it sets balance to 0
+    await waitFor(() => {
+      expect(result.current).toEqual({
+        displayBalance: '0',
+        atomicBalance: BigNumber.from('0'),
+      });
     });
   });
 
@@ -661,4 +662,233 @@ describe('useLatestBalance', () => {
       });
     });
   });
+
+  describe('balance reset when token address changes', () => {
+    it('resets balance to undefined when token address changes', async () => {
+      let tokenAddress = '0x1234567890123456789012345678901234567890';
+
+      const { result, rerender } = renderHookWithProvider(
+        () =>
+          useLatestBalance({
+            address: tokenAddress,
+            decimals: 6,
+            chainId: '0x1' as Hex,
+            balance: '10.0',
+          }),
+        { state: initialState },
+      );
+
+      await waitFor(() => {
+        expect(result.current?.displayBalance).toBe('1.0');
+      });
+
+      tokenAddress = '0xabcdefabcdefabcdefabcdefabcdefabcdefabcd';
+      rerender({ state: initialState });
+
+      expect(result.current?.displayBalance).toBe('10.0');
+    });
+
+    it('returns cached balance immediately after token address changes', async () => {
+      let tokenAddress = '0x1111111111111111111111111111111111111111';
+      let tokenBalance = '100.5';
+
+      const { result, rerender } = renderHookWithProvider(
+        () =>
+          useLatestBalance({
+            address: tokenAddress,
+            decimals: 6,
+            chainId: '0x1' as Hex,
+            balance: tokenBalance,
+          }),
+        { state: initialState },
+      );
+
+      expect(result.current?.displayBalance).toBe('100.5');
+
+      tokenAddress = '0x2222222222222222222222222222222222222222';
+      tokenBalance = '250.75';
+      rerender({ state: initialState });
+
+      expect(result.current?.displayBalance).toBe('250.75');
+      expect(result.current?.atomicBalance).toEqual(
+        BigNumber.from('250750000'),
+      );
+    });
+
+    it('fetches new balance after token address changes', async () => {
+      let tokenAddress = constants.AddressZero;
+
+      const { result, rerender } = renderHookWithProvider(
+        () =>
+          useLatestBalance({
+            address: tokenAddress,
+            decimals: 18,
+            chainId: '0x1' as Hex,
+          }),
+        { state: initialState },
+      );
+
+      await waitFor(() => {
+        expect(result.current?.displayBalance).toBe('1.0');
+      });
+
+      jest.clearAllMocks();
+
+      tokenAddress = '0x9999999999999999999999999999999999999999';
+      rerender({ state: initialState });
+
+      await waitFor(() => {
+        expect(getProviderByChainId).toHaveBeenCalled();
+      });
+    });
+
+    it('resets balance state before fetching new balance for different token', async () => {
+      let tokenAddress = '0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa';
+      let tokenBalance = '50.0';
+
+      const { result, rerender } = renderHookWithProvider(
+        () =>
+          useLatestBalance({
+            address: tokenAddress,
+            decimals: 6,
+            chainId: '0x1' as Hex,
+            balance: tokenBalance,
+          }),
+        { state: initialState },
+      );
+
+      await waitFor(() => {
+        expect(result.current?.displayBalance).toBe('1.0');
+      });
+
+      const previousBalance = result.current;
+
+      tokenAddress = '0xbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb';
+      tokenBalance = '75.0';
+      rerender({ state: initialState });
+
+      expect(result.current?.displayBalance).toBe('75.0');
+      expect(result.current).not.toBe(previousBalance);
+    });
+  });
+
+  describe('non-EVM token not found in balance controller', () => {
+    it('sets balance to 0 when Solana token is not found in nonEvmTokens', async () => {
+      const state = cloneDeep(initialState);
+      state.engine.backgroundState.AccountsController.internalAccounts.selectedAccount =
+        solanaAccountId;
+
+      const { result } = renderHookWithProvider(
+        () =>
+          useLatestBalance({
+            chainId: SolScope.Mainnet,
+            address: 'TokenNotInBalanceController123',
+            decimals: 9,
+          }),
+        { state },
+      );
+
+      await waitFor(() => {
+        expect(result.current).toEqual({
+          displayBalance: '0',
+          atomicBalance: BigNumber.from('0'),
+        });
+      });
+    });
+
+    it('sets balance to 0 when Solana token is not found in nonEvmTokens list', async () => {
+      const state = cloneDeep(initialState);
+      state.engine.backgroundState.AccountsController.internalAccounts.selectedAccount =
+        solanaAccountId;
+
+      const nonExistentTokenAddress =
+        'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/token:NonExistentToken123' as CaipAssetId;
+
+      const { result } = renderHookWithProvider(
+        () =>
+          useLatestBalance({
+            chainId: SolScope.Mainnet,
+            address: nonExistentTokenAddress,
+            decimals: 9,
+          }),
+        { state },
+      );
+
+      await waitFor(() => {
+        expect(result.current).toEqual({
+          displayBalance: '0',
+          atomicBalance: BigNumber.from('0'),
+        });
+      });
+    });
+
+    it('sets balance to 0 when non-EVM token has undefined balance in controller', async () => {
+      const state = cloneDeep(initialState);
+      state.engine.backgroundState.AccountsController.internalAccounts.selectedAccount =
+        solanaAccountId;
+
+      const { result } = renderHookWithProvider(
+        () =>
+          useLatestBalance({
+            chainId: SolScope.Mainnet,
+            address: 'UnknownSolanaToken456',
+            decimals: 6,
+          }),
+        { state },
+      );
+
+      await waitFor(() => {
+        expect(result.current).toEqual({
+          displayBalance: '0',
+          atomicBalance: BigNumber.from('0'),
+        });
+      });
+    });
+
+    it('returns valid balance when Solana token is found in nonEvmTokens', async () => {
+      const state = cloneDeep(initialState);
+      state.engine.backgroundState.AccountsController.internalAccounts.selectedAccount =
+        solanaAccountId;
+
+      const { result } = renderHookWithProvider(
+        () =>
+          useLatestBalance({
+            chainId: SolScope.Mainnet,
+            address: solanaToken2Address,
+            decimals: 6,
+          }),
+        { state },
+      );
+
+      await waitFor(() => {
+        expect(result.current).toEqual({
+          displayBalance: '20000.456',
+          atomicBalance: BigNumber.from('20000456000'),
+        });
+      });
+    });
+
+    it('handles empty balance string from controller by setting to 0', async () => {
+      const state = cloneDeep(initialState);
+      state.engine.backgroundState.AccountsController.internalAccounts.selectedAccount =
+        solanaAccountId;
+
+      const { result } = renderHookWithProvider(
+        () =>
+          useLatestBalance({
+            chainId: SolScope.Mainnet,
+            address: 'TokenWithEmptyBalance789',
+            decimals: 9,
+          }),
+        { state },
+      );
+
+      await waitFor(() => {
+        expect(result.current).toEqual({
+          displayBalance: '0',
+          atomicBalance: BigNumber.from('0'),
+        });
+      });
+    });
+  });
 });