From 05c05fe28d8e44ba7606650b7c7d400275ac6fd3 Mon Sep 17 00:00:00 2001
From: ost-ptk <gestap9@gmail.com>
Date: Mon, 16 Sep 2024 15:32:12 +0300
Subject: [PATCH] fix issue with token balance update on transfer flow

---
 .../components/token-swticher-row.tsx         | 44 ++++++-------------
 src/apps/popup/pages/transfer/token-step.tsx  | 20 ++++++++-
 src/hooks/use-active-account-erc20-tokens.ts  |  7 ++-
 src/libs/services/erc20-service/constants.ts  |  5 ---
 .../services/erc20-service/erc20-service.ts   | 33 +-------------
 5 files changed, 38 insertions(+), 71 deletions(-)

diff --git a/src/apps/popup/pages/transfer/components/token-swticher-row.tsx b/src/apps/popup/pages/transfer/components/token-swticher-row.tsx
index 607496f2d..aee3ab65d 100644
--- a/src/apps/popup/pages/transfer/components/token-swticher-row.tsx
+++ b/src/apps/popup/pages/transfer/components/token-swticher-row.tsx
@@ -2,23 +2,13 @@ import React from 'react';
 import { Trans, useTranslation } from 'react-i18next';
 import styled from 'styled-components';
 
-import {
-  AlignedFlexRow,
-  FlexColumn,
-  ParagraphContainer,
-  SpaceBetweenFlexRow,
-  SpacingSize
-} from '@libs/layout';
+import { AlignedFlexRow, SpaceBetweenFlexRow, SpacingSize } from '@libs/layout';
 import { SvgIcon, Tile, Typography } from '@libs/ui/components';
 
 const Container = styled(SpaceBetweenFlexRow)`
   padding: 16px;
 `;
 
