From 369b75bc065e709fa8972954b87aefc81d8cdc90 Mon Sep 17 00:00:00 2001
From: BenRey <44082144+Ben-Rey@users.noreply.github.com>
Date: Mon, 14 Oct 2024 12:52:11 +0200
Subject: [PATCH] Add network management to account store and update state
 handling (#491)

---
 .../ConnectMassaWallets/store/accountStore.ts | 27 ++++++++++---------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/src/lib/ConnectMassaWallets/store/accountStore.ts b/src/lib/ConnectMassaWallets/store/accountStore.ts
index d5c0acb7..b2e4b43a 100644
--- a/src/lib/ConnectMassaWallets/store/accountStore.ts
+++ b/src/lib/ConnectMassaWallets/store/accountStore.ts
@@ -34,11 +34,13 @@ export interface AccountStoreState {
     unsubscribe: () => void;
   };
   chainId?: bigint;
+  network?: string;
 
   setCurrentWallet: (wallet?: Wallet) => void;
   setWallets: (wallets: Wallet[]) => void;
 
   setConnectedAccount: (account?: Provider) => void;
+  setCurrentNetwork: () => void;
 }
 
 export const useAccountStore = create<AccountStoreState>((set, get) => ({
@@ -50,6 +52,7 @@ export const useAccountStore = create<AccountStoreState>((set, get) => ({
   wallets: [],
   isFetching: false,
   chainId: undefined,
+  network: undefined,
 
   setCurrentWallet: (currentWallet?: Wallet) => {
     try {
@@ -73,9 +76,7 @@ export const useAccountStore = create<AccountStoreState>((set, get) => ({
 
       if (!get().networkObserver) {
         const networkObserver = currentWallet.listenNetworkChanges(async () => {
-          set({
-            chainId: await currentWallet.networkInfos().then((n) => n.chainId),
-          });
+          get().setCurrentNetwork();
         });
         set({ networkObserver });
       }
@@ -84,21 +85,13 @@ export const useAccountStore = create<AccountStoreState>((set, get) => ({
         currentWallet
           .connect()
           .then(() => {
-            // get current network
-            currentWallet
-              .networkInfos()
-              .then((infos) => {
-                set({ chainId: infos.chainId });
-              })
-              .catch((error) => {
-                console.warn('error getting network from bearby', error);
-              });
             // subscribe to network events
             const observer = currentWallet.listenAccountChanges(
               (newAddress: string) => {
                 handleBearbyAccountChange(newAddress, get());
               },
             );
+
             set({ currentWallet, accountObserver: observer });
 
             // get connected account
@@ -121,6 +114,8 @@ export const useAccountStore = create<AccountStoreState>((set, get) => ({
 
       set({ currentWallet });
 
+      get().setCurrentNetwork();
+
       currentWallet
         .accounts()
         .then((accounts) => {
@@ -154,4 +149,12 @@ export const useAccountStore = create<AccountStoreState>((set, get) => ({
   setConnectedAccount: async (connectedAccount?: Provider) => {
     set({ connectedAccount });
   },
+
+  setCurrentNetwork: () => {
+    get()
+      .currentWallet?.networkInfos()
+      .then((infos) => {
+        set({ chainId: infos.chainId, network: infos.name });
+      });
+  },
 }));