fix: ensure StakingBalance is shown or hidden appropriately per asset…

app/components/UI/AssetOverview/TokenDetails/TokenDetails.tsx

@@ -146,7 +146,7 @@ const TokenDetails: React.FC<TokenDetailsProps> = ({ asset }) => {
 
   return (
     <View style={styles.tokenDetailsContainer}>
-      {asset.isETH && <StakingEarnings />}
+      {asset.isETH && <StakingEarnings asset={asset} />}
       {(asset.isETH || tokenMetadata) && (
         <TokenDetailsList tokenDetails={tokenDetails} />
       )}

app/components/UI/Stake/__mocks__/mockData.ts

@@ -10,10 +10,12 @@ import { Contract } from 'ethers';
 import { Stake } from '../sdk/stakeSdkProvider';
 
 export const MOCK_STAKED_ETH_ASSET = {
+  chainId: '0x1',
   balance: '4.9999 ETH',
   balanceFiat: '$13,292.20',
   name: 'Staked Ethereum',
   symbol: 'ETH',
+  isETH: true,
 } as TokenI;
 
 export const MOCK_GET_POOLED_STAKES_API_RESPONSE: PooledStakes = {
@@ -67,17 +69,18 @@ export const MOCK_GET_POOLED_STAKES_API_RESPONSE: PooledStakes = {
   exchangeRate: '1.010906701603882254',
 };
 
-export const MOCK_GET_POOLED_STAKES_API_RESPONSE_HIGH_ASSETS_AMOUNT: PooledStakes = {
-  accounts: [
-    {
-      account: '0x0111111111abcdef2222222222abcdef33333333',
-      lifetimeRewards: '0',
-      assets: '99999999990000000000000',
-      exitRequests: [],
-    },
-  ],
-  exchangeRate: '1.010906701603882254',
-};
+export const MOCK_GET_POOLED_STAKES_API_RESPONSE_HIGH_ASSETS_AMOUNT: PooledStakes =
+  {
+    accounts: [
+      {
+        account: '0x0111111111abcdef2222222222abcdef33333333',
+        lifetimeRewards: '0',
+        assets: '99999999990000000000000',
+        exitRequests: [],
+      },
+    ],
+    exchangeRate: '1.010906701603882254',
+  };
 
 export const MOCK_GET_VAULT_RESPONSE: VaultData = {
   apy: '2.853065141088762750393474836309926',

app/components/UI/Stake/components/StakingBalance/StakingBalance.test.tsx

@@ -14,6 +14,7 @@ import { createMockAccountsControllerState } from '../../../../../util/test/acco
 import { backgroundState } from '../../../../../util/test/initial-root-state';
 // eslint-disable-next-line import/no-namespace
 import * as networks from '../../../../../util/networks';
+import { mockNetworkState } from '../../../../../util/test/network';
 
 const MOCK_ADDRESS_1 = '0x0';
 
@@ -168,4 +169,39 @@ describe('StakingBalance', () => {
       screen: Routes.STAKING.UNSTAKE,
     });
   });
+
+  it('should not render if asset chainId is not a staking supporting chain', () => {
+    const { queryByText, queryByTestId } = renderWithProvider(
+      <StakingBalance asset={{ ...MOCK_STAKED_ETH_ASSET, chainId: '0x4' }} />,
+      { state: mockInitialState },
+    );
+    expect(queryByTestId('staking-balance-container')).toBeNull();
+    expect(queryByText(strings('stake.stake_more'))).toBeNull();
+    expect(queryByText(strings('stake.unstake'))).toBeNull();
+    expect(queryByText(strings('stake.claim'))).toBeNull();
+  });
+
+  it('should not render claim link or action buttons if asset.chainId is not selected chainId', () => {
+    const { queryByText, queryByTestId } = renderWithProvider(
+      <StakingBalance asset={MOCK_STAKED_ETH_ASSET} />,
+      {
+        state: {
+          ...mockInitialState,
+          engine: {
+            ...mockInitialState.engine,
+            backgroundState: {
+              ...mockInitialState.engine.backgroundState,
+              NetworkController: {
+                ...mockNetworkState({ chainId: '0x4268' }),
+              },
+            },
+          },
+        },
+      },
+    );
+    expect(queryByTestId('staking-balance-container')).toBeTruthy();
+    expect(queryByText(strings('stake.stake_more'))).toBeNull();
+    expect(queryByText(strings('stake.unstake'))).toBeNull();
+    expect(queryByText(strings('stake.claim'))).toBeNull();
+  });
 });

