Refactor MASBalance component to use balance from account store and s…

…treamline balance fetching logic
massalabs · Jan 20, 2025 · 4678475 · 4678475
1 parent c8ef05c
commit 4678475
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 17 deletions.
diff --git a/src/lib/ConnectMassaWallets/components/MASBalance.tsx b/src/lib/ConnectMassaWallets/components/MASBalance.tsx
@@ -1,8 +1,6 @@
 // eslint-disable-next-line @typescript-eslint/ban-ts-comment
 // @ts-ignore
-import React from 'react';
-
-import { useEffect, useState } from 'react';
+import React, { useEffect } from 'react';
 
 import Intl from '../i18n';
 import { useAccountStore } from '../store';
@@ -11,18 +9,7 @@ import { massaToken } from '../../massa-react/utils/const';
 import { formatAmount } from '../../util/parseAmount';
 
 export function MASBalance() {
-  const [balance, setBalance] = useState<bigint>();
-
-  const { connectedAccount, currentWallet, network } = useAccountStore();
-
-  useEffect(() => {
-    if (!connectedAccount) return;
-    const fetchBalance = async () => {
-      const balance = await connectedAccount.balance(false);
-      setBalance(balance);
-    };
-    fetchBalance();
-  }, [connectedAccount, setBalance, currentWallet, network]);
+  const { balance } = useAccountStore();
 
   const formattedBalance = formatAmount(balance?.toString() || '0', 9).full;
 

diff --git a/src/lib/ConnectMassaWallets/store/accountStore.ts b/src/lib/ConnectMassaWallets/store/accountStore.ts
@@ -4,6 +4,7 @@ import { Wallet, WalletName } from '@massalabs/wallet-provider';
 
 export interface AccountStoreState {
   connectedAccount?: Provider;
+  balance?: bigint;
   accounts?: Provider[];
   currentWallet?: Wallet;
   wallets: Wallet[];
@@ -20,12 +21,15 @@ export interface AccountStoreState {
   setWallets: (wallets: Wallet[]) => void;
   setConnectedAccount: (account?: Provider) => void;
   setCurrentNetwork: (network: Network) => void;
+  fetchBalance: (provider: Provider, final: boolean) => Promise<void>;
+  refreshBalance: (final: boolean) => void;
 }
 
 export const useAccountStore = create<AccountStoreState>((set, get) => ({
   accounts: undefined,
   network: undefined,
   connectedAccount: undefined,
+  balance: undefined,
   accountObserver: undefined,
   networkObserver: undefined,
   currentWallet: undefined,
@@ -78,7 +82,7 @@ export const useAccountStore = create<AccountStoreState>((set, get) => ({
   setAccounts: async (wallet: Wallet, account?: Provider) => {
     const accounts = await wallet.accounts();
     set({ accounts });
-    set({ connectedAccount: account || accounts[0] });
+    get().setConnectedAccount(account || accounts[0]);
   },
 
   setWallets: (wallets: Wallet[]) => {
@@ -88,14 +92,27 @@ export const useAccountStore = create<AccountStoreState>((set, get) => ({
     }
   },
 
-  setConnectedAccount: (connectedAccount?: Provider) => {
+  setConnectedAccount: async (connectedAccount?: Provider) => {
     set({ connectedAccount });
+    if (!connectedAccount) return;
+    get().fetchBalance(connectedAccount, false);
   },
 
   setCurrentNetwork: (network: Network) => {
     if (network === get().network) return;
     set({ network });
   },
+
+  fetchBalance: async (provider: Provider) => {
+    const balance = await provider.balance(false);
+    set({ balance });
+  },
+
+  refreshBalance: async (final: boolean) => {
+    const { connectedAccount } = get();
+    if (!connectedAccount) return;
+    get().fetchBalance(connectedAccount, final);
+  },
 }));
 
 function resetObservers(