From 35cdf27aef83bc011202a1dc366377ccfb8511e2 Mon Sep 17 00:00:00 2001
From: Nenad Vracar <34926005+nenadV91@users.noreply.github.com>
Date: Mon, 31 Jan 2022 11:52:59 +0100
Subject: [PATCH] Fix change or disconnect wallet after claim confirm (#2337)

* Fix change or disconnect wallet after claim confirm

* Update after merge

* Handle account change better

* Reset isSearchUsed on switch to connected account
---
 src/custom/pages/Claim/ClaimNav.tsx |  9 +++--
 src/custom/pages/Claim/index.tsx    | 54 +++++++++++++++++++++--------
 2 files changed, 46 insertions(+), 17 deletions(-)

diff --git a/src/custom/pages/Claim/ClaimNav.tsx b/src/custom/pages/Claim/ClaimNav.tsx
index 3386f52d9..4ef395621 100644
--- a/src/custom/pages/Claim/ClaimNav.tsx
+++ b/src/custom/pages/Claim/ClaimNav.tsx
@@ -11,7 +11,7 @@ type ClaimNavProps = Pick<ClaimCommonTypes, 'account' | 'handleChangeAccount'>
 
 export default function ClaimNav({ account, handleChangeAccount }: ClaimNavProps) {
   const { activeClaimAccount, activeClaimAccountENS, claimStatus, investFlowStep } = useClaimState()
-  const { setActiveClaimAccount } = useClaimDispatchers()
+  const { setActiveClaimAccount, setIsSearchUsed } = useClaimDispatchers()
 
   const isDefaultStatus = claimStatus === ClaimStatus.DEFAULT
   const isConfirmed = claimStatus === ClaimStatus.CONFIRMED
@@ -36,7 +36,12 @@ export default function ClaimNav({ account, handleChangeAccount }: ClaimNavProps
           ) : (
             !!account &&
             allowToChangeAccount && (
-              <ButtonSecondary onClick={() => setActiveClaimAccount(account)}>
+              <ButtonSecondary
+                onClick={() => {
+                  setActiveClaimAccount(account)
+                  setIsSearchUsed(false)
+                }}
+              >
                 Switch to connected account
               </ButtonSecondary>
             )
diff --git a/src/custom/pages/Claim/index.tsx b/src/custom/pages/Claim/index.tsx
index 72b0e4729..8985a386f 100644
--- a/src/custom/pages/Claim/index.tsx
+++ b/src/custom/pages/Claim/index.tsx
@@ -27,6 +27,8 @@ import FooterNavButtons from './FooterNavButtons'
 import InvestmentFlow from './InvestmentFlow'
 import { ClaimSummary } from './ClaimSummary'
 
+import usePrevious from 'hooks/usePrevious'
+
 /* TODO: Replace URLs with the actual final URL destinations */
 export const COW_LINKS = {
   vCowPost: 'https://cow.fi/',
@@ -37,6 +39,9 @@ export default function Claim() {
   const { account, chainId } = useActiveWeb3React()
   const { error } = useWeb3React()
 
+  // get previous account
+  const previousAccount = usePrevious(account)
+
   const {
     // address/ENS address
     inputAddress,
@@ -95,20 +100,32 @@ export default function Claim() {
   const { claimCallback, estimateGasCallback } = useClaimCallback(activeClaimAccount)
 
   // reset claim state
-  const resetClaimState = useCallback(() => {
-    setClaimStatus(ClaimStatus.DEFAULT)
-    setActiveClaimAccount('')
-    setActiveClaimAccountENS('')
-    setSelected([])
-  }, [setActiveClaimAccount, setActiveClaimAccountENS, setClaimStatus, setSelected])
-
-  // handle change account
-  const handleChangeAccount = () => {
+  const resetClaimState = useCallback(
+    (account = '', ens = '') => {
+      setClaimStatus(ClaimStatus.DEFAULT)
+      setActiveClaimAccount(account)
+      setActiveClaimAccountENS(ens)
+      setSelected([])
+    },
+    [setActiveClaimAccount, setActiveClaimAccountENS, setClaimStatus, setSelected]
+  )
+
+  // handle account change
+  const handleAccountChange = useCallback(
+    (account = '') => {
+      resetClaimState(account)
+      setIsSearchUsed(false)
+    },
+    [resetClaimState, setIsSearchUsed]
+  )
+
+  // handle change account click
+  const handleChangeClick = useCallback(() => {
     resetClaimState()
     setIsSearchUsed(true)
-  }
+  }, [resetClaimState, setIsSearchUsed])
 
-  // check claim
+  // handle
   const handleCheckClaim = () => {
     setActiveClaimAccount(resolvedAddress || '')
     setActiveClaimAccountENS(resolvedENS || '')
@@ -182,7 +199,7 @@ export default function Claim() {
 
     // handle unsupported network
     if (error instanceof UnsupportedChainIdError) {
-      resetClaimState()
+      handleAccountChange()
     }
 
     // properly reset the user to the claims table and initial investment flow
@@ -197,9 +214,16 @@ export default function Claim() {
     setActiveClaimAccount,
     resetClaimUi,
     error,
-    resetClaimState,
+    handleAccountChange,
   ])
 
+  // handle account disconnect or account change after claim is confirmed
+  useEffect(() => {
+    if (!account || (account !== previousAccount && claimStatus === ClaimStatus.CONFIRMED)) {
+      handleAccountChange(account || '')
+    }
+  }, [account, claimStatus, previousAccount, handleAccountChange])
+
   // Transaction confirmation modal
   const { TransactionConfirmationModal, openModal, closeModal } = useTransactionConfirmationModal(
     OperationType.APPROVE_TOKEN
@@ -222,7 +246,7 @@ export default function Claim() {
             {/* If claim is confirmed > trigger confetti effect */}
             <Confetti start={claimStatus === ClaimStatus.CONFIRMED} />
             {/* Top nav buttons */}
-            <ClaimNav account={account} handleChangeAccount={handleChangeAccount} />
+            <ClaimNav account={account} handleChangeAccount={handleChangeClick} />
             {/* Show general title OR total to claim (user has airdrop or airdrop+investment) --------------------------- */}
             <EligibleBanner hasClaims={hasClaims} />
             {/* Show total to claim (user has airdrop or airdrop+investment) */}
@@ -233,7 +257,7 @@ export default function Claim() {
             <CanUserClaimMessage
               hasClaims={hasClaims}
               isAirdropOnly={isAirdropOnly}
-              handleChangeAccount={handleChangeAccount}
+              handleChangeAccount={handleChangeClick}
             />
 
             {/* Try claiming or inform successful claim */}