Fix tests and update onPersist and onRehydrate functions

Author: tommasoberlose

src/core/reducer.ts

@@ -18,26 +18,51 @@ import { deepGetByPath, REHYDRATE, writePersistedStorage } from './utils';
 
 /**
  * Creates a persisted reducer that wraps Redux Toolkit's `createReducer`.
- * This function enhances a standard reducer with automatic state persistence,
- * saving its state to storage and rehydrating it on app startup.
  *
- * This function must be used with a store configured by `configurePersistedStore`.
+ * This function enhances a standard reducer with automatic state persistence,
+ * saving its state to storage and rehydrating it when the application starts.
+ * It must be used with a store configured by `configurePersistedStore`.
  *
- * @param reducerName - A unique string name for the reducer. This name serves as the key in both the root state and the storage.
- * @param initialState - The initial state for the reducer.
- * @param mapOrBuilderCallback - A callback that receives a `builder` object to define case reducers, similar to the original `createReducer`.
- * @param nestedPath - An optional dot-separated string path indicating where the reducer's state is located within the root state. If not provided, `reducerName` is used.
  * @public
+ * @param reducerName - A unique name for the reducer. This name is used as the key
+ * in both the root state and the underlying storage.
+ * @param initialState - The initial state for the reducer, same as in `createReducer`.
+ * @param mapOrBuilderCallback - A callback that receives a `builder` object to define
+ * case reducers. This builder is wrapped to automatically track state changes for persistence.
+ * @param persistenceOptions - Optional configuration for customizing persistence behavior.
+ * @returns A reducer function enhanced with persistence capabilities.
+ *
+ * @example
+ * ```typescript
+ * const counterReducer = createPersistedReducer(
+ * 'counter',
+ * { value: 0 },
+ * (builder) => {
+ * builder.addCase(increment, (state) => {
+ * state.value++;
+ * });
+ * },
+ * {
+ * // Optional: Specify a different path in the root state.
+ * nestedPath: 'nested.counter',
+ * // Optional: Transform state before saving.
+ * onPersist: (state) => ({ value: state.value.toString() }),
+ * // Optional: Transform state after loading from storage.
+ * onRehydrate: (persistedState) => ({ value: parseInt(persistedState.value, 10) }),
+ * }
+ * );
+ * ```
  */
 export const createPersistedReducer = <
   ReducerName extends string,
   S extends NotFunction<any>,
