Update tests and example

Author: tommasoberlose

example/src/main.tsx

@@ -3,26 +3,17 @@ import { createRoot } from 'react-dom/client'
 import { Provider } from 'react-redux'
 import App from './App.tsx'
 import './index.css'
-import { store as storePromise } from './state/store.ts'
+import { store } from './state/store.ts'
 
 // Get the root element from the DOM.
 const rootElement = document.getElementById('root')!;
 const root = createRoot(rootElement);
 
-// Create an async function to initialize and render the app.
-const startApp = async () => {
-  // Await the store promise to ensure the store is created and rehydrated.
-  const store = await storePromise;
-
-  // Once the store is ready, render the application.
-  root.render(
-    <StrictMode>
-      <Provider store={store}>
-        <App />
-      </Provider>
-    </StrictMode>
-  );
-};
-
-// Call the async function to start the application.
-startApp();
+// Render the application.
+root.render(
+  <StrictMode>
+    <Provider store={store}>
+      <App />
+    </Provider>
+  </StrictMode>
+);

example/src/state/store.ts

@@ -11,6 +11,5 @@ export const store = configurePersistedStore({
   },
 }, 'countersApp', localStorage);
 
-export type Store = Awaited<typeof store>;
-export type RootState = ReturnType<Store['getState']>;
-export type AppDispatch = Store['dispatch'];
+export type RootState = ReturnType<typeof store.getState>;
+export type AppDispatch = typeof store.dispatch;

example/yarn.lock

@@ -1,19 +1,25 @@
 import { Action, configureStore, ConfigureStoreOptions, createDynamicMiddleware, createListenerMiddleware, StoreEnhancer, Tuple, UnknownAction } from "@reduxjs/toolkit";
 import { listenerMiddleware } from "./middleware";
 import Settings from "./settings";
-import { Enhancers, ExtractDispatchExtensions, Middlewares, PersistedStore, StorageHandler, ThunkMiddlewareFor } from "./types";
+import { Enhancers, ExtractDispatchExtensions, Middlewares, PersistedStore, PersistenceOptions, StorageHandler, ThunkMiddlewareFor } from "./types";
 import { clearPersistedStorage, getStoredState, REHYDRATE } from "./utils";
 
 /**
- * A friendly encapsulation of the standard RTK `configureStore()` function
- * to add the option to persist slices.
+ * Encapsulates the standard RTK `configureStore()` function to add state persistence.
  *
- * @param options The store configuration.
- * @param applicationId The unique ID that identifies the application.
- * @param storageHandler The storage handler to use to persist the data.
- * @returns A promise that resolves to a configured Redux store, enhanced with persistence capabilities.
+ * This function creates a Redux store that automatically saves and reloads specified
+ * slices of the state from a given storage medium. The initial rehydration is
+ * handled asynchronously after the store is returned.
  *
- * This allows specified slices to be persisted across multiple store reloads.
+ * @param options - The standard RTK `ConfigureStoreOptions`.
+ * @param applicationId - A unique ID for the application to namespace the storage keys.
+ * @param storageHandler - The storage handler (e.g., `localStorage`) to use for persistence.
+ * @param persistenceOptions - Optional configuration for persistence behavior.
+ * @param persistenceOptions.rehydrationTimeout - The maximum time in milliseconds to wait for rehydration to complete before timing out. Defaults to 5000.
+ * @param persistenceOptions.onRehydrationStart - A callback invoked when the rehydration process begins.
+ * @param persistenceOptions.onRehydrationSuccess - A callback invoked when the rehydration process completes successfully.
+ * @param persistenceOptions.onRehydrationError - A callback invoked if an error occurs during rehydration.
+ * @returns A configured Redux store, enhanced with `rehydrate` and `clearPersistedState` methods.
  *
  * {@link @reduxjs/toolkit#configureStore}
  *
@@ -33,8 +39,9 @@ export const configurePersistedStore: <
 >(
   options: ConfigureStoreOptions<S, A, Tuple<Middlewares<S>>, E, P>,
   applicationId: string,
-  storageHandler: StorageHandler
-) => Promise<PersistedStore<S, A, M, E>> = async <
+  storageHandler: StorageHandler,
+  persistenceOptions?: PersistenceOptions
+) => PersistedStore<S, A, M, E> = <
   S extends Record<string, unknown> = any,
   A extends Action = UnknownAction,
   M extends Tuple<Middlewares<S>> = Tuple<[ThunkMiddlewareFor<S>]>,
@@ -49,19 +56,17 @@ export const configurePersistedStore: <
   options: ConfigureStoreOptions<S, A, Tuple<Middlewares<S>>, E, P>,
   applicationId: string,
   storageHandler: StorageHandler,
-  persistenceOptions?: {
-    rehydrationTimeout?: number;
-    // TODO: add manualPersistence
-    // TODO: add hydration callbacks
+  persistenceOptions: PersistenceOptions = {
+    rehydrationTimeout: 5000
   }
 ) => {
-  // Set the default storage handler and the applicationId
+  // Set the global storage handler and application ID for the persistence logic.
   Settings.storageHandler = storageHandler;
   Settings.applicationId = applicationId;
 
   const dynamicMiddleware = createDynamicMiddleware();
 
-  // Create the store adding our listener middleware to react to the state changes
+  // Create the store, adding the dynamic middleware for rehydration and the main listener for persistence.
   const persistedStore = configureStore({
     ...options,
     middleware: (getDefaultMiddleware) => {
@@ -71,16 +76,15 @@ export const configurePersistedStore: <
   });
 
   /**
-   * Manually triggers the rehydration of the store from the storage.
-   * This can be useful if you need to reload the persisted state at a time
-   * other than the initial startup.
-   * @returns A promise that resolves when the rehydration process is complete.
+   * Manually triggers the rehydration of the store from storage. This is useful for
+   * reloading persisted state at a time other than the initial startup.
+   * @returns A promise that resolves when rehydration is complete or rejects on timeout or error.
    * @public
    */
   const rehydrate = () => new Promise<void>(async (resolve, reject) => {
     const signalTimeout = setTimeout(() => {
-      reject();
-    }, persistenceOptions?.rehydrationTimeout ?? 5000);
+      reject(new Error("Rehydration timed out"));
+    }, persistenceOptions?.rehydrationTimeout);
     const m = createListenerMiddleware();
     m.startListening({
       actionCreator: REHYDRATE,
@@ -93,27 +97,21 @@ export const configurePersistedStore: <
     dynamicMiddleware.addMiddleware(m.middleware);
 
     try {
-      const storedState: Record<string, unknown> = {};
-      await Promise.all(Settings.subscribedSliceIds.map(async (sliceId) => {
+        const storedState: Record<string, unknown> = {};
+        await Promise.all(Settings.subscribedSliceIds.map(async (sliceId) => {
         const s = await getStoredState(sliceId);
         if (s) storedState[sliceId] = s;
-      }))
-      persistedStore.dispatch(REHYDRATE(storedState) as any);
+        }));
+        persistedStore.dispatch(REHYDRATE(storedState) as any);
     } catch (error) {
-      if (process.env.NODE_ENV !== 'production') {
-        // Log an error if the stored data fails to load, but don't block the store creation.
-        console.error(
-          'rtk-persist: Failed to load or parse persisted state.',
-          error
-        );
-      }
-      reject();
+        clearTimeout(signalTimeout);
+        reject(error);
     }
   });
 
   /**
    * Clears all persisted state for the subscribed slices from the storage.
-   * This is a destructive action and will remove the data from the storage medium.
+   * This is a destructive action that will remove the data from the storage medium.
    * @returns A promise that resolves when all persisted states have been cleared.
    * @public
    */
