From f6fee2c31b254b33dec06addecba2218d6dc8cd1 Mon Sep 17 00:00:00 2001
From: Eduardo San Martin Morote <posva13@gmail.com>
Date: Fri, 16 Jul 2021 19:09:39 +0200
Subject: [PATCH] fix(hmr): add state properties to the store

---
 __tests__/hmr.spec.ts | 103 ++++++++++++++++++++++++++++++++++++++++++
 src/store.ts          |   9 ++--
 2 files changed, 107 insertions(+), 5 deletions(-)
 create mode 100644 __tests__/hmr.spec.ts

diff --git a/__tests__/hmr.spec.ts b/__tests__/hmr.spec.ts
new file mode 100644
index 00000000..09e6ea05
--- /dev/null
+++ b/__tests__/hmr.spec.ts
@@ -0,0 +1,103 @@
+import { watch } from 'vue'
+import {
+  createPinia,
+  defineSetupStore,
+  defineStore,
+  DefineStoreOptions,
+  setActivePinia,
+  StateTree,
+} from '../src'
+
+function defineOptions<
+  O extends DefineStoreOptions<string, StateTree, any, any>
+>(options: O): O {
+  return options
+}
+
+describe('HMR', () => {
+  const baseOptions = defineOptions({
+    id: 'main',
+    state: () => ({
+      n: 0,
+    }),
+
+    actions: {
+      increment(amount = 1) {
+        // @ts-ignore
+        this.n += amount
+      },
+    },
+
+    getters: {
+      double: (state: any) => state.n * 2,
+    },
+  })
+
+  beforeEach(() => {
+    setActivePinia(createPinia())
+  })
+
+  it('adds new state properties', () => {
+    const useStore = defineStore(baseOptions)
+    const store: any = useStore()
+    store.n++
+
+    // simulate a hmr
+    defineStore({
+      ...baseOptions,
+      state: () => ({ newOne: 'hey', n: 0 }),
+    })(null, store)
+
+    expect(store.$state).toEqual({ n: 1, newOne: 'hey' })
+    expect(store.n).toBe(1)
+    expect(store.newOne).toBe('hey')
+
+    defineStore({
+      ...baseOptions,
+      state: () => ({ other: 'new', n: 0 }),
+    })(null, store)
+  })
+
+  it('keeps state reactive', () => {
+    const useStore = defineStore(baseOptions)
+    const store: any = useStore()
+
+    const directSpy = jest.fn()
+    const $stateSpy = jest.fn()
+
+    watch(() => store.n, directSpy, { flush: 'sync' })
+    watch(() => store.$state.n, $stateSpy, { flush: 'sync' })
+
+    // simulate a hmr
+    defineStore({
+      ...baseOptions,
+      state: () => ({ newOne: 'hey', n: 0 }),
+    })(null, store)
+
+    expect(directSpy).toHaveBeenCalledTimes(0)
+    expect($stateSpy).toHaveBeenCalledTimes(0)
+
+    store.n++
+    expect(directSpy).toHaveBeenCalledTimes(1)
+    expect($stateSpy).toHaveBeenCalledTimes(1)
+
+    store.$state.n++
+    expect(directSpy).toHaveBeenCalledTimes(2)
+    expect($stateSpy).toHaveBeenCalledTimes(2)
+
+    defineStore({
+      ...baseOptions,
+      state: () => ({ other: 'new', n: 0 }),
+    })(null, store)
+
+    store.n++
+    expect(directSpy).toHaveBeenCalledTimes(3)
+    expect($stateSpy).toHaveBeenCalledTimes(3)
+
+    store.$state.n++
+    expect(directSpy).toHaveBeenCalledTimes(4)
+    expect($stateSpy).toHaveBeenCalledTimes(4)
+  })
+
+  it.todo('handles nested objects updates')
+})
diff --git a/src/store.ts b/src/store.ts
index c637dc4d..fb6ee7f7 100644
--- a/src/store.ts
+++ b/src/store.ts
@@ -463,12 +463,11 @@ function createSetupStore<
             // ---
             // @ts-expect-error
             store.$state[stateKey]
-
-          // patch direct access properties to allow store.stateProperty to work as
-          // store.$state.stateProperty
-          // @ts-expect-error
-          store[stateKey] = toRef(newStore.$state, stateKey)
         }
+        // patch direct access properties to allow store.stateProperty to work as
+        // store.$state.stateProperty
+        // @ts-expect-error
+        store[stateKey] = toRef(newStore.$state, stateKey)
       })
 
       pinia.state.value[$id] = toRef(newStore._hmrPayload, 'hotState')
-- 
2.47.2