Refactor ConnectMassaWallet component and account store for improved …

…state management and wallet handling
massalabs · Nov 6, 2024 · 417bb3d · 417bb3d
1 parent fa5bfe3
commit 417bb3d
Show file tree

Hide file tree

Showing 2 changed files with 118 additions and 171 deletions.
diff --git a/src/lib/ConnectMassaWallets/components/ConnectMassaWallet.tsx b/src/lib/ConnectMassaWallets/components/ConnectMassaWallet.tsx
@@ -1,9 +1,3 @@
-// eslint-disable-next-line @typescript-eslint/ban-ts-comment
-// @ts-ignore
-import React from 'react';
-
-import { useEffect, useState } from 'react';
-
 import { BearbySvg } from './BearbySvg';
 import BearbyWallet from './BearbyWallet';
 import SelectMassaWallet from './SelectMassaWallet';
@@ -19,19 +13,8 @@ export const ConnectMassaWallet = () => {
   const { currentWallet, wallets, setCurrentWallet, isFetching } =
     useAccountStore();
 
-  const [selectedWallet, setSelectedWallet] = useState<WalletName | undefined>(
-    currentWallet?.name(),
-  );
-
-  useEffect(() => {
-    const provider = wallets.find((p) => p.name() === selectedWallet);
-    if (provider && !currentWallet) {
-      setCurrentWallet(provider);
-    }
-  }, [wallets, selectedWallet, currentWallet, setCurrentWallet]);
-
   function renderWallet() {
-    switch (selectedWallet) {
+    switch (currentWallet?.name()) {
       case WalletName.MassaStation:
         return <StationWallet />;
       case WalletName.Bearby:
@@ -43,7 +26,7 @@ export const ConnectMassaWallet = () => {
   }
 
   function renderSelectedWallet() {
-    switch (selectedWallet) {
+    switch (currentWallet?.name()) {
       case WalletName.MassaStation:
         return (
           <>
@@ -61,54 +44,45 @@ export const ConnectMassaWallet = () => {
     }
   }
 
-  const noWalletSelected = !selectedWallet || isFetching;
-
-  function renderNoWalletSelected() {
+  if (!currentWallet) {
     return (
-      <SelectMassaWallet
-        onClick={(providerName) => {
-          setSelectedWallet(providerName);
-          const provider = wallets.find((p) => p.name() === providerName);
-          if (provider) {
-            setCurrentWallet(provider);
-          }
-        }}
-      />
+      <div className="text-f-primary">
+        <SelectMassaWallet
+          onClick={async (providerName) => {
+            const provider = wallets.find((p) => p.name() === providerName);
+            if (provider) {
+              await setCurrentWallet(provider);
+            }
+          }}
+        />
+      </div>
     );
   }
 
   return (
     <div className="text-f-primary">
-      {noWalletSelected ? (
-        renderNoWalletSelected()
-      ) : (
-        <>
-          <div
-            data-testid="connect-massa-wallet"
-            className="flex justify-between items-center mb-4"
-          >
-            <div className="flex gap-2 items-center">
-              {renderSelectedWallet()}
-              <ChainStatus />
-              {selectedWallet === WalletName.Bearby && (
-                <Tooltip
-                  customClass="mas-caption w-fit whitespace-nowrap"
-                  body={Intl.t(
-                    'connect-wallet.card-destination.non-massa-wallet',
-                  )}
-                />
-              )}
-            </div>
-            <SwitchWalletButton
-              onClick={() => {
-                setSelectedWallet(undefined);
-                setCurrentWallet();
-              }}
+      <div
+        data-testid="connect-massa-wallet"
+        className="flex justify-between items-center mb-4"
+      >
+        <div className="flex gap-2 items-center">
+          {renderSelectedWallet()}
+          <ChainStatus />
+          {currentWallet?.name() === WalletName.Bearby && (
+            <Tooltip
+              customClass="mas-caption w-fit whitespace-nowrap"
+              body={Intl.t('connect-wallet.card-destination.non-massa-wallet')}
             />
-          </div>
-          {renderWallet()}
-        </>
-      )}
+          )}
+        </div>
+        <SwitchWalletButton
+          onClick={() => {
+            setCurrentWallet();
+          }}
+        />
+      </div>
+
+      {!isFetching && renderWallet()}
     </div>
   );
 };
diff --git a/src/lib/ConnectMassaWallets/store/accountStore.ts b/src/lib/ConnectMassaWallets/store/accountStore.ts
@@ -1,26 +1,7 @@
-import { Provider } from '@massalabs/massa-web3';
 import { create } from 'zustand';
+import { Network, Provider } from '@massalabs/massa-web3';
 import { Wallet, WalletName } from '@massalabs/wallet-provider';
 
-async function handleBearbyAccountChange(
-  newAddress: string,
-  store: AccountStoreState,
-) {
-  const { connectedAccount, currentWallet, setConnectedAccount } = store;
-
-  const oldAddress = connectedAccount?.address;
-
-  if (newAddress !== oldAddress) {
-    const newAccounts = await currentWallet?.accounts();
-
-    if (newAccounts?.length) {
-      // Bearby returns only one account
-      const newAccount = newAccounts[0];
-      setConnectedAccount(newAccount);
-    }
-  }
-}
-
 export interface AccountStoreState {
   connectedAccount?: Provider;
   accounts?: Provider[];
@@ -33,128 +14,120 @@ export interface AccountStoreState {
   networkObserver?: {
     unsubscribe: () => void;
   };
-  chainId?: bigint;
-  network?: string;
+  network?: Network;
 
-  setCurrentWallet: (wallet?: Wallet) => void;
+  setCurrentWallet: (wallet?: Wallet, account?: Provider) => Promise<void>;
   setWallets: (wallets: Wallet[]) => void;
-
   setConnectedAccount: (account?: Provider) => void;
-  setCurrentNetwork: () => void;
+  setCurrentNetwork: (network: Network) => void;
 }
 
 export const useAccountStore = create<AccountStoreState>((set, get) => ({
   accounts: undefined,
+  network: undefined,
   connectedAccount: undefined,
   accountObserver: undefined,
   networkObserver: undefined,
   currentWallet: undefined,
   wallets: [],
   isFetching: false,
-  chainId: undefined,
-  network: undefined,
 
-  setCurrentWallet: (currentWallet?: Wallet) => {
-    try {
-      set({ isFetching: true });
+  setCurrentWallet: async (wallet?: Wallet, account?: Provider) => {
+    set({ isFetching: true });
+    if (!wallet) {
+      resetAll(get, set);
+      return;
+    }
 
-      const previousWallet = get().currentWallet;
+    if (get().currentWallet?.name() !== wallet.name()) resetObservers(get, set);
 
-      if (previousWallet?.name() !== currentWallet?.name()) {
-        get().accountObserver?.unsubscribe();
-        get().networkObserver?.unsubscribe();
-        set({ accountObserver: undefined, networkObserver: undefined });
-      }
-      if (!currentWallet) {
-        set({
-          currentWallet: undefined,
-          connectedAccount: undefined,
-          accounts: undefined,
-        });
+    if (wallet.name() === WalletName.Bearby) {
+      try {
+        await setupBearbyWallet(wallet, set, get);
+      } catch (error) {
         return;
       }
+    }
 
-      if (!get().networkObserver) {
-        const networkObserver = currentWallet.listenNetworkChanges(async () => {
-          get().setCurrentNetwork();
-        });
-        set({ networkObserver });
-      }
-
-      if (currentWallet?.name() === WalletName.Bearby) {
-        currentWallet
-          .connect()
-          .then(() => {
-            // subscribe to network events
-            const observer = currentWallet.listenAccountChanges(
-              (newAddress: string) => {
-                handleBearbyAccountChange(newAddress, get());
-              },
-            );
-
-            set({ currentWallet, accountObserver: observer });
-
-            // get connected account
-            currentWallet
-              .accounts()
-              .then((accounts) => {
-                // bearby expose only 1 account
-                get().setConnectedAccount(accounts[0]);
-                set({ accounts });
-              })
-              .catch((error) => {
-                console.warn('error getting accounts from bearby', error);
-              });
-          })
-          .catch((error) => {
-            console.warn('error connecting to bearby', error);
-          });
-        return;
-      }
-
-      set({ currentWallet });
-
-      get().setCurrentNetwork();
+    if (!get().networkObserver) {
+      const networkObserver = wallet.listenNetworkChanges(async () => {
+        get().setCurrentNetwork(await wallet.networkInfos());
+      });
+      set({ networkObserver });
+    }
 
-      currentWallet
-        .accounts()
-        .then((accounts) => {
-          set({ accounts });
+    set({ currentWallet: wallet });
+    const network = await wallet.networkInfos();
+    get().setCurrentNetwork(network);
+    const accounts = await wallet.accounts();
+    set({ accounts });
+    get().setConnectedAccount(account || accounts[0]);
 
-          const selectedAccount = accounts[0];
-          get().setConnectedAccount(selectedAccount);
-        })
-        .catch((error) => {
-          console.warn('error getting accounts from wallet', error);
-        });
-    } finally {
-      set({ isFetching: false });
-    }
+    set({ isFetching: false });
   },
 
   setWallets: (wallets: Wallet[]) => {
     set({ wallets });
-
-    // if current wallet is not in the new list of wallets, unset it
     if (!wallets.some((p) => p.name() === get().currentWallet?.name())) {
-      set({
-        currentWallet: undefined,
-        connectedAccount: undefined,
-        accounts: undefined,
-      });
+      resetWallet(set);
     }
   },
 
-  // set the connected account, and update the massa client
-  setConnectedAccount: async (connectedAccount?: Provider) => {
+  setConnectedAccount: (connectedAccount?: Provider) => {
     set({ connectedAccount });
   },
 
-  setCurrentNetwork: () => {
-    get()
-      .currentWallet?.networkInfos()
-      .then((infos) => {
-        set({ chainId: infos.chainId, network: infos.name });
-      });
+  setCurrentNetwork: (network: Network) => {
+    if (network === get().network) return;
+    set({ network });
   },
 }));
+
+function resetObservers(
+  get: () => AccountStoreState,
+  set: (partial: Partial<AccountStoreState>) => void,
+) {
+  get().accountObserver?.unsubscribe();
+  get().networkObserver?.unsubscribe();
+  set({ accountObserver: undefined, networkObserver: undefined });
+}
+
+function resetWallet(set: (partial: Partial<AccountStoreState>) => void) {
+  set({
+    accounts: undefined,
+    connectedAccount: undefined,
+    currentWallet: undefined,
+    network: undefined,
+  });
+}
+
+function resetAll(
+  get: () => AccountStoreState,
+  set: (partial: Partial<AccountStoreState>) => void,
+) {
+  resetObservers(get, set);
+  resetWallet(set);
+  set({ isFetching: false });
+}
+
+async function setupBearbyWallet(
+  wallet: Wallet,
+  set: (partial: Partial<AccountStoreState>) => void,
+  get: () => AccountStoreState,
+) {
+  await wallet.connect();
+
+  const observer = wallet.listenAccountChanges(async (newAddress: string) => {
+    const { connectedAccount, currentWallet, setConnectedAccount } = get();
+
+    if (!currentWallet || !connectedAccount) return;
+
+    if (newAddress !== connectedAccount.address) {
+      const accounts = await currentWallet.accounts();
+      const newAccount = accounts.find((acc) => acc.address === newAddress);
+      setConnectedAccount(newAccount);
+    }
+  });
+
+  set({ accountObserver: observer });
+}