fix(token-input): display correct USD value for selected balance

src/components/pageComponents/home/Examples/demos/TokenInput/index.test.tsx

@@ -35,6 +35,8 @@ vi.mock('@/src/components/sharedComponents/TokenInput/useTokenInput', () => ({
     balance: 0n,
     balanceError: null,
     isLoadingBalance: false,
+    isLoadingPrice: false,
+    priceUSD: undefined,
     selectedToken: undefined,
     setTokenSelected: vi.fn(),
   })),
@@ -47,6 +49,7 @@ vi.mock('@/src/hooks/useTokens', () => ({
     tokens: [],
     tokensByChainId: { 1: [], 10: [], 42161: [], 137: [], 11155111: [] },
     isLoadingBalances: false,
+    isLoadingPrices: false,
   })),
 }))
 

src/components/sharedComponents/TokenInput/index.tsx

@@ -85,6 +85,8 @@ const TokenInput: FC<Props> = ({
     balance,
     balanceError,
     isLoadingBalance,
+    isLoadingPrice,
+    priceUSD,
     selectedToken,
     setAmount,
     setAmountError,
@@ -105,12 +107,10 @@ const TokenInput: FC<Props> = ({
 
   const estimatedUSDValue = useMemo(() => {
     if (isTestnetChain) return null
-    if (!selectedToken) return 0
-    const priceUSD = selectedToken.extensions?.priceUSD
-    if (priceUSD === undefined || priceUSD === null) return 0
-    const tokenAmount = Number.parseFloat(formatUnits(amount, selectedToken.decimals ?? 0))
-    return Number.parseFloat(priceUSD as string) * tokenAmount
-  }, [isTestnetChain, selectedToken, amount])
+    if (!selectedToken || !priceUSD || !balance) return 0
+    const tokenBalance = Number.parseFloat(formatUnits(balance, selectedToken.decimals ?? 0))
+    return Number.parseFloat(priceUSD) * tokenBalance
+  }, [isTestnetChain, selectedToken, priceUSD, balance])
 
   const selectIconSize = 24
   const decimals = selectedToken ? selectedToken.decimals : 2
@@ -188,7 +188,13 @@ const TokenInput: FC<Props> = ({
         </TopRow>
         <BottomRow>
           <EstimatedUSDValue>
-            {estimatedUSDValue === null ? NO_PRICE_DATA_LABEL : `~$${estimatedUSDValue.toFixed(2)}`}
+            {estimatedUSDValue === null ? (
+              NO_PRICE_DATA_LABEL
+            ) : selectedToken && (isLoadingPrice || isLoadingBalance) ? (
+              <Spinner size="sm" />
+            ) : (
+              `~$${estimatedUSDValue.toFixed(2)}`
+            )}
           </EstimatedUSDValue>
           <Balance>
             <BalanceValue>

src/components/sharedComponents/TokenInput/useTokenInput.test.ts

@@ -11,6 +11,7 @@ const walletAddress = '0x71C7656EC7ab88b098defB751B7401B5f6d8976F' as const
 const mockUseAccount = vi.fn()
 const mockUsePublicClient = vi.fn()
 const mockGetBalance = vi.fn()
+const mockUseTokens = vi.fn()
 
 vi.mock('wagmi', () => ({
   useAccount: () => mockUseAccount(),
@@ -24,6 +25,10 @@ vi.mock('@/src/hooks/useErc20Balance', () => ({
   useErc20Balance: () => ({ balance: undefined, balanceError: null, isLoadingBalance: false }),
 }))
 
+vi.mock('@/src/hooks/useTokens', () => ({
+  useTokens: (args: unknown) => mockUseTokens(args),
+}))
+
 vi.mock('@/src/env', () => ({
   env: { PUBLIC_NATIVE_TOKEN_ADDRESS: zeroAddress.toLowerCase() },
 }))
@@ -56,6 +61,12 @@ describe('useTokenInput', () => {
     mockUseAccount.mockReturnValue({ address: walletAddress })
     mockUsePublicClient.mockClear()
     mockGetBalance.mockReset()
+    mockUseTokens.mockReturnValue({
+      tokensByChainId: {},
+      isLoadingBalances: false,
+      isLoadingPrices: false,
+      tokens: [],
+    })
   })
 
   it('rebinds the native public client to the selected token chain when the user switches chains', async () => {
@@ -98,4 +109,59 @@ describe('useTokenInput', () => {
     expect(result.current.isLoadingBalance).toBe(false)
     expect(mockGetBalance).not.toHaveBeenCalled()
   })
+
+  it('exposes priceUSD for the selected token from useTokens', async () => {
+    mockUseTokens.mockReturnValue({
+      tokensByChainId: {
+        1: [{ ...mainnetUsdc, extensions: { priceUSD: '1.00' } }],
+      },
+      isLoadingBalances: false,
+      isLoadingPrices: false,
+      tokens: [],
+    })
+
+    const { result } = renderHook(() => useTokenInput(mainnetUsdc), { wrapper })
+
+    await waitFor(() => expect(result.current.priceUSD).toBe('1.00'))
+  })
+
+  it('exposes isLoadingPrice as true while useTokens is loading', () => {
+    mockUseTokens.mockReturnValue({
+      tokensByChainId: {},
+      isLoadingBalances: true,
+      isLoadingPrices: true,
+      tokens: [],
+    })
+
+    const { result } = renderHook(() => useTokenInput(mainnetUsdc), { wrapper })
+
+    expect(result.current.isLoadingPrice).toBe(true)
+  })
+
+  it('updates priceUSD when a different token is selected', async () => {
+    const mainnetEth: Token = {
+      address: zeroAddress,
+      chainId: 1,
+      decimals: 18,
+      name: 'Ether',
+      symbol: 'ETH',
+    }
+    mockUseTokens.mockReturnValue({
+      tokensByChainId: {
+        1: [{ ...mainnetEth, extensions: { priceUSD: '3000.00' } }],
+      },
+      isLoadingBalances: false,
+      isLoadingPrices: false,
+      tokens: [],
+    })
+    mockGetBalance.mockResolvedValue(1000000000000000000n)
+
+    const { result } = renderHook(() => useTokenInput(mainnetUsdc), { wrapper })
+
+    act(() => {
+      result.current.setTokenSelected(mainnetEth)
+    })
+
+    await waitFor(() => expect(result.current.priceUSD).toBe('3000.00'))
+  })
 })

src/components/sharedComponents/TokenInput/useTokenInput.tsx

@@ -3,6 +3,7 @@ import { useEffect, useState } from 'react'
 import { getAddress } from 'viem'
 import { useAccount, usePublicClient } from 'wagmi'
 import { useErc20Balance } from '@/src/hooks/useErc20Balance'
+import { useTokens } from '@/src/hooks/useTokens'
 import type { Token } from '@/src/types/token'
 import { isNativeToken } from '@/src/utils/address'
 
@@ -25,6 +26,8 @@ export type UseTokenInputReturnType = ReturnType<typeof useTokenInput>
  * @returns {bigint} returns.balance - Current token balance (ERC20 or native)
  * @returns {Error|null} returns.balanceError - Error from balance fetching
  * @returns {boolean} returns.isLoadingBalance - Loading state for balance
+ * @returns {string|undefined} returns.priceUSD - USD price of the selected token (from useTokens)
+ * @returns {boolean} returns.isLoadingPrice - Loading state for the selected token's USD price
  * @returns {Token|undefined} returns.selectedToken - Currently selected token
  * @returns {function} returns.setTokenSelected - Function to update selected token
  *
@@ -49,6 +52,16 @@ export function useTokenInput(token?: Token) {
   }, [token])
 
   const { address: userWallet } = useAccount()
+  const { tokensByChainId, isLoadingPrices: isLoadingPrice } = useTokens({
+    chainId: selectedToken?.chainId,
+    withBalance: true,
+  })
+  const priceUSD = selectedToken
+    ? (tokensByChainId[selectedToken.chainId]?.find(
+        (t) => t.address.toLowerCase() === selectedToken.address.toLowerCase(),
+      )?.extensions?.priceUSD as string | undefined)
+    : undefined
+
   const { balance, balanceError, isLoadingBalance } = useErc20Balance({
     address: userWallet ? getAddress(userWallet) : undefined,
     token: selectedToken,
@@ -75,6 +88,8 @@ export function useTokenInput(token?: Token) {
     balance: isNative ? nativeBalance : balance,
     balanceError: isNative ? nativeBalanceError : balanceError,
     isLoadingBalance: isNative ? isLoadingNativeBalance : isLoadingBalance,
+    isLoadingPrice,
+    priceUSD,
     selectedToken,
     setTokenSelected,
   }

src/hooks/useTokens.ts

@@ -47,6 +47,7 @@ export const lifiConfig = createConfig({
  * @returns {Token[]} returns.tokens - Array of tokens with price and balance information
  * @returns {Record<number, Token[]>} returns.tokensByChainId - Tokens organized by chain ID
  * @returns {boolean} returns.isLoadingBalances - Loading state for token balances and prices
+ * @returns {boolean} returns.isLoadingPrices - Loading state for token prices only
  *
  * @example
  * ```tsx
@@ -150,6 +151,7 @@ export const useTokens = (
   return {
     ...cache,
     isLoadingBalances: Boolean(isLoadingChains || isLoadingBalances || isLoadingPrices),
+    isLoadingPrices: Boolean(isLoadingChains || isLoadingPrices),
   }
 }