From 455ad955ea4220d8434c327c6e54521a1ca93141 Mon Sep 17 00:00:00 2001
From: Eduardo San Martin Morote <posva13@gmail.com>
Date: Sat, 10 Apr 2021 17:13:25 +0200
Subject: [PATCH] fix(subscribe): remove subscription when unmounted

---
 __tests__/lifespan.spec.ts      |  1 -
 __tests__/subscriptions.spec.ts | 58 ++++++++++++++++++++++++++++++++-
 src/store.ts                    |  9 ++++-
 3 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/__tests__/lifespan.spec.ts b/__tests__/lifespan.spec.ts
index af84547b..869868df 100644
--- a/__tests__/lifespan.spec.ts
+++ b/__tests__/lifespan.spec.ts
@@ -39,7 +39,6 @@ describe('Store Lifespan', () => {
 
   it.skip('what', async () => {
     const localVue = createLocalVue()
-    localVue.use(VueCompositionAPI)
     const n = 0
     const Component = defineComponent({
       render: (h) => h('p'),
diff --git a/__tests__/subscriptions.spec.ts b/__tests__/subscriptions.spec.ts
index 5aabbbdc..3277360d 100644
--- a/__tests__/subscriptions.spec.ts
+++ b/__tests__/subscriptions.spec.ts
@@ -1,5 +1,13 @@
+import { nextTick } from '@vue/composition-api'
+import { createLocalVue, mount } from '@vue/test-utils'
 import Vue from 'vue'
-import { defineStore, setActivePinia, createPinia, Pinia } from '../src'
+import {
+  defineStore,
+  setActivePinia,
+  createPinia,
+  Pinia,
+  PiniaPlugin,
+} from '../src'
 
 describe('Subscriptions', () => {
   let pinia: Pinia
@@ -57,6 +65,9 @@ describe('Subscriptions', () => {
     })
 
     it('triggers subscribe only once', async () => {
+      const pinia = createPinia()
+      pinia.Vue = Vue
+      setActivePinia(pinia)
       const s1 = useStore()
       const s2 = useStore()
 
@@ -74,5 +85,50 @@ describe('Subscriptions', () => {
       expect(spy1).toHaveBeenCalledTimes(1)
       expect(spy2).toHaveBeenCalledTimes(1)
     })
+
+    it('removes on unmount', async () => {
+      const pinia = createPinia()
+      pinia.Vue = Vue
+      const spy1 = jest.fn()
+      const spy2 = jest.fn()
+
+      const localVue = createLocalVue()
+      localVue.use(PiniaPlugin)
+
+      const wrapper = mount(
+        {
+          setup() {
+            const s1 = useStore()
+            s1.$subscribe(spy1)
+          },
+          template: `<p/>`,
+        },
+        { pinia, localVue }
+      )
+
+      const s1 = useStore()
+      const s2 = useStore()
+
+      s2.$subscribe(spy2)
+
+      expect(spy1).toHaveBeenCalledTimes(0)
+      expect(spy2).toHaveBeenCalledTimes(0)
+
+      s1.name = 'Edu'
+
+      expect(spy2).toHaveBeenCalledTimes(1)
+      expect(spy1).toHaveBeenCalledTimes(1)
+
+      s1.$patch({ name: 'a' })
+      expect(spy1).toHaveBeenCalledTimes(2)
+      expect(spy2).toHaveBeenCalledTimes(2)
+
+      wrapper.destroy()
+      await nextTick()
+
+      s1.$patch({ name: 'b' })
+      expect(spy1).toHaveBeenCalledTimes(2)
+      expect(spy2).toHaveBeenCalledTimes(3)
+    })
   })
 })
diff --git a/src/store.ts b/src/store.ts
index 491d1e13..e6453222 100644
--- a/src/store.ts
+++ b/src/store.ts
@@ -6,6 +6,7 @@ import {
   getCurrentInstance,
   markRaw,
   inject,
+  onUnmounted,
 } from '@vue/composition-api'
 import {
   StateTree,
@@ -141,13 +142,19 @@ function initStore<Id extends string, S extends StateTree>(
       }
     )
 
-    return () => {
+    const removeSubscription = () => {
       const idx = subscriptions.indexOf(callback)
       if (idx > -1) {
         subscriptions.splice(idx, 1)
         stopWatcher()
       }
     }
+
+    if (getCurrentInstance()) {
+      onUnmounted(removeSubscription)
+    }
+
+    return removeSubscription
   }
 
   function $reset() {
-- 
2.47.3