+  SavedState,
   Nesting extends NestedPath<ReducerName> = ReducerName,
 >(
   reducerName: ReducerName,
   initialState: S | (() => S),
   mapOrBuilderCallback: (builder: ActionReducerMapBuilder<S>) => void,
-  persistenceOptions?: ReducerPersistenceOptions<ReducerName, S, Nesting>,
+  persistenceOptions?: ReducerPersistenceOptions<ReducerName, S, SavedState, Nesting>,
 ): PersistedReducer<S, ReducerName, Nesting> => {
   // Register the reducer for persistence tracking.
   Settings.subscribeSlice(reducerName);
@@ -51,6 +76,7 @@ export const createPersistedReducer = <
   /**
    * The full dot-separated path to the reducer's state within the root state object.
    * Defaults to the reducer's name if `nestedPath` is not provided.
+   * @internal
    */
   const finalNestedPath = (persistenceOptions?.nestedPath ?? reducerName) as Nesting;
 
@@ -64,10 +90,16 @@ export const createPersistedReducer = <
     if (debounceTimeout) clearTimeout(debounceTimeout);
     debounceTimeout = setTimeout(() => {
       const reducerState = deepGetByPath(state, finalNestedPath);
-      if (persistenceOptions && 'onDump' in persistenceOptions) {
-        writePersistedStorage(persistenceOptions.onDump(reducerState), reducerName);
+      if (reducerState === null) {
+        if (process.env.NODE_ENV !== 'production') {
+          console.error(`DUMP: No state found for ${reducerName}, check if the nestedPath is corrected.`);
+        }
       } else {
-        writePersistedStorage(reducerState, reducerName);
+        if (persistenceOptions && 'onPersist' in persistenceOptions) {
+          writePersistedStorage(persistenceOptions.onPersist(reducerState), reducerName);
+        } else {
+          writePersistedStorage(reducerState, reducerName);
+        }
       }
     }, 100);
   };
@@ -87,14 +119,14 @@ export const createPersistedReducer = <
     // Add a case to handle the rehydration of state from storage.
     builder.addCase(
       REHYDRATE.toString(),
-      (_state, action: PayloadAction<RehydrateActionPayload<ReducerName, S>>):
+      (_state, action: PayloadAction<RehydrateActionPayload<ReducerName, S | SavedState>>):
         | void
         | S => {
         if (action.payload?.[reducerName]) {
           if (persistenceOptions && 'onRehydrate' in persistenceOptions) {
-            return persistenceOptions.onRehydrate(action.payload[reducerName]);
+            return persistenceOptions.onRehydrate(action.payload[reducerName] as SavedState);
           } else {
-            return action.payload[reducerName];
+            return action.payload[reducerName] as S;
           }
         }
       },

src/core/slice.ts

@@ -8,7 +8,12 @@ import {
 import { Builder } from './extraReducersBuilder';
 import { listenerMiddleware } from './middleware';
 import Settings from './settings';
-import { NestedPath, PersistedSlice, RehydrateActionPayload, SlicePersistenceOptions } from './types';
+import {
+  NestedPath,
+  PersistedSlice,
+  RehydrateActionPayload,
+  SlicePersistenceOptions,
+} from './types';
 import UpdatedAtHelper from './updatedAtHelper';
 import { deepGetByPath, REHYDRATE, writePersistedStorage } from './utils';
 
@@ -19,17 +24,48 @@ import { deepGetByPath, REHYDRATE, writePersistedStorage } from './utils';
  *
  * This function must be used with a store configured by `configurePersistedStore`.
  *
+ * @public
  * @param sliceOptions - The standard `CreateSliceOptions` object from Redux Toolkit.
- * @param nestedPath - An optional dot-separated string path indicating where the
- * slice's state is located within the root state. If not provided, `reducerPath`
- * or `name` from `sliceOptions` is used.
+ * @param persistenceOptions - Optional configuration for persistence behavior.
+ * @param persistenceOptions.nestedPath - A dot-separated string path indicating where the
+ * slice's state is located within the root state. If not provided, it defaults to
+ * `reducerPath` or `name` from `sliceOptions`.
+ * @param persistenceOptions.onPersist - A function that transforms the slice's state
+ * *before* it is saved to storage. This is useful for saving a different version of the state
+ * than what is used in the application.
+ * @param persistenceOptions.onRehydrate - A function that transforms the state *after* it is
+ * loaded from storage but *before* it is placed in the Redux store. This is useful for
+ * migrating old state shapes or re-instantiating complex objects.
  * @returns A Redux slice object with persistence enabled, augmented with a
- * `nestedPath` property for state tracking.
+ * `nestedPath` property for internal state tracking.
  *
- * @public
+ * @example
+ * // Basic usage
+ * const counterSlice = createPersistedSlice({
+ * name: 'counter',
+ * initialState: { value: 0 },
+ * reducers: {
+ * increment: (state) => {
+ * state.value += 1;
+ * },
+ * },
+ * });
+ *
+ * // Usage with persistence options
+ * const userSlice = createPersistedSlice({
+ * name: 'user',
+ * initialState: { data: null, loadedAt: null },
+ * reducers: {
+ * // ...
+ * },
+ * }, {
+ * onRehydrate: (savedState) => ({ ...savedState, loadedAt: new Date() }),
+ * onPersist: (state) => ({ data: state.data }), // Only persist the 'data' field
+ * });
  */
 export const createPersistedSlice = <
   SliceState,
+  SavedState,
   Name extends string,
   PCR extends SliceCaseReducers<SliceState>,
   ReducerPath extends string = Name,
@@ -43,7 +79,13 @@ export const createPersistedSlice = <
     ReducerPath,
     PersistedSelectors
   >,
-  persistenceOptions?: SlicePersistenceOptions<SliceState, Name, ReducerPath, Nesting>,
+  persistenceOptions?: SlicePersistenceOptions<
+    SliceState,
+    SavedState,
+    Name,
+    ReducerPath,
+    Nesting
+  >,
 ): PersistedSlice<
   SliceState,
   PCR,
@@ -63,6 +105,7 @@ export const createPersistedSlice = <
 
   /**
    * The full dot-separated path to the slice's state within the root state object.
+   * @internal
    */
   const finalNestedPath = (persistenceOptions?.nestedPath ??
     sliceOptions.reducerPath ??
@@ -78,10 +121,19 @@ export const createPersistedSlice = <
     if (debounceTimeout) clearTimeout(debounceTimeout);
     debounceTimeout = setTimeout(() => {
       const sliceState = deepGetByPath(state, finalNestedPath);
-      if (persistenceOptions && 'onDump' in persistenceOptions) {
-        writePersistedStorage(persistenceOptions.onDump(sliceState), sliceOptions.name);
+      if (sliceState === null) {
+        if (process.env.NODE_ENV !== 'production') {
+          console.error(`DUMP: No state found for ${sliceOptions.name}, check if the nestedPath is corrected.`);
+        }
       } else {
-        writePersistedStorage(sliceState, sliceOptions.name);
+        if (persistenceOptions && 'onPersist' in persistenceOptions) {
+          writePersistedStorage(
+            persistenceOptions.onPersist(sliceState),
+            sliceOptions.name,
+          );
+        } else {
+          writePersistedStorage(sliceState, sliceOptions.name);
+        }
       }
     }, 100);
   };
@@ -105,13 +157,17 @@ export const createPersistedSlice = <
         REHYDRATE.toString(),
         (
           _state,
-          action: PayloadAction<RehydrateActionPayload<Name, SliceState>>,
+          action: PayloadAction<
+            RehydrateActionPayload<Name, SliceState | SavedState>
+          >,
         ): void | SliceState => {
           if (action.payload?.[sliceOptions.name]) {
             if (persistenceOptions && 'onRehydrate' in persistenceOptions) {
-              return persistenceOptions.onRehydrate(action.payload[sliceOptions.name]);
+              return persistenceOptions.onRehydrate(
+                action.payload[sliceOptions.name] as SavedState,
+              );
             } else {
-              return action.payload[sliceOptions.name];
+              return action.payload[sliceOptions.name] as SliceState;
             }
           }
         },