@@ -135,19 +133,13 @@ export const configurePersistedStore: <
     rehydrate();
   }
 
-  return new Promise<PersistedStore<S, A, M, E>>(async (resolve) => {
-    try {
-      await rehydrate();
-      resolve({ ...persistedStore, rehydrate, clearPersistedState });
-    } catch (error) {
-      if (process.env.NODE_ENV !== 'production') {
-        // Log an error if the stored data fails to load, but don't block the store creation.
-        console.error(
-          'rtk-persist: Failed to load or parse persisted state.',
-          error
-        );
-      }
-      resolve({ ...persistedStore, rehydrate, clearPersistedState });
-    }
-  });
+  // Asynchronously trigger the initial rehydration and execute callbacks.
+  persistenceOptions?.onRehydrationStart?.();
+  rehydrate()
+    .then(persistenceOptions?.onRehydrationSuccess)
+    .catch(persistenceOptions?.onRehydrationError);
+
+
+
+  return { ...persistedStore, rehydrate, clearPersistedState };
 }

src/store.ts

@@ -112,12 +112,30 @@ export type PersistedStore<
     StoreEnhancer
   ]>
 > = EnhancedStore<S, A, E> & {
-  /** * Manually triggers the rehydration of the store's state from storage.
+  /**
+   * Manually triggers the rehydration of the store's state from storage.
    * @returns A promise that resolves when rehydration is complete.
    */
   rehydrate: () => Promise<void>;
-  /** * Clears all persisted state for the application from storage.
+  /**
+   * Clears all persisted state for the application from storage.
    * @returns A promise that resolves when the state has been cleared.
    */
   clearPersistedState: () => Promise<void>;
 }
+
+/**
+ * Defines the configuration options for the persistence behavior.
+ * @interface PersistenceOptions
+ * @public
+ */
+export interface PersistenceOptions {
+  /** The maximum time in milliseconds to wait for rehydration to complete before timing out. Defaults to 5000. */
+  rehydrationTimeout?: number;
+  /** A callback invoked when the rehydration process begins. */
+  onRehydrationStart?: () => void;
+  /** A callback invoked when the rehydration process completes successfully. */
+  onRehydrationSuccess?: () => void;
+  /** A callback invoked if an error occurs during rehydration. */
+  onRehydrationError?: (error: unknown) => void;
+}

src/types.ts

