From f929b2402a97307e292cba137ce3f34d9bc51744 Mon Sep 17 00:00:00 2001
From: Atatakai <atatakai@Atatakais-MacBook-Pro.local>
Date: Fri, 24 May 2024 17:01:58 +0400
Subject: [PATCH] Fix balances getters

---
 frontend/components/Main/MainGasBalance.tsx   |  8 ++---
 frontend/components/Main/MainHeader.tsx       | 20 ++++++-----
 frontend/components/Main/MainNeedsFunds.tsx   | 30 +++++-----------
 .../Settings/SettingsAddBackupWallet.tsx      |  8 ++---
 .../Setup/Create/SetupEoaFunding.tsx          | 34 +++++++------------
 frontend/context/BalanceProvider.tsx          | 16 +++++++++
 frontend/hooks/useBalance.ts                  |  4 +++
 frontend/types/Records.ts                     |  7 +++-
 8 files changed, 66 insertions(+), 61 deletions(-)

diff --git a/frontend/components/Main/MainGasBalance.tsx b/frontend/components/Main/MainGasBalance.tsx
index f92751e80..0ac8dd8cc 100644
--- a/frontend/components/Main/MainGasBalance.tsx
+++ b/frontend/components/Main/MainGasBalance.tsx
@@ -34,19 +34,19 @@ const LowDot = styled(Dot)`
 `;
 
 const BalanceStatus = () => {
-  const { totalEthBalance } = useBalance();
+  const { safeBalance } = useBalance();
 
   const status = useMemo(() => {
-    if (!totalEthBalance || totalEthBalance === 0) {
+    if (!safeBalance || safeBalance.ETH === 0) {
       return { statusName: 'Empty', StatusComponent: EmptyDot };
     }
 
-    if (totalEthBalance < LOW_BALANCE) {
+    if (safeBalance.ETH < LOW_BALANCE) {
       return { statusName: 'Low', StatusComponent: LowDot };
     }
 
     return { statusName: 'Fine', StatusComponent: FineDot };
-  }, [totalEthBalance]);
+  }, [safeBalance]);
 
   const { statusName, StatusComponent } = status;
   return (
diff --git a/frontend/components/Main/MainHeader.tsx b/frontend/components/Main/MainHeader.tsx
index 79934f29f..fa2de48ed 100644
--- a/frontend/components/Main/MainHeader.tsx
+++ b/frontend/components/Main/MainHeader.tsx
@@ -30,8 +30,10 @@ export const MainHeader = () => {
   const { getServiceTemplates } = useServiceTemplates();
   const { wallets, masterSafeAddress } = useWallet();
   const {
+    safeBalance,
     totalOlasBalance,
     totalEthBalance,
+    isBalanceLoaded,
     setIsPaused: setIsBalancePollingPaused,
   } = useBalance();
 
@@ -44,14 +46,14 @@ export const MainHeader = () => {
   );
 
   useEffect(() => {
-    if (totalEthBalance && totalEthBalance < LOW_BALANCE) {
+    if (safeBalance && safeBalance.ETH < LOW_BALANCE) {
       setTrayIcon?.('low-gas');
     } else if (serviceStatus === DeploymentStatus.DEPLOYED) {
       setTrayIcon?.('running');
     } else if (serviceStatus === DeploymentStatus.STOPPED) {
       setTrayIcon?.('paused');
     }
-  }, [totalEthBalance, serviceStatus, setTrayIcon]);
+  }, [safeBalance, serviceStatus, setTrayIcon]);
 
   const agentHead = useMemo(() => {
     if (
@@ -185,7 +187,7 @@ export const MainHeader = () => {
       );
     }
 
-    if (totalOlasBalance === undefined || totalEthBalance === undefined) {
+    if (!isBalanceLoaded) {
       return (
         <Button type="primary" size="large" disabled>
           Start agent
@@ -213,8 +215,9 @@ export const MainHeader = () => {
     );
 
     if (
-      totalOlasBalance < olasCostOfBond + olasRequiredToStake ||
-      totalEthBalance < monthlyGasEstimate
+      (totalOlasBalance &&
+        totalOlasBalance < olasCostOfBond + olasRequiredToStake) ||
+      (totalEthBalance && totalEthBalance < monthlyGasEstimate)
     ) {
       return (
         <Button type="default" size="large" disabled>
@@ -229,12 +232,13 @@ export const MainHeader = () => {
       </Button>
     );
   }, [
+    handlePause,
+    handleStart,
+    isBalanceLoaded,
     serviceButtonState,
     serviceStatus,
-    totalOlasBalance,
     totalEthBalance,
-    handleStart,
-    handlePause,
+    totalOlasBalance,
   ]);
 
   return (
diff --git a/frontend/components/Main/MainNeedsFunds.tsx b/frontend/components/Main/MainNeedsFunds.tsx
index ce6e1662d..fd6e67bb1 100644
--- a/frontend/components/Main/MainNeedsFunds.tsx
+++ b/frontend/components/Main/MainNeedsFunds.tsx
@@ -17,7 +17,7 @@ const COVER_PREV_BLOCK_BORDER_STYLE = { marginTop: '-1px' };
 const useNeedsFunds = () => {
   const serviceTemplate = SERVICE_TEMPLATES[0];
   const { storeState } = useStore();
-  const { totalEthBalance, totalOlasBalance } = useBalance();
+  const { safeBalance } = useBalance();
 
   const isInitialFunded = storeState?.isInitialFunded as boolean | undefined;
 
@@ -45,13 +45,13 @@ const useNeedsFunds = () => {
   ]);
 
   const hasEnoughEth = useMemo(
-    () => (totalEthBalance || 0) >= (serviceFundRequirements?.eth || 0),
-    [serviceFundRequirements?.eth, totalEthBalance],
+    () => (safeBalance?.ETH || 0) >= (serviceFundRequirements?.eth || 0),
+    [serviceFundRequirements?.eth, safeBalance],
   );
 
   const hasEnoughOlas = useMemo(
-    () => (totalOlasBalance || 0) >= (serviceFundRequirements?.olas || 0),
-    [serviceFundRequirements?.olas, totalOlasBalance],
+    () => (safeBalance?.OLAS || 0) >= (serviceFundRequirements?.olas || 0),
+    [serviceFundRequirements?.olas, safeBalance],
   );
 
   return {
@@ -63,7 +63,7 @@ const useNeedsFunds = () => {
 };
 
 export const MainNeedsFunds = () => {
-  const { isBalanceLoaded, totalEthBalance, totalOlasBalance } = useBalance();
+  const { isBalanceLoaded } = useBalance();
   const {
     hasEnoughEth,
     hasEnoughOlas,
@@ -75,23 +75,10 @@ export const MainNeedsFunds = () => {
 
   const isVisible: boolean = useMemo(() => {
     if (isInitialFunded) return false;
-    if (
-      [totalEthBalance, totalOlasBalance].some(
-        (balance) => balance === undefined,
-      )
-    ) {
-      return false;
-    }
-
+    if (!isBalanceLoaded) return false;
     if (hasEnoughEth && hasEnoughOlas) return false;
     return true;
-  }, [
-    hasEnoughEth,
-    hasEnoughOlas,
-    isInitialFunded,
-    totalEthBalance,
-    totalOlasBalance,
-  ]);
+  }, [hasEnoughEth, hasEnoughOlas, isBalanceLoaded, isInitialFunded]);
 
   const message: ReactNode = useMemo(
     () => (
@@ -130,7 +117,6 @@ export const MainNeedsFunds = () => {
   }, [electronApi.store, hasEnoughEth, hasEnoughOlas, isInitialFunded]);
 
   if (!isVisible) return null;
-  if (!isBalanceLoaded) return null;
 
   return (
     <CardSection style={COVER_PREV_BLOCK_BORDER_STYLE}>
diff --git a/frontend/components/Settings/SettingsAddBackupWallet.tsx b/frontend/components/Settings/SettingsAddBackupWallet.tsx
index bc31d5e12..2555b8f0f 100644
--- a/frontend/components/Settings/SettingsAddBackupWallet.tsx
+++ b/frontend/components/Settings/SettingsAddBackupWallet.tsx
@@ -12,17 +12,17 @@ import { CardTitle } from '../common/CardTitle';
 import { CardFlex } from '../styled/CardFlex';
 
 export const SettingsAddBackupWallet = () => {
-  const { totalEthBalance } = useBalance();
+  const { eoaBalance } = useBalance();
   const { goto } = useSettings();
 
   const [form] = Form.useForm();
 
   const isFunded = useMemo<boolean>(() => {
-    if (!totalEthBalance) return false;
+    if (!eoaBalance) return false;
     return (
-      totalEthBalance >= MIN_ETH_BALANCE_THRESHOLDS[Chain.GNOSIS].safeAddSigner
+      eoaBalance.ETH >= MIN_ETH_BALANCE_THRESHOLDS[Chain.GNOSIS].safeAddSigner
     );
-  }, [totalEthBalance]);
+  }, [eoaBalance]);
 
   return (
     <CardFlex
diff --git a/frontend/components/Setup/Create/SetupEoaFunding.tsx b/frontend/components/Setup/Create/SetupEoaFunding.tsx
index 6ee024f32..6f0d8ad4c 100644
--- a/frontend/components/Setup/Create/SetupEoaFunding.tsx
+++ b/frontend/components/Setup/Create/SetupEoaFunding.tsx
@@ -27,22 +27,16 @@ import { UNICODE_SYMBOLS } from '@/constants/unicode';
 import { SetupScreen } from '@/enums';
 import { useBalance, useSetup } from '@/hooks';
 import { useWallet } from '@/hooks/useWallet';
-import { Address } from '@/types';
 
 import { SetupCreateHeader } from './SetupCreateHeader';
 
 export const SetupEoaFunding = () => {
-  const { masterEoaAddress } = useWallet();
-  const { walletBalances } = useBalance();
+  const { eoaBalance } = useBalance();
   const { goto } = useSetup();
 
-  const masterEaoEthBalance =
-    masterEoaAddress && walletBalances?.[masterEoaAddress]?.ETH;
-
   const isFundedMasterEoa =
-    masterEaoEthBalance &&
-    masterEaoEthBalance >=
-      MIN_ETH_BALANCE_THRESHOLDS[Chain.GNOSIS].safeCreation;
+    eoaBalance?.ETH &&
+    eoaBalance.ETH >= MIN_ETH_BALANCE_THRESHOLDS[Chain.GNOSIS].safeCreation;
 
   const statusMessage = useMemo(() => {
     if (isFundedMasterEoa) {
@@ -80,18 +74,14 @@ export const SetupEoaFunding = () => {
           Status: {statusMessage}
         </Typography.Text>
       </CardSection>
-      {!isFundedMasterEoa && (
-        <SetupEoaFundingWaiting masterEoa={masterEoaAddress} />
-      )}
+      {!isFundedMasterEoa && <SetupEoaFundingWaiting />}
     </CardFlex>
   );
 };
 
-const SetupEoaFundingWaiting = ({
-  masterEoa,
-}: {
-  masterEoa: Address | undefined;
-}) => {
+const SetupEoaFundingWaiting = () => {
+  const { masterEoaAddress } = useWallet();
+
   return (
     <>
       <CardSection>
@@ -121,21 +111,21 @@ const SetupEoaFundingWaiting = ({
               <CopyOutlined
                 style={ICON_STYLE}
                 onClick={() =>
-                  masterEoa &&
-                  copyToClipboard(masterEoa).then(() =>
+                  masterEoaAddress &&
+                  copyToClipboard(masterEoaAddress).then(() =>
                     message.success('Address copied!'),
                   )
                 }
               />
             </Tooltip>
 
-            {/* {masterEoa && (
+            {/* {masterEoaAddress && (
                 <Popover
                   title="Scan QR code"
                   content={
                     <QRCode
                       size={250}
-                      value={`https://metamask.app.link/send/${masterEoa}@${100}`}
+                      value={`https://metamask.app.link/send/${masterEoaAddress}@${100}`}
                     />
                   }
                 >
@@ -146,7 +136,7 @@ const SetupEoaFundingWaiting = ({
         </Flex>
 
         <span className="can-select-text break-word">
-          {`GNO: ${masterEoa}`}
+          {`GNO: ${masterEoaAddress}`}
         </span>
         <Alert
           type="info"
diff --git a/frontend/context/BalanceProvider.tsx b/frontend/context/BalanceProvider.tsx
index 1ae94c1ee..7aa9d0c6e 100644
--- a/frontend/context/BalanceProvider.tsx
+++ b/frontend/context/BalanceProvider.tsx
@@ -1,6 +1,7 @@
 import { message } from 'antd';
 import { isAddress } from 'ethers/lib/utils';
 import { isNumber } from 'lodash';
+import { ValueOf } from 'next/dist/shared/lib/constants';
 import {
   createContext,
   Dispatch,
@@ -36,6 +37,8 @@ export const BalanceContext = createContext<{
   isBalanceLoaded: boolean;
   olasBondBalance?: number;
   olasDepositBalance?: number;
+  eoaBalance?: ValueOf<WalletAddressNumberRecord>;
+  safeBalance?: ValueOf<WalletAddressNumberRecord>;
   totalEthBalance?: number;
   totalOlasBalance?: number;
   wallets?: Wallet[];
@@ -49,6 +52,8 @@ export const BalanceContext = createContext<{
   isBalanceLoaded: false,
   olasBondBalance: undefined,
   olasDepositBalance: undefined,
+  eoaBalance: undefined,
+  safeBalance: undefined,
   totalEthBalance: undefined,
   totalOlasBalance: undefined,
   wallets: undefined,
@@ -179,6 +184,15 @@ export const BalanceProvider = ({ children }: PropsWithChildren) => {
     }
   }, [masterEoaAddress, masterSafeAddress, serviceAddresses, services]);
 
+  const eoaBalance = useMemo(
+    () => masterEoaAddress && walletBalances[masterEoaAddress],
+    [masterEoaAddress, walletBalances],
+  );
+  const safeBalance = useMemo(
+    () => masterSafeAddress && walletBalances[masterSafeAddress],
+    [masterSafeAddress, walletBalances],
+  );
+
   useInterval(
     () => {
       updateBalances();
@@ -194,6 +208,8 @@ export const BalanceProvider = ({ children }: PropsWithChildren) => {
         isBalanceLoaded,
         olasBondBalance,
         olasDepositBalance,
+        eoaBalance,
+        safeBalance,
         totalEthBalance,
         totalOlasBalance,
         wallets,
diff --git a/frontend/hooks/useBalance.ts b/frontend/hooks/useBalance.ts
index e40740f5c..4132572a1 100644
--- a/frontend/hooks/useBalance.ts
+++ b/frontend/hooks/useBalance.ts
@@ -7,6 +7,8 @@ export const useBalance = () => {
     isLoaded,
     setIsLoaded,
     isBalanceLoaded,
+    eoaBalance,
+    safeBalance,
     totalEthBalance,
     totalOlasBalance,
     wallets,
@@ -20,6 +22,8 @@ export const useBalance = () => {
     isLoaded,
     setIsLoaded,
     isBalanceLoaded,
+    eoaBalance,
+    safeBalance,
     totalEthBalance,
     totalOlasBalance,
     wallets,
diff --git a/frontend/types/Records.ts b/frontend/types/Records.ts
index 98b6ab184..616119e0d 100644
--- a/frontend/types/Records.ts
+++ b/frontend/types/Records.ts
@@ -1,7 +1,12 @@
+import { Token } from '@/enums';
+
 import { Address } from '.';
 
 export type AddressNumberRecord = Record<Address, number>;
 export type AddressBooleanRecord = Record<Address, boolean>;
 
 // defines token balances in a wallet by token name
-export type WalletAddressNumberRecord = Record<Address, Record<string, number>>;
+export type WalletAddressNumberRecord = Record<
+  Address,
+  Record<Token.ETH | Token.OLAS, number>
+>;