From ffaa688671e717183b7112b88190fb129b0f388f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?D=C6=B0=C6=A1ng=20T=E1=BA=A5n=20Hu=E1=BB=B3nh=20Phong?=
 <alphanolucifer@gmail.com>
Date: Sat, 19 Jun 2021 06:34:51 +0700
Subject: [PATCH] make type NamedSet compatible with type SetState

---
 package.json                   |   1 +
 src/middleware.ts              |   9 +-
 tests/middlewareTypes.test.tsx | 201 +++++++++++++++++++++++++++++++++
 yarn.lock                      |   5 +
 4 files changed, 214 insertions(+), 2 deletions(-)
 create mode 100644 tests/middlewareTypes.test.tsx

diff --git a/package.json b/package.json
index e870294657..54cc6269c3 100644
--- a/package.json
+++ b/package.json
@@ -142,6 +142,7 @@
     "eslint-plugin-react": "^7.24.0",
     "eslint-plugin-react-hooks": "^4.2.0",
     "husky": "^6.0.0",
+    "immer": "^9.0.3",
     "jest": "^27.0.4",
     "json": "^11.0.0",
     "lint-staged": "^11.0.0",
diff --git a/src/middleware.ts b/src/middleware.ts
index 021ffb1fc2..8a60f1d89e 100644
--- a/src/middleware.ts
+++ b/src/middleware.ts
@@ -31,8 +31,13 @@ export const redux =
   }
 
 export type NamedSet<T extends State> = {
-  <K extends keyof T>(
-    partial: PartialState<T, K>,
+  <
+    K1 extends keyof T,
+    K2 extends keyof T = K1,
+    K3 extends keyof T = K2,
+    K4 extends keyof T = K3
+  >(
+    partial: PartialState<T, K1, K2, K3, K4>,
     replace?: boolean,
     name?: string
   ): void
diff --git a/tests/middlewareTypes.test.tsx b/tests/middlewareTypes.test.tsx
new file mode 100644
index 0000000000..2834082ca1
--- /dev/null
+++ b/tests/middlewareTypes.test.tsx
@@ -0,0 +1,201 @@
+import { produce } from 'immer'
+import type { Draft } from 'immer'
+import create, { UseStore } from '../src'
+import { State, StateCreator } from '../src/vanilla'
+import { devtools, NamedSet, persist } from '../src/middleware'
+
+type TImmerConfigFn<T extends State> = (fn: (draft: Draft<T>) => void) => void
+type TImmerConfig<T extends State> = StateCreator<T, TImmerConfigFn<T>>
+
+interface ISelectors<T> {
+  use: {
+    [key in keyof T]: () => T[key]
+  }
+}
+
+const immer = <T extends State>(
+  config: TImmerConfig<T>
+): StateCreator<T, NamedSet<T>> => {
+  return (set, get, api) => {
+    return config((fn) => set(produce<T>(fn)), get, api)
+  }
+}
+
+const createSelectorHooks = <T extends State>(store: UseStore<T>) => {
+  const storeAsSelectors = store as unknown as ISelectors<T>
+  storeAsSelectors.use = {} as ISelectors<T>['use']
+
+  Object.keys(store.getState()).forEach((key) => {
+    const storeKey = key as keyof T
+    const selector = (state: T) => state[storeKey]
+
+    storeAsSelectors.use[storeKey] = () => store(selector)
+  })
+
+  return store as UseStore<T> & ISelectors<T>
+}
+
+interface ITestStateProps {
+  testKey: string
+  setTestKey: (testKey: string) => void
+}
+
+it('should have correct type when creating store with devtool', () => {
+  const createStoreWithDevtool = <T extends State>(
+    createState: StateCreator<T>,
+    prefix = 'prefix'
+  ): UseStore<T> & ISelectors<T> => {
+    return createSelectorHooks(create(devtools(createState, prefix)))
+  }
+
+  const testDevtoolStore = createStoreWithDevtool<ITestStateProps>(
+    (set) => ({
+      testKey: 'test',
+      setTestKey: (testKey: string) => {
+        set((state) => {
+          state.testKey = testKey
+        })
+      },
+    }),
+    'test'
+  )
+
+  const TestComponent = (): JSX.Element => {
+    testDevtoolStore.use.testKey()
+    testDevtoolStore.use.setTestKey()
+
+    return <></>
+  }
+})
+
+it('should have correct type when creating store with devtool and immer', () => {
+  const createStoreWithImmer = <T extends State>(
+    createState: TImmerConfig<T>,
+    prefix = 'prefix'
+  ): UseStore<T> & ISelectors<T> => {
+    return createSelectorHooks(create(devtools(immer(createState), prefix)))
+  }
+
+  const testImmerStore = createStoreWithImmer<ITestStateProps>(
+    (set) => ({
+      testKey: 'test',
+      setTestKey: (testKey: string) => {
+        set((state) => {
+          state.testKey = testKey
+        })
+      },
+    }),
+    'test'
+  )
+
+  const TestComponent = (): JSX.Element => {
+    testImmerStore.use.testKey()
+    testImmerStore.use.setTestKey()
+
+    return <></>
+  }
+})
+
+it('should have correct type when creating store with devtool and persist', () => {
+  const createStoreWithPersist = <T extends State>(
+    createState: StateCreator<T>,
+    prefix = 'prefix',
+    persistName = 'persist'
+  ): UseStore<T> & ISelectors<T> => {
+    return createSelectorHooks(
+      create(devtools(persist(createState, { name: persistName }), prefix))
+    )
+  }
+
+  const testPersistStore = createStoreWithPersist<ITestStateProps>(
+    (set) => ({
+      testKey: 'test',
+      setTestKey: (testKey: string) => {
+        set((state) => {
+          state.testKey = testKey
+        })
+      },
+    }),
+    'test',
+    'persist'
+  )
+
+  const TestComponent = (): JSX.Element => {
+    testPersistStore.use.testKey()
+    testPersistStore.use.setTestKey()
+
+    return <></>
+  }
+})
+
+it('should have correct type when creating store without middleware', () => {
+  const testStore = create<ITestStateProps>((set) => ({
+    testKey: 'test',
+    setTestKey: (testKey: string) => {
+      set((state) => {
+        state.testKey = testKey
+      })
+    },
+  }))
+
+  const TestComponent = (): JSX.Element => {
+    testStore((state) => state.testKey)
+    testStore((state) => state.setTestKey)
+
+    return <></>
+  }
+})
+
+it('should have correct type when creating store with persist', () => {
+  const createStoreWithPersist = <T extends State>(
+    createState: StateCreator<T>,
+    persistName = 'persist'
+  ): UseStore<T> & ISelectors<T> => {
+    return createSelectorHooks(
+      create(persist(createState, { name: persistName }))
+    )
+  }
+
+  const testPersistStore = createStoreWithPersist<ITestStateProps>(
+    (set) => ({
+      testKey: 'test',
+      setTestKey: (testKey: string) => {
+        set((state) => {
+          state.testKey = testKey
+        })
+      },
+    }),
+    'persist'
+  )
+
+  const TestComponent = (): JSX.Element => {
+    testPersistStore.use.testKey()
+    testPersistStore.use.setTestKey()
+
+    return <></>
+  }
+})
+
+it('should have correct type when creating store with immer', () => {
+  const createStoreWithImmer = <T extends State>(
+    createState: TImmerConfig<T>
+  ): UseStore<T> & ISelectors<T> => {
+    return createSelectorHooks(create(immer(createState)))
+  }
+
+  const testImmerStore = createStoreWithImmer<ITestStateProps>((set) => ({
+    testKey: 'test',
+    setTestKey: (testKey: string) => {
+      set((state) => {
+        state.testKey = testKey
+      })
+    },
+  }))
+
+  const TestComponent = (): JSX.Element => {
+    testImmerStore.use.testKey()
+    testImmerStore.use.setTestKey()
+
+    return <></>
+  }
+})
diff --git a/yarn.lock b/yarn.lock
index 591e084b6f..e36c100dfc 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -3755,6 +3755,11 @@ ignore@^5.1.4:
   resolved "https://registry.yarnpkg.com/ignore/-/ignore-5.1.8.tgz#f150a8b50a34289b33e22f5889abd4d8016f0e57"
   integrity sha512-BMpfD7PpiETpBl/A6S498BaIJ6Y/ABT93ETbby2fP00v4EbvPBXWEoaR1UBPKs3iR53pJY7EtZk5KACI57i1Uw==
 
+immer@^9.0.3:
+  version "9.0.3"
+  resolved "https://registry.yarnpkg.com/immer/-/immer-9.0.3.tgz#146e2ba8b84d4b1b15378143c2345559915097f4"
+  integrity sha512-mONgeNSMuyjIe0lkQPa9YhdmTv8P19IeHV0biYhcXhbd5dhdB9HSK93zBpyKjp6wersSUgT5QyU0skmejUVP2A==
+
 import-fresh@^3.0.0, import-fresh@^3.2.1:
   version "3.3.0"
   resolved "https://registry.yarnpkg.com/import-fresh/-/import-fresh-3.3.0.tgz#37162c25fcb9ebaa2e6e53d5b4d88ce17d9e0c2b"