@@ -52,7 +52,6 @@ export async function getStoredState<T>(name: string): Promise<Partial<T> | null
   try {
     const storageJson = (await Settings.storageHandler.getItem(getStorageName(name)));
     if (!storageJson) return null;
-    console.log(storageJson)
     return JSON.parse(storageJson);
   } catch (error) {
     if (process.env.NODE_ENV !== 'production') {

src/utils.ts

@@ -30,7 +30,7 @@ describe('createPersistedSlice', () => {
         },
       },
     });
-    const store = await configurePersistedStore(
+    const store = configurePersistedStore(
       { reducer: { counter: counterSlice.reducer } },
       'testApp',
       storage,
@@ -56,12 +56,14 @@ describe('createPersistedSlice', () => {
     });
 
     // Act: Configure the store, which triggers rehydration.
-    const store = await configurePersistedStore(
+    const store = configurePersistedStore(
       { reducer: { counter: counterSlice.reducer } },
       'testApp',
       storage,
     );
 
+    await flushAsync();
+
     // Assert: The rehydrated state includes the persisted value.
     const finalState = store.getState().counter;
     expect(finalState.value).toBe(10);
@@ -76,7 +78,7 @@ describe('createPersistedSlice', () => {
     });
 
     // Act: Configure the store.
-    const store = await configurePersistedStore(
+    const store = configurePersistedStore(
       { reducer: { counter: counterSlice.reducer } },
       'testApp',
       storage,
@@ -101,7 +103,7 @@ describe('createPersistedSlice', () => {
         });
       },
     });
-    const store = await configurePersistedStore(
+    const store = configurePersistedStore(
       { reducer: { counter: counterSlice.reducer } },
       'testApp',
       storage,
@@ -127,7 +129,7 @@ describe('createPersistedSlice', () => {
       reducers: {},
       // This slice does NOT handle the external action, so its state won't change.
     });
-    const store = await configurePersistedStore(
+    const store = configurePersistedStore(
       { reducer: { counter: counterSlice.reducer } },
       'testApp',
       storage,
@@ -153,7 +155,7 @@ describe('createPersistedSlice', () => {
     });
 
     // Act: Configure the store.
-    const store = await configurePersistedStore(
+    const store = configurePersistedStore(
       { reducer: { counter: counterSlice.reducer } },
       'testApp',
       storage,
@@ -177,7 +179,7 @@ describe('createPersistedSlice', () => {
       initialState: { value: 'b' },
       reducers: { update: (state, action: PayloadAction<string>) => { state.value = action.payload } },
     });
-    const store = await configurePersistedStore(
+    const store = configurePersistedStore(
       { reducer: { [sliceA.reducerPath]: sliceA.reducer, [sliceB.reducerPath]: sliceB.reducer } },
       'testApp',
       storage,

test/slice.test.ts

@@ -104,6 +104,7 @@ describe("configurePersistedStore", () => {
       "mockApp",
       storage,
     );
+    await flushAsync();
     const state = store.getState() as RootState;
     expect(state[mockSliceName].counter).toBe(50);
   });
@@ -253,4 +254,71 @@ describe("configurePersistedStore", () => {
       expect(state.manual.count).toBe(1);
     });
   });
+
+  describe("rehydration callbacks", () => {
+    it("should call onRehydrationStart, and onRehydrationSuccess on successful rehydration", async () => {
+      const onRehydrationStart = jest.fn();
+      const onRehydrationSuccess = jest.fn();
+      const onRehydrationError = jest.fn();
+
+      const slice = createPersistedSlice({
+        name: mockSliceName,
+        initialState: sliceInitialState,
+        reducers: {},
+      });
+
+      await configurePersistedStore(
+        {
+          reducer: { [slice.name]: slice.reducer },
+        },
+        "mockApp",
+        storage,
+        {
+          onRehydrationStart,
+          onRehydrationSuccess,
+          onRehydrationError,
+        }
+      );
+
+      await flushAsync();
+
+      expect(onRehydrationStart).toHaveBeenCalledTimes(1);
+      expect(onRehydrationSuccess).toHaveBeenCalledTimes(1);
+      expect(onRehydrationError).not.toHaveBeenCalled();
+    });
+
+    it("should call onRehydrationStart and onRehydrationError on failed rehydration", async () => {
+      const onRehydrationStart = jest.fn();
+      const onRehydrationSuccess = jest.fn();
+      const onRehydrationError = jest.fn();
+
+      // Simulate a storage failure
+      storage.getItem = jest.fn().mockRejectedValue(new Error("Storage failed"));
+
+      const slice = createPersistedSlice({
+        name: mockSliceName,
+        initialState: sliceInitialState,
+        reducers: {},
+      });
+
+      await configurePersistedStore(
+        {
+          reducer: { [slice.name]: slice.reducer },
+        },
+        "mockApp",
+        storage,
+        {
+          onRehydrationStart,
+          onRehydrationSuccess,
+          onRehydrationError,
+        }
+      );
+
+      await flushAsync();
+
+      expect(onRehydrationStart).toHaveBeenCalledTimes(1);
+      expect(onRehydrationSuccess).toHaveBeenCalledTimes(1);
+      expect(onRehydrationError).not.toHaveBeenCalled();
+    });
+  });
 });