app/components/UI/Stake/components/StakingBalance/StakingBalance.tsx

@@ -1,4 +1,5 @@
 import React, { useEffect, useMemo, useState } from 'react';
+import { Hex } from '@metamask/utils';
 import Badge, {
   BadgeVariant,
 } from '../../../../../component-library/components/Badges/Badge';
@@ -35,7 +36,7 @@ import {
 import { multiplyValueByPowerOfTen } from '../../utils/bignumber';
 import StakingCta from './StakingCta/StakingCta';
 import useStakingEligibility from '../../hooks/useStakingEligibility';
-import useStakingChain from '../../hooks/useStakingChain';
+import { useStakingChainByChainId } from '../../hooks/useStakingChain';
 import usePooledStakes from '../../hooks/usePooledStakes';
 import useVaultData from '../../hooks/useVaultData';
 import { StakeSDKProvider } from '../../sdk/stakeSdkProvider';
@@ -64,7 +65,9 @@ const StakingBalanceContent = ({ asset }: StakingBalanceProps) => {
 
   const { isEligible: isEligibleForPooledStaking } = useStakingEligibility();
 
-  const { isStakingSupportedChain } = useStakingChain();
+  const { isStakingSupportedChain } = useStakingChainByChainId(
+    asset.chainId as Hex,
+  );
 
   const { trackEvent, createEventBuilder } = useMetrics();
 
@@ -81,7 +84,7 @@ const StakingBalanceContent = ({ asset }: StakingBalanceProps) => {
   const {
     formattedStakedBalanceETH: stakedBalanceETH,
     formattedStakedBalanceFiat: stakedBalanceFiat,
-  } = useBalance();
+  } = useBalance(asset.chainId as Hex);
 
   const { unstakingRequests, claimableRequests } = useMemo(() => {
     const exitRequests = pooledStakesData?.exitRequests ?? [];
@@ -129,6 +132,9 @@ const StakingBalanceContent = ({ asset }: StakingBalanceProps) => {
   }
 
   const renderStakingContent = () => {
+    if (chainId !== asset.chainId) {
+      return <></>;
+    }
     if (isLoadingPooledStakesData) {
       return (
         <SkeletonPlaceholder>
@@ -197,7 +203,7 @@ const StakingBalanceContent = ({ asset }: StakingBalanceProps) => {
   };
 
   return (
-    <View>
+    <View testID="staking-balance-container">
       {hasEthToUnstake && !isLoadingPooledStakesData && (
         <AssetElement
           asset={asset}
@@ -209,7 +215,10 @@ const StakingBalanceContent = ({ asset }: StakingBalanceProps) => {
             badgeElement={
               <Badge
                 variant={BadgeVariant.Network}
-                imageSource={NetworkBadgeSource(chainId, asset.symbol)}
+                imageSource={NetworkBadgeSource(
+                  asset.chainId as Hex,
+                  asset.ticker || asset.symbol,
+                )}
                 name={networkName}
               />
             }

app/components/UI/Stake/components/StakingBalance/__snapshots__/StakingBalance.test.tsx.snap

@@ -1,7 +1,9 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
 exports[`StakingBalance render matches snapshot 1`] = `
-<View>
+<View
+  testID="staking-balance-container"
+>
   <TouchableOpacity
     onLongPress={[Function]}
     onPress={[Function]}
@@ -469,7 +471,9 @@ exports[`StakingBalance render matches snapshot 1`] = `
 `;
 
 exports[`StakingBalance should match the snapshot when portfolio view is enabled   1`] = `
-<View>
+<View
+  testID="staking-balance-container"
+>
   <TouchableOpacity
     onLongPress={[Function]}
     onPress={[Function]}

app/components/UI/Stake/components/StakingEarnings/StakingEarnings.test.tsx

@@ -3,6 +3,7 @@ import StakingEarnings from './';
 import renderWithProvider from '../../../../../util/test/renderWithProvider';
 import { strings } from '../../../../../../locales/i18n';
 import { mockNetworkState } from '../../../../../util/test/network';
+import { TokenI } from '../../../Tokens/types';
 
 const mockNavigate = jest.fn();
 
@@ -66,9 +67,12 @@ jest.mock('../../../../../core/Engine', () => ({
 
 describe('Staking Earnings', () => {
   it('should render correctly', () => {
-    const { toJSON, getByText } = renderWithProvider(<StakingEarnings />, {
-      state: STATE_MOCK,
-    });
+    const { toJSON, getByText } = renderWithProvider(
+      <StakingEarnings asset={{ chainId: '0x1' } as TokenI} />,
+      {
+        state: STATE_MOCK,
+      },
+    );
 
     expect(getByText(strings('stake.your_earnings'))).toBeDefined();
     expect(getByText(strings('stake.annual_rate'))).toBeDefined();

app/components/UI/Stake/components/StakingEarnings/index.tsx

@@ -1,5 +1,6 @@
 import React from 'react';
 import { View } from 'react-native';
+import { Hex } from '@metamask/utils';
 import Text, {
   TextColor,
   TextVariant,
@@ -15,15 +16,20 @@ import ButtonIcon, {
 } from '../../../../../component-library/components/Buttons/ButtonIcon';
 import useTooltipModal from '../../../../../components/hooks/useTooltipModal';
 import { strings } from '../../../../../../locales/i18n';
-import useStakingChain from '../../hooks/useStakingChain';
+import { useStakingChainByChainId } from '../../hooks/useStakingChain';
 import { StakeSDKProvider } from '../../sdk/stakeSdkProvider';
 import useStakingEarnings from '../../hooks/useStakingEarnings';
 import SkeletonPlaceholder from 'react-native-skeleton-placeholder';
 import { withMetaMetrics } from '../../utils/metaMetrics/withMetaMetrics';
 import { MetaMetricsEvents } from '../../../../hooks/useMetrics';
 import { getTooltipMetricProperties } from '../../utils/metaMetrics/tooltipMetaMetricsUtils';
+import { TokenI } from '../../../Tokens/types';
 
-const StakingEarningsContent = () => {
+export interface StakingEarningsProps {
+  asset: TokenI;
+}
+
+const StakingEarningsContent = ({ asset }: StakingEarningsProps) => {
   const { styles } = useStyles(styleSheet, {});
 
   const { openTooltipModal } = useTooltipModal();
@@ -38,14 +44,16 @@ const StakingEarningsContent = () => {
     hasStakedPositions,
   } = useStakingEarnings();
 
+  const { isStakingSupportedChain } = useStakingChainByChainId(
+    asset.chainId as Hex,
+  );
+
   const onDisplayAnnualRateTooltip = () =>
     openTooltipModal(
       strings('stake.annual_rate'),
       strings('tooltip_modal.reward_rate.tooltip'),
     );
 
-  const { isStakingSupportedChain } = useStakingChain();
-
   if (!isStakingSupportedChain || !hasStakedPositions) return <></>;
 
   return (
@@ -175,9 +183,9 @@ const StakingEarningsContent = () => {
   );
 };
 
-export const StakingEarnings = () => (
+export const StakingEarnings = ({ asset }: StakingEarningsProps) => (
   <StakeSDKProvider>
-    <StakingEarningsContent />
+    <StakingEarningsContent asset={asset} />
   </StakeSDKProvider>
 );
 

app/components/UI/Stake/hooks/useBalance.test.tsx

@@ -22,7 +22,20 @@ const initialState = {
       AccountTrackerController: {
         accountsByChainId: {
           '0x1': {
-            [MOCK_ADDRESS_1]: { balance: toHex('12345678909876543210000000'), stakedBalance: toHex(MOCK_GET_POOLED_STAKES_API_RESPONSE.accounts[0].assets) },
+            [MOCK_ADDRESS_1]: {
+              balance: toHex('12345678909876543210000000'),
+              stakedBalance: toHex(
+                MOCK_GET_POOLED_STAKES_API_RESPONSE.accounts[0].assets,
+              ),
+            },
+          },
+          '0x4268': {
+            [MOCK_ADDRESS_1]: {
+              balance: toHex('22345678909876543210000000'),
+              stakedBalance: toHex(
+                MOCK_GET_POOLED_STAKES_API_RESPONSE.accounts[0].assets,
+              ),
+            },
           },
         },
       },
@@ -139,4 +152,21 @@ describe('useBalance', () => {
     expect(result.current.stakedBalanceFiatNumber).toBe(319999999.968); // Staked balance in fiat number
     expect(result.current.formattedStakedBalanceFiat).toBe('$319999999.96'); // should round to floor
   });
+
+  it('returns correct stake amounts and fiat values when chainId is overriden', async () => {
+    const { result } = renderHookWithProvider(() => useBalance('0x4268'), {
+      state: initialState,
+    });
+
+    expect(result.current.balanceETH).toBe('22345678.90988');
+    expect(result.current.balanceWei.toString()).toBe(
+      '22345678909876543210000000',
+    );
+    expect(result.current.balanceFiat).toBe('$71506172511.60'); // Fiat balance
+    expect(result.current.balanceFiatNumber).toBe(71506172511.6); // Fiat number balance
+    expect(result.current.stakedBalanceWei).toBe('5791332670714232000');
+    expect(result.current.formattedStakedBalanceETH).toBe('5.79133 ETH'); // Formatted ETH balance
+    expect(result.current.stakedBalanceFiatNumber).toBe(18532.26454); // Staked balance in fiat number
+    expect(result.current.formattedStakedBalanceFiat).toBe('$18532.26'); //
+  });
 });

app/components/UI/Stake/hooks/useBalance.ts

@@ -1,5 +1,6 @@
 import { useMemo } from 'react';
 import { useSelector } from 'react-redux';
+import { Hex } from '@metamask/utils';
 import { selectSelectedInternalAccountFormattedAddress } from '../../../../selectors/accountsController';
 import { selectAccountsByChainId } from '../../../../selectors/accountTrackerController';
 import {
@@ -14,21 +15,22 @@ import {
   weiToFiatNumber,
 } from '../../../../util/number';
 
-const useBalance = () => {
+const useBalance = (chainId?: Hex) => {
   const accountsByChainId = useSelector(selectAccountsByChainId);
-  const chainId = useSelector(selectChainId);
+  const selectedChainId = useSelector(selectChainId);
   const selectedAddress = useSelector(
     selectSelectedInternalAccountFormattedAddress,
   );
   const currentCurrency = useSelector(selectCurrentCurrency);
   const currencyRates = useSelector(selectCurrencyRates);
+  const balanceChainId = chainId || selectedChainId;
   const conversionRate = currencyRates?.ETH?.conversionRate ?? 1;
   const rawAccountBalance = selectedAddress
-    ? accountsByChainId[chainId]?.[selectedAddress]?.balance
+    ? accountsByChainId[balanceChainId]?.[selectedAddress]?.balance
     : '0';
 
   const stakedBalance = selectedAddress
-    ? accountsByChainId[chainId]?.[selectedAddress]?.stakedBalance || '0'
+    ? accountsByChainId[balanceChainId]?.[selectedAddress]?.stakedBalance || '0'
     : '0';
 
   const balanceETH = useMemo(