Add support for setting the storage dynamically

src/store/actions/account/accountActions.ts

@@ -1,4 +1,4 @@
-import { store } from 'store/store';
+import { getStore } from 'store/store';
 import { AccountType } from 'types/account.types';
 import { emptyAccount } from 'store/slices/account/emptyAccount';
 import {
@@ -7,12 +7,12 @@ import {
 } from 'store/slices/account/account.types';
 
 export const setAddress = (address: string) =>
-  store.setState(({ account: state }) => {
+  getStore().setState(({ account: state }) => {
     state.address = address;
   });
 
 export const setAccount = (account: AccountType) =>
-  store.setState(({ account: state }) => {
+  getStore().setState(({ account: state }) => {
     const isSameAddress = state.address === account.address;
     state.accounts = {
       [state.address]: isSameAddress ? account : emptyAccount
@@ -21,7 +21,7 @@ export const setAccount = (account: AccountType) =>
 
 // TODO: check if needed
 export const setLedgerAccount = (ledgerAccount: LedgerAccountType | null) =>
-  store.setState(({ account: state }) => {
+  getStore().setState(({ account: state }) => {
     state.ledgerAccount = ledgerAccount;
   });
 
@@ -33,28 +33,28 @@ export const updateLedgerAccount = ({
   index: LedgerAccountType['index'];
   address: LedgerAccountType['address'];
 }) =>
-  store.setState(({ account: state }) => {
+  getStore().setState(({ account: state }) => {
     if (state.ledgerAccount) {
       state.ledgerAccount.address = address;
       state.ledgerAccount.index = index;
     }
   });
 
 export const setWalletConnectAccount = (walletConnectAccount: string | null) =>
-  store.setState(({ account: state }) => {
+  getStore().setState(({ account: state }) => {
     state.walletConnectAccount = walletConnectAccount;
   });
 
 export const setWebsocketEvent = (message: string) =>
-  store.setState(({ account: state }) => {
+  getStore().setState(({ account: state }) => {
     state.websocketEvent = {
       timestamp: Date.now(),
       message
     };
   });
 
 export const setWebsocketBatchEvent = (data: BatchTransactionsWSResponseType) =>
-  store.setState(({ account: state }) => {
+  getStore().setState(({ account: state }) => {
     state.websocketBatchEvent = {
       timestamp: Date.now(),
       data

src/store/actions/loginInfo/loginInfoActions.ts

@@ -1,52 +1,52 @@
 import { LoginMethodsEnum } from 'types/enums.types';
-import { store } from '../../store';
 import { TokenLoginType } from 'types/login.types';
 import {
   LedgerLoginType,
   LoginInfoType,
   WalletConnectLoginType
 } from 'store/slices/loginInfo/loginInfo.types';
+import { getStore } from 'store/store';
 
 export const setLoginMethod = (loginMethod: LoginMethodsEnum) =>
-  store.setState(({ loginInfo: state }) => {
+  getStore().setState(({ loginInfo: state }) => {
     state.loginMethod = loginMethod;
   });
 
 export const setTokenLogin = (tokenLogin: TokenLoginType) =>
-  store.setState(({ loginInfo: state }) => {
+  getStore().setState(({ loginInfo: state }) => {
     state.tokenLogin = tokenLogin;
   });
 
 export const setTokenLoginSignature = (signature: string) =>
-  store.setState(({ loginInfo: state }) => {
+  getStore().setState(({ loginInfo: state }) => {
     if (state?.tokenLogin != null) {
       state.tokenLogin.signature = signature;
     }
   });
 
 export const setWalletLogin = (walletLogin: LoginInfoType | null) =>
-  store.setState(({ loginInfo: state }) => {
+  getStore().setState(({ loginInfo: state }) => {
     state.walletLogin = walletLogin;
   });
 
 export const setWalletConnectLogin = (
   walletConnectLogin: WalletConnectLoginType | null
 ) =>
-  store.setState(({ loginInfo: state }) => {
+  getStore().setState(({ loginInfo: state }) => {
     state.walletConnectLogin = walletConnectLogin;
   });
 
 export const setLedgerLogin = (ledgerLogin: LedgerLoginType | null) =>
-  store.setState(({ loginInfo: state }) => {
+  getStore().setState(({ loginInfo: state }) => {
     state.ledgerLogin = ledgerLogin;
   });
 
 export const setLogoutRoute = (logoutRoute: string | undefined) =>
-  store.setState(({ loginInfo: state }) => {
+  getStore().setState(({ loginInfo: state }) => {
     state.logoutRoute = logoutRoute;
   });
 
 export const setIsWalletConnectV2Initialized = (isInitialized: boolean) =>
-  store.setState(({ loginInfo: state }) => {
+  getStore().setState(({ loginInfo: state }) => {
     state.isWalletConnectV2Initialized = isInitialized;
   });

src/store/actions/network/networkActions.ts

@@ -1,8 +1,8 @@
 import { NetworkType } from 'types/network.types';
-import { store } from '../../store';
+import { getStore } from '../../store';
 
 export const initializeNetworkConfig = (newNetwork: NetworkType) =>
-  store.setState(({ network: state }) => {
+  getStore().setState(({ network: state }) => {
     const walletConnectV2RelayAddress =
       newNetwork.walletConnectV2RelayAddresses[
         Math.floor(
@@ -19,7 +19,7 @@ export const initializeNetworkConfig = (newNetwork: NetworkType) =>
   });
 
 export const setCustomWalletAddress = (customWalletAddress: string) =>
-  store.setState(({ network: state }) => {
+  getStore().setState(({ network: state }) => {
     state.network.customWalletAddress = customWalletAddress;
   });
 

src/store/actions/sharedActions/sharedActions.ts

@@ -1,16 +1,16 @@
 import { Address } from '@multiversx/sdk-core/out';
-import { store } from '../../store';
+import { getStore } from '../../store';
 import { LoginMethodsEnum } from 'types/enums.types';
 import { resetStore } from 'store/middleware/logoutMiddleware';
 
-export const logoutAction = () => store.setState(resetStore);
+export const logoutAction = () => getStore().setState(resetStore);
 export interface LoginActionPayloadType {
   address: string;
   loginMethod: LoginMethodsEnum;
 }
 
 export const loginAction = ({ address, loginMethod }: LoginActionPayloadType) =>
-  store.setState(({ account, loginInfo }) => {
+  getStore().setState(({ account, loginInfo }) => {
     account.address = address;
     account.publicKey = new Address(address).hex();
     loginInfo.loginMethod = loginMethod;

src/store/selectors/hooks/useSelector.ts

@@ -1,10 +1,11 @@
 import { StoreType } from 'store/store.types';
-import { useStore } from '../../store';
+import { getStoreHook } from '../../store';
 
-type ExtractState<S> = S extends { getState: () => infer X } ? X : StoreType;
+type ExtractState<S> = S extends { getState: () => infer T } ? T : StoreType;
 
-export const useSelector = <T>(
+export function useSelector<T>(
   selector: (state: ExtractState<StoreType>) => T
-) => {
+) {
+  const useStore = getStoreHook();
   return useStore(selector);
-};
+}

src/store/storage/inMemoryStorage.ts

@@ -0,0 +1,32 @@
+interface InMemoryStorageType {
+    [key: string]: string;
+}
+
+export class InMemoryStorage {
+    private storage: InMemoryStorageType = {};
+
+    setItem(key: string, value: string) {
+        this.storage[key] = value;
+    }
+
+    getItem(key: string): string | null {
+        return this.storage.hasOwnProperty(key) ? this.storage[key] : null;
+    }
+
+    removeItem(key: string) {
+        delete this.storage[key];
+    }
+
+    clear() {
+        this.storage = {} as Storage;
+    }
+
+    get length() {
+        return Object.keys(this.storage).length;
+    }
+
+    key(index: number): string | null {
+        const keys = Object.keys(this.storage);
+        return keys[index] || null;
+    }
+}

src/store/storage/index.ts

@@ -0,0 +1,2 @@
+export * from "./storageCallback";
+export * from "./inMemoryStorage";

src/store/storage/storageCallback.ts

@@ -0,0 +1,5 @@
+import { StateStorage } from 'zustand/middleware';
+
+export type StorageCallback = () => StateStorage;
+
+export const defaultStorageCallback: StorageCallback = () => localStorage;

src/store/store.ts

@@ -1,6 +1,11 @@
 import { createStore } from 'zustand/vanilla';
 import { createJSONStorage, devtools, persist } from 'zustand/middleware';
 import { immer } from 'zustand/middleware/immer';
+import {
+  InMemoryStorage,
+  defaultStorageCallback,
+  StorageCallback
+} from './storage';
 import { networkSlice } from './slices/network/networkSlice';
 import { accountSlice } from './slices/account/accountSlice';
 import { createBoundedUseStore } from './createBoundedStore';
@@ -20,24 +25,59 @@ export type MutatorsOut = [
   ['zustand/immer', never]
 ];
 
-export const store = createStore<StoreType, MutatorsOut>(
-  devtools(
-    persist(
-      immer((...args) => ({
-        network: networkSlice(...args),
-        account: accountSlice(...args),
-        loginInfo: loginInfoSlice(...args)
-      })),
-      {
-        name: 'sdk-dapp-store',
-        storage: createJSONStorage(() => localStorage)
-      }
+export const createDAppStore = (getStorageCallback: StorageCallback) => {
+  const store = createStore<StoreType, MutatorsOut>(
+    devtools(
+      persist(
+        immer((...args) => ({
+          network: networkSlice(...args),
+          account: accountSlice(...args),
+          loginInfo: loginInfoSlice(...args)
+        })),
+        {
+          name: 'sdk-dapp-store',
+          storage: createJSONStorage(getStorageCallback)
+        }
+      )
     )
-  )
-);
+  );
+  applyMiddleware(store);
 
-applyMiddleware(store);
+  return store;
+};
 
-export const getState = () => store.getState();
+export type StoreApi = ReturnType<typeof createDAppStore>;
 
-export const useStore = createBoundedUseStore(store);
+let store: StoreApi;
+
+export const getStore = () => {
+  if (!store) {
+    setDAppStore(createDAppStore(() => new InMemoryStorage()));
+  }
+  return store;
+};
+
+export const setDAppStore = (_store: StoreApi) => {
+  store = _store;
+};
+
+/**
+ * Initialize store with the preferred storage by passing a callback.
+ * Default storage is localStorage.
+ * You can pass your own storage.
+ * Call this function before using store, ideally before app bootstrapping.
+ * @param getStorageCallback
+ * @default () => localStorage
+ * @returns persistent store instance
+ * @example
+ * initStore(() => window.localStorage);
+ * initStore(() => window.sessionStorage);
+ * initStore(() => new InMemoryStorage());
+ * */
+export const initStore = (getStorageCallback = defaultStorageCallback) => {
+  return setDAppStore(createDAppStore(getStorageCallback));
+};
+
+export const getState = () => getStore().getState();
+
+export const getStoreHook = () => createBoundedUseStore(getStore());