-const RowContainer = styled(FlexColumn)`
-  width: 100%;
-`;
-
 interface TokenSwitcherRowProps {
   tokenName: string | undefined;
   iconUrl: string | null | undefined;
@@ -31,26 +21,18 @@ export const TokenSwitcherRow = ({
   const { t } = useTranslation();
 
   return (
-    <RowContainer gap={SpacingSize.Small}>
-      <ParagraphContainer top={SpacingSize.XXL}>
-        <Typography type="bodySemiBold">
-          <Trans t={t}>Token</Trans>
-        </Typography>
-      </ParagraphContainer>
-
-      <Tile>
-        <Container gap={SpacingSize.Medium}>
-          <AlignedFlexRow gap={SpacingSize.Large} flexGrow={1}>
-            <SvgIcon src={iconUrl || ''} size={24} />
-            <Typography dataTestId="token-row" type="body">
-              {tokenName}
-            </Typography>
-          </AlignedFlexRow>
-          <Typography type="bodySemiBold" color="contentAction">
-            <Trans t={t}>Change</Trans>
+    <Tile style={{ marginTop: '8px' }}>
+      <Container gap={SpacingSize.Medium}>
+        <AlignedFlexRow gap={SpacingSize.Large} flexGrow={1}>
+          <SvgIcon src={iconUrl || ''} size={24} />
+          <Typography dataTestId="token-row" type="body">
+            {tokenName}
           </Typography>
-        </Container>
-      </Tile>
-    </RowContainer>
+        </AlignedFlexRow>
+        <Typography type="bodySemiBold" color="contentAction">
+          <Trans t={t}>Change</Trans>
+        </Typography>
+      </Container>
+    </Tile>
   );
 };
diff --git a/src/apps/popup/pages/transfer/token-step.tsx b/src/apps/popup/pages/transfer/token-step.tsx
index 7e3e80ad9..15f3c9116 100644
--- a/src/apps/popup/pages/transfer/token-step.tsx
+++ b/src/apps/popup/pages/transfer/token-step.tsx
@@ -37,11 +37,13 @@ export const TokenStep = ({
   const { t } = useTranslation();
 
   const casperToken = useCasperToken();
-  const { tokens } = useActiveAccountErc20Tokens();
+  const { tokens, isLoading } = useActiveAccountErc20Tokens();
 
   useEffect(() => {
     const tokensList: TokenType[] = [];
 
+    if (isLoading) return;
+
     if (casperToken) {
       tokensList.push(casperToken);
     }
@@ -51,7 +53,15 @@ export const TokenStep = ({
     }
 
     setTokenList(tokensList);
-  }, [casperToken, tokens]);
+
+    const token = tokensList.find(token => token.id === selectedToken?.id);
+
+    if (token) {
+      setSelectedToken(token);
+    } else {
+      setSelectedToken(casperToken);
+    }
+  }, [casperToken, tokens, isLoading]);
 
   return (
     <ContentContainer>
@@ -61,6 +71,12 @@ export const TokenStep = ({
         </Typography>
       </ParagraphContainer>
 
+      <ParagraphContainer top={SpacingSize.XXL}>
+        <Typography type="bodySemiBold">
+          <Trans t={t}>Token</Trans>
+        </Typography>
+      </ParagraphContainer>
+
       <Modal
         style={{ height: '528px' }}
         renderContent={({ closeModal }) => (
diff --git a/src/hooks/use-active-account-erc20-tokens.ts b/src/hooks/use-active-account-erc20-tokens.ts
index 5ee7b947f..a6588542b 100644
--- a/src/hooks/use-active-account-erc20-tokens.ts
+++ b/src/hooks/use-active-account-erc20-tokens.ts
@@ -22,6 +22,7 @@ export const useActiveAccountErc20Tokens = () => {
   const effectTimeoutRef = useRef<NodeJS.Timeout>();
   const forceUpdate = useForceUpdate();
   const [tokens, setTokens] = useState<TokenType[] | undefined>(undefined);
+  const [isLoading, setIsLoading] = useState(true);
   const { t } = useTranslation();
 
   const activeAccount = useSelector(selectVaultActiveAccount);
@@ -30,6 +31,7 @@ export const useActiveAccountErc20Tokens = () => {
   );
 
   useEffect(() => {
+    setIsLoading(true);
     dispatchFetchErc20TokensRequest(
       getAccountHashFromPublicKey(activeAccount?.publicKey)
     )
@@ -40,7 +42,8 @@ export const useActiveAccountErc20Tokens = () => {
       })
       .catch(error => {
         console.error('Balance request failed:', error);
-      });
+      })
+      .finally(() => setIsLoading(false));
 
     // will cause effect to run again after timeout
     effectTimeoutRef.current = setTimeout(() => {
@@ -53,5 +56,5 @@ export const useActiveAccountErc20Tokens = () => {
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [activeAccount?.publicKey, casperClarityApiUrl, t, forceUpdate]);
 
-  return { tokens: tokens };
+  return { tokens, isLoading };
 };
diff --git a/src/libs/services/erc20-service/constants.ts b/src/libs/services/erc20-service/constants.ts
index dfc5d1132..604bf4611 100644
--- a/src/libs/services/erc20-service/constants.ts
+++ b/src/libs/services/erc20-service/constants.ts
@@ -3,8 +3,3 @@ export const getErc20TokensUrl = (
   accountHash: string
 ) =>
   `${casperClarityApiUrl}/accounts/${accountHash}/erc20-tokens?fields=latest_contract,contract_package`;
-
-export const getContractPackageUrl = (
-  casperClarityApiUrl: string,
-  contractPackageHash: string
-) => `${casperClarityApiUrl}/contract-packages/${contractPackageHash}`;
diff --git a/src/libs/services/erc20-service/erc20-service.ts b/src/libs/services/erc20-service/erc20-service.ts
index d0e36fe57..c1e71fa21 100644
--- a/src/libs/services/erc20-service/erc20-service.ts
+++ b/src/libs/services/erc20-service/erc20-service.ts
@@ -7,12 +7,8 @@ import { queryClient } from '@libs/services/query-client';
 import { DataResponse, Payload } from '@libs/services/types';
 import { handleError, toJson } from '@libs/services/utils';
 
-import { getContractPackageUrl, getErc20TokensUrl } from './constants';
-import {
-  ContractPackage,
-  ContractPackageWithBalance,
-  Erc20Token
-} from './types';
+import { getErc20TokensUrl } from './constants';
+import { ContractPackageWithBalance, Erc20Token } from './types';
 
 export const erc20TokensRequest = (
   casperClarityApiUrl: string,
@@ -37,31 +33,6 @@ export const fetchErc20Tokens = ({
     { staleTime: TOKENS_REFRESH_RATE }
   );
 
-export const contractPackageRequest = (
-  casperClarityApiUrl: string,
-  contractPackageHash: string,
-  signal?: AbortSignal
-): Promise<ContractPackage> =>
-  fetch(getContractPackageUrl(casperClarityApiUrl, contractPackageHash), {
-    signal
-  })
-    .then(toJson)
-    .catch(handleError);
-
-export const fetchContractPackage = ({
-  casperClarityApiUrl,
-  contractPackageHash
-}: {
-  casperClarityApiUrl: string;
-  contractPackageHash: string;
-}) =>
-  queryClient.fetchQuery(
-    ['contractPackageRequest', casperClarityApiUrl, contractPackageHash],
-    ({ signal }) =>
-      contractPackageRequest(casperClarityApiUrl, contractPackageHash, signal),
-    { staleTime: TOKENS_REFRESH_RATE }
-  );
-
 export const dispatchFetchErc20TokensRequest = (
   accountHash: string
 ): Promise<Payload<ContractPackageWithBalance[]>> =>