fix: useStore selector was not working on derived stores

schummar · Dec 10, 2024 · 300adfd · 300adfd
1 parent f1ae38c
commit 300adfd
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 1 deletion.
diff --git a/src/react/useStore.ts b/src/react/useStore.ts
@@ -38,6 +38,7 @@ export function useStore<T, S>(store: Store<T>, argument1?: any, argument2?: any
   const selector = makeSelector<T, S>(
     typeof argument1 === 'function' || typeof argument1 === 'string' ? argument1 : undefined,
   );
+
   const allOptions = (
     typeof argument1 === 'object' ? argument1 : (argument2 ?? {})
   ) as UseStoreOptions<S>;
@@ -51,7 +52,7 @@ export function useStore<T, S>(store: Store<T>, argument1?: any, argument2?: any
         for (const selector of store.derivedFrom!.selectors) {
           value = makeSelector(selector)(value);
         }
-        return value;
+        return selector(value);
       },
       allOptions,
     );

diff --git a/test/react/registerReactMethods.test.tsx b/test/react/registerReactMethods.test.tsx
@@ -20,6 +20,36 @@ describe('register react methods', () => {
     expect(div.textContent).toBe('0');
   });
 
+  test('useStore with selector', async () => {
+    const store = createStore({ x: 0 });
+
+    const Component = vi.fn(function Component() {
+      const value = store.useStore((x) => x.x + 1);
+
+      return <div data-testid="div">{value}</div>;
+    });
+
+    render(<Component />);
+    const div = screen.getByTestId('div');
+
+    expect(div.textContent).toBe('1');
+  });
+
+  test('useStore with selector and mapping', async () => {
+    const store = createStore({ x: 0 });
+
+    const Component = vi.fn(function Component() {
+      const value = store.map('x').useStore((x) => x + 1);
+
+      return <div data-testid="div">{value}</div>;
+    });
+
+    render(<Component />);
+    const div = screen.getByTestId('div');
+
+    expect(div.textContent).toBe('1');
+  });
+
   test('useProp', async () => {
     const store = createStore({ x: 0 });
 

diff --git a/test/react/useStore.test.tsx b/test/react/useStore.test.tsx
@@ -124,6 +124,49 @@ describe('useStore', () => {
     expect(Component.mock.calls.length).toBe(2);
   });
 
+  test('with selector as argument', async () => {
+    const store = createStore({ a: 1, b: 2 });
+
+    const Component = vi.fn(function Component() {
+      const value = useStore(store, (s) => s.a);
+
+      return <div data-testid="div">{value}</div>;
+    });
+
+    render(<Component />);
+    const div = screen.getByTestId('div');
+
+    act(() => {
+      store.set({ a: 2, b: 2 });
+    });
+
+    expect(div.textContent).toBe('2');
+    expect(Component.mock.calls.length).toBe(2);
+  });
+
+  test('with mapping and selector as argument', async () => {
+    const store = createStore({ a: 1, b: 2 });
+
+    const Component = vi.fn(function Component() {
+      const value = useStore(
+        store.map((s) => s.a),
+        (s) => s + 1,
+      );
+
+      return <div data-testid="div">{value}</div>;
+    });
+
+    render(<Component />);
+    const div = screen.getByTestId('div');
+
+    act(() => {
+      store.set({ a: 2, b: 2 });
+    });
+
+    expect(div.textContent).toBe('3');
+    expect(Component.mock.calls.length).toBe(2);
+  });
+
   test('fall back to store equals', async () => {
     const store = createStore({ a: 1 }, { equals: strictEqual });