fix: move wagmi wallet functions into useWallet hook

oxen-io · Jul 8, 2024 · e530e76 · e530e76
1 parent 4404f7c
commit e530e76
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 27 deletions.
diff --git a/apps/staking/app/faucet/AuthModule.tsx b/apps/staking/app/faucet/AuthModule.tsx
@@ -13,7 +13,7 @@ import ActionModule, { ActionModuleDivider } from '../stake/ActionModule';
 import { FAUCET, FAUCET_ERROR } from '@/lib/constants';
 import { ButtonDataTestId } from '@/testing/data-test-ids';
 import { zodResolver } from '@hookform/resolvers/zod';
-import { CHAIN, chains } from '@session/contracts';
+import { CHAIN } from '@session/contracts';
 import { ArrowDownIcon } from '@session/ui/icons/ArrowDownIcon';
 import { toast } from '@session/ui/lib/sonner';
 import {
@@ -26,11 +26,10 @@ import {
 } from '@session/ui/ui/form';
 import { Tooltip, TooltipContent, TooltipTrigger } from '@session/ui/ui/tooltip';
 import { collapseString } from '@session/util/string';
-import { useWalletChain } from '@session/wallet/hooks/wallet-hooks';
+import { WALLET_STATUS, useWallet, useWalletChain } from '@session/wallet/hooks/wallet-hooks';
 import { ReactNode, useEffect, useMemo, useState } from 'react';
 import { useForm } from 'react-hook-form';
 import { Address, formatEther, isAddress } from 'viem';
-import { useAccount, useBalance, useDisconnect } from 'wagmi';
 import { z } from 'zod';
 import { transferTestTokens } from './actions';
 
@@ -71,14 +70,8 @@ export const AuthModule = () => {
   const [faucetError, setFaucetError] = useState<FAUCET_ERROR | null>(null);
   const { data, status: authStatus } = useSession();
   const [transactionHash, setTransactionHash] = useState<Address | null>(null);
-  const { address, status: accountStatus } = useAccount();
-  const { disconnect } = useDisconnect();
+  const { address, disconnect, status: walletStatus, ethBalance } = useWallet();
   const { chain, switchChain } = useWalletChain();
-  const { data: ethBalance } = useBalance({
-    address,
-    chainId: chains[CHAIN.TESTNET].id,
-    query: { enabled: !!address },
-  });
 
   const FormSchema = getFaucetFormSchema();
 
@@ -159,14 +152,14 @@ export const AuthModule = () => {
   };
 
   const ethAmount = useMemo(() => {
-    if (!address || ethBalance === undefined) {
-      return null;
+    if (ethBalance) {
+      return parseFloat(formatEther(ethBalance));
     }
-    return parseFloat(formatEther(ethBalance.value));
-  }, [ethBalance, address]);
+    return null;
+  }, [ethBalance]);
 
   useEffect(() => {
-    if (ethAmount !== null && ethAmount < FAUCET.MIN_ETH_BALANCE) {
+    if (address && ethAmount !== null && ethAmount < FAUCET.MIN_ETH_BALANCE) {
       form.setError('root', {
         type: 'deps',
         message: dictionary('error.insufficientEth', {
@@ -175,10 +168,10 @@ export const AuthModule = () => {
         }),
       });
     }
-  }, [address, ethAmount]);
+  }, [address, ethAmount, form]);
 
   useEffect(() => {
-    if (accountStatus === 'connected') {
+    if (walletStatus === WALLET_STATUS.CONNECTED && address) {
       form.clearErrors();
       form.setValue('walletAddress', address, {
         shouldValidate: true,
@@ -189,12 +182,12 @@ export const AuthModule = () => {
       if (chain !== CHAIN.TESTNET) {
         switchChain(CHAIN.TESTNET);
       }
-    } else if (accountStatus === 'disconnected') {
+    } else if (walletStatus === WALLET_STATUS.DISCONNECTED) {
       form.reset({ walletAddress: '' });
       form.clearErrors();
       setTransactionHash(null);
     }
-  }, [accountStatus, address, form]);
+  }, [walletStatus, address, form]);
 
   useEffect(() => {
     if (authStatus === 'authenticated') {

diff --git a/apps/staking/app/mystakes/modules/StakedNodesModule.tsx b/apps/staking/app/mystakes/modules/StakedNodesModule.tsx
@@ -17,11 +17,10 @@ import {
 } from '@session/ui/components/ModuleGrid';
 import { Button } from '@session/ui/ui/button';
 import { Switch } from '@session/ui/ui/switch';
-import { useToggleWalletModal } from '@session/wallet/hooks/wallet-hooks';
+import { useToggleWalletModal, useWallet } from '@session/wallet/hooks/wallet-hooks';
 import { useTranslations } from 'next-intl';
 import Link from 'next/link';
 import { useMemo } from 'react';
-import { useAccount } from 'wagmi';
 
 function StakedNodesWithAddress({ address }: { address: string }) {
   const showMockNodes = useFeatureFlag(FEATURE_FLAG.MOCK_STAKED_NODES);
@@ -65,7 +64,7 @@ function StakedNodesWithAddress({ address }: { address: string }) {
 
 export default function StakedNodesModule() {
   const dictionary = useTranslations('modules.stakedNodes');
-  const { address } = useAccount();
+  const { address } = useWallet();
   return (
     <>
       <ModuleGridHeader>

diff --git a/packages/wallet/components/WalletModalButton.tsx b/packages/wallet/components/WalletModalButton.tsx
@@ -16,7 +16,7 @@ export type WalletModalButtonProps = Omit<ButtonProps, 'data-testid'> & {
 };
 
 export default function WalletModalButton(props: WalletModalButtonProps) {
-  const { address, ensName, arbName, tokenBalance, ethBalance, status, isConnected } = useWallet();
+  const { address, ensName, arbName, tokenBalance, status, isConnected } = useWallet();
   const { open } = useWeb3Modal();
   const { open: isOpen } = useWeb3ModalState();
   const { isBalanceVisible } = useWalletButton();
@@ -40,13 +40,12 @@ export default function WalletModalButton(props: WalletModalButtonProps) {
       ensName={ensName}
       arbName={arbName}
       tokenBalance={tokenBalance}
-      ethBalance={ethBalance}
     />
   );
 }
 
 export type WalletButtonProps = WalletModalButtonProps &
-  Omit<ReturnType<typeof useWallet>, 'isConnected'> & {
+  Omit<ReturnType<typeof useWallet>, 'isConnected' | 'disconnect'> & {
     handleClick: () => void;
     fallbackName: string;
     isLoading?: boolean;

diff --git a/packages/wallet/hooks/wallet-hooks.tsx b/packages/wallet/hooks/wallet-hooks.tsx
@@ -5,7 +5,7 @@ import { useEns } from '@session/contracts/hooks/ens';
 import { useWeb3Modal } from '@web3modal/wagmi/react';
 import { useMemo, useState } from 'react';
 import { createWalletClient, custom, type Address } from 'viem';
-import { useAccount, useBalance, useConfig } from 'wagmi';
+import { useAccount, useBalance, useConfig, useDisconnect } from 'wagmi';
 import { switchChain as switchChainWagmi } from 'wagmi/actions';
 import { getEthereumWindowProvider } from '../lib/eth';
 import { useArbName } from './arb';
@@ -68,10 +68,13 @@ type UseWalletType = {
   status: WALLET_STATUS;
   /** Whether the wallet is connected. */
   isConnected: boolean;
+  /** Disconnect the wallet. */
+  disconnect: () => void;
 };
 
 export function useWallet(): UseWalletType {
   const { address, isConnected, isConnecting, isDisconnected, isReconnecting } = useAccount();
+  const { disconnect } = useDisconnect();
   const { balance: tokenBalanceData } = useSENTBalanceQuery({
     startEnabled: Boolean(address),
     args: [address!],
@@ -102,7 +105,17 @@ export function useWallet(): UseWalletType {
     [isConnected, isConnecting, isDisconnected, isReconnecting]
   );
 
-  return { address, ensName, ensAvatar, arbName, status, isConnected, tokenBalance, ethBalance };
+  return {
+    address,
+    ensName,
+    ensAvatar,
+    arbName,
+    status,
+    isConnected,
+    tokenBalance,
+    ethBalance,
+    disconnect,
+  };
 }
 
 export function useWalletChain() {