Add token

Author: matthewwalsh0

app/components/Views/confirmations/components/info/predict-withdraw-info/predict-withdraw-info.tsx

@@ -3,11 +3,21 @@ import { strings } from '../../../../../../../locales/i18n';
 import useNavbar from '../../../hooks/ui/useNavbar';
 import { CustomAmountInfo } from '../custom-amount-info';
 import { PredictWithdrawBalance } from '../../predict-confirmations/predict-withdraw-balance/predict-withdraw-balance';
-import { PREDICT_CURRENCY } from '../../../constants/predict';
+import { POLYGON_USDCE, PREDICT_CURRENCY } from '../../../constants/predict';
+import { useAddToken } from '../../../hooks/tokens/useAddToken';
+import { CHAIN_IDS } from '@metamask/transaction-controller';
 
 export function PredictWithdrawInfo() {
   useNavbar(strings('confirm.title.predict_withdraw'));
 
+  useAddToken({
+    chainId: CHAIN_IDS.POLYGON,
+    decimals: POLYGON_USDCE.decimals,
+    name: POLYGON_USDCE.name,
+    symbol: POLYGON_USDCE.symbol,
+    tokenAddress: POLYGON_USDCE.address,
+  });
+
   return (
     <CustomAmountInfo disablePay currency={PREDICT_CURRENCY}>
       <PredictWithdrawBalance />

app/components/Views/confirmations/hooks/tokens/useTokenWithBalance.ts

@@ -3,19 +3,29 @@ import { useSelector } from 'react-redux';
 import { selectSingleTokenByAddressAndChainId } from '../../../../../selectors/tokensController';
 import { RootState } from '../../../../../reducers';
 import { selectSingleTokenBalance } from '../../../../../selectors/tokenBalancesController';
-import { selectSelectedInternalAccountAddress } from '../../../../../selectors/accountsController';
 import { toChecksumAddress } from '../../../../../util/address';
 import { useTokenFiatRate } from './useTokenFiatRates';
 import { BigNumber } from 'bignumber.js';
 import { useMemo } from 'react';
 import useFiatFormatter from '../../../../UI/SimulationDetails/FiatDisplay/useFiatFormatter';
-import { selectAccountBalanceByChainId } from '../../../../../selectors/accountTrackerController';
+import { selectAccountsByChainId } from '../../../../../selectors/accountTrackerController';
 import { selectConversionRateByChainId } from '../../../../../selectors/currencyRateController';
 import { selectTickerByChainId } from '../../../../../selectors/networkController';
 import { getNativeTokenAddress } from '../../utils/asset';
+import { selectSelectedAccountGroupInternalAccounts } from '../../../../../selectors/multichainAccounts/accountTreeController';
+import { selectSelectedInternalAccountAddress } from '../../../../../selectors/accountsController';
 
 export function useTokenWithBalance(tokenAddress: Hex, chainId: Hex) {
   const selectedAddress = useSelector(selectSelectedInternalAccountAddress);
+
+  const internalAccounts = useSelector(
+    selectSelectedAccountGroupInternalAccounts,
+  );
+
+  const selectedEvmAddress =
+    (internalAccounts.find((a) => a.type === 'eip155:eoa')?.address as Hex) ??
+    selectedAddress;
+
   const fiatFormatter = useFiatFormatter();
   const nativeTokenAddress = getNativeTokenAddress(chainId);
 
@@ -26,7 +36,7 @@ export function useTokenWithBalance(tokenAddress: Hex, chainId: Hex) {
   const tokenBalanceResult = useSelector((state: RootState) =>
     selectSingleTokenBalance(
       state,
-      selectedAddress as Hex,
+      selectedEvmAddress,
       chainId,
       toChecksumAddress(tokenAddress),
     ),
@@ -38,9 +48,13 @@ export function useTokenWithBalance(tokenAddress: Hex, chainId: Hex) {
 
   const tokenFiatRate = useTokenFiatRate(tokenAddress, chainId);
 
-  const nativeBalanceResult = useSelector((state: RootState) =>
-    selectAccountBalanceByChainId(state, chainId),
-  );
+  const accountsByChainId = useSelector(selectAccountsByChainId);
+
+  const nativeBalanceResult = Object.entries(
+    accountsByChainId?.[chainId] ?? {},
+  ).find(
+    ([address]) => address.toLowerCase() === selectedEvmAddress.toLowerCase(),
+  )?.[1];
 
   const conversionRate = useSelector((state: RootState) =>
     selectConversionRateByChainId(state, chainId),

app/selectors/tokensController.test.ts

@@ -23,6 +23,7 @@ describe('TokensController Selectors', () => {
     allTokens: {
       '0x1': {
         '0xAddress1': [mockToken],
+        '0xAddress2': [mockToken2],
       },
     },
     allDetectedTokens: {
@@ -241,7 +242,10 @@ describe('TokensController Selectors', () => {
 
   describe('selectAllTokensFlat', () => {
     it('returns all tokens as a flat array', () => {
-      expect(selectAllTokensFlat(mockRootState)).toStrictEqual([mockToken]);
+      expect(selectAllTokensFlat(mockRootState)).toStrictEqual([
+        mockToken,
+        mockToken2,
+      ]);
     });
 
     it('returns an empty array if no tokens are present', () => {
@@ -361,10 +365,19 @@ describe('TokensController Selectors', () => {
     it('returns undefined if no token exists for the given address and chain ID', () => {
       const token = selectSingleTokenByAddressAndChainId(
         mockRootState,
-        '0xAddress1',
+        '0xToken3',
         '0x2',
       );
       expect(token).toBeUndefined();
     });
+
+    it('returns token not from selected address', () => {
+      const token = selectSingleTokenByAddressAndChainId(
+        mockRootState,
+        '0xToken2',
+        '0x1',
+      );
+      expect(token).toStrictEqual(mockToken2);
+    });
   });
 });

app/selectors/tokensController.ts

@@ -216,16 +216,12 @@ export const selectTransformedTokens = createSelector(
 
 export const selectSingleTokenByAddressAndChainId = createSelector(
   selectAllTokens,
-  selectSelectedInternalAccountAddress,
   (_state: RootState, tokenAddress: Hex) => tokenAddress,
   (_state: RootState, _tokenAddress: Hex, chainId: Hex) => chainId,
-  (allTokens, selectedAddress, tokenAddress, chainId) => {
-    if (!selectedAddress) return undefined;
-
-    const tokensForAddressAndChain =
-      allTokens[chainId]?.[selectedAddress] ?? [];
+  (allTokens, tokenAddress, chainId) => {
+    const chainTokens = Object.values(allTokens[chainId] ?? {}).flat();
 
-    return tokensForAddressAndChain.find(
+    return chainTokens.find(
       (token) => token.address.toLowerCase() === tokenAddress.toLowerCase(),
     );
   },