Merge pull request #120 from udecode/fix/function-options-handling

Author: zbeyens

.changeset/fix-function-options.md

@@ -0,0 +1,5 @@
+---
+'zustand-x': patch
+---
+
+Fix assigning function-valued options so callbacks still work.

packages/zustand-x/src/internal/createBaseApi.ts

@@ -1,8 +1,7 @@
 import { StoreMutatorIdentifier } from 'zustand';
 
 import { TBaseStateApiForBuilder, TStoreApiGet } from '../types/baseStore';
-import { TState } from '../types/utils';
-import { TCreatedStoreMutateType } from '../types/utils';
+import { TCreatedStoreMutateType, TState } from '../types/utils';
 
 import type {
   AnyFunction,
@@ -44,7 +43,13 @@ export const createBaseApi = <
     const typedKey = key as keyof StateType;
     const prevValue = store.getState()[typedKey];
 
-    if (typeof value === 'function') {
+    const shouldInvokeUpdater =
+      typeof value === 'function' &&
+      prevValue !== undefined &&
+      prevValue !== null &&
+      typeof prevValue !== 'function';
+
+    if (shouldInvokeUpdater) {
       value = value(prevValue);
     }
     if (prevValue === value) return;

packages/zustand-x/src/tests/createStore.spec.ts

@@ -126,6 +126,43 @@ describe('zustandX', () => {
       expect(starsListener).toHaveBeenCalledTimes(2);
       expect(starsListener).toHaveBeenLastCalledWith(3, 0);
     });
+
+    it('should support callback updates for primitive values', () => {
+      expect(store.get('stars')).toBe(0);
+      store.set('stars', (currentStars) => currentStars + 5);
+      expect(store.get('stars')).toBe(5);
+      store.set('state', {
+        name: 'zustandX',
+        stars: 0,
+      });
+    });
+
+    it('should store function values without invoking them', () => {
+      type FunctionStoreState = {
+        handler: () => string;
+        optionalHandler: ((value: number) => number) | null;
+      };
+
+      const functionStore = createStore<FunctionStoreState>(
+        {
+          handler: () => 'initial',
+          optionalHandler: null,
+        },
+        {
+          name: 'function-store',
+        }
+      );
+
+      const newHandler = vi.fn(() => 'updated');
+      expect(() => functionStore.set('handler', newHandler)).not.toThrow();
+      expect(functionStore.get('handler')).toBe(newHandler);
+      expect(newHandler).not.toHaveBeenCalled();
+
+      const nextOptional = vi.fn((value: number) => value * 2);
+      functionStore.set('optionalHandler', nextOptional);
+      expect(functionStore.get('optionalHandler')).toBe(nextOptional);
+      expect(nextOptional).not.toHaveBeenCalled();
+    });
   });
 
   describe('when using hooks', () => {