From b88ef45102152648bc4a189c5eed0c650fb3e2dc Mon Sep 17 00:00:00 2001
From: Illya Klymov <xanf@xanf.me>
Date: Tue, 5 Oct 2021 23:41:33 +0300
Subject: [PATCH] feat: add enableAutoUnmount feature

- port enableAutoDestroy feature from VTU v1
---
 docs/api/index.md         | 22 ++++++++++++++++++++++
 src/index.ts              |  3 +++
 src/mount.ts              |  5 ++++-
 src/utils/autoUnmount.ts  | 32 ++++++++++++++++++++++++++++++++
 tests/autoUnmount.spec.ts | 37 +++++++++++++++++++++++++++++++++++++
 5 files changed, 98 insertions(+), 1 deletion(-)
 create mode 100644 src/utils/autoUnmount.ts
 create mode 100644 tests/autoUnmount.spec.ts

diff --git a/docs/api/index.md b/docs/api/index.md
index c5549bda9..167bcf47a 100644
--- a/docs/api/index.md
+++ b/docs/api/index.md
@@ -1829,6 +1829,28 @@ function shallowMount(Component, options?: MountingOptions): VueWrapper
 
 `shallowMount` behaves exactly like `mount`, but it stubs all child components by default. Essentially, `shallowMount(Component)` is an alias of `mount(Component, { shallow: true })`.
 
+## enableAutoUnmount
+
+
+**Signature:**
+```ts
+enableAutoUnmount(hook: Function));
+disableAutoUnmount(): void;
+```
+
+**Details:**
+
+`enableAutoUnmount` allows to automatically destroy Vue wrappers. Destroy logic is passed as callback to `hook` Function.
+Common usage is to use `enableAutoUnmount` with teardown helper functions provided by your test framework, such as `afterEach`:
+
+```ts
+import { enableAutoUnmount } from '@vue/test-utils'
+
+enableAutoUnmount(afterEach)
+```
+
+`disableAutoUnmount` might be useful if you want this behavior only in specific subset of your test suite and you want to explicitly disable this behavior
+
 ## flushPromises
 
 **Signature:**
diff --git a/src/index.ts b/src/index.ts
index 3ed2e7a5b..1f9e8dc75 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -6,10 +6,13 @@ import { DOMWrapper } from './domWrapper'
 import { createWrapperError } from './errorWrapper'
 import { config } from './config'
 import { flushPromises } from './utils/flushPromises'
+import { enableAutoUnmount, disableAutoUnmount } from './utils/autoUnmount'
 
 export {
   mount,
   shallowMount,
+  enableAutoUnmount,
+  disableAutoUnmount,
   RouterLinkStub,
   VueWrapper,
   DOMWrapper,
diff --git a/src/mount.ts b/src/mount.ts
index 8deca1982..f19a40923 100644
--- a/src/mount.ts
+++ b/src/mount.ts
@@ -45,6 +45,7 @@ import {
   isLegacyFunctionalComponent,
   unwrapLegacyVueExtendComponent
 } from './utils/vueCompatSupport'
+import { trackInstance } from './utils/autoUnmount'
 
 // NOTE this should come from `vue`
 type PublicProps = VNodeProps & AllowedComponentProps & ComponentCustomProps
@@ -452,7 +453,9 @@ export function mount(
     return Reflect.has(appRef, property)
   }
   console.warn = warnSave
-  return createWrapper(app, appRef, setProps)
+  const wrapper = createWrapper(app, appRef, setProps)
+  trackInstance(wrapper)
+  return wrapper
 }
 
 export const shallowMount: typeof mount = (component: any, options?: any) => {
diff --git a/src/utils/autoUnmount.ts b/src/utils/autoUnmount.ts
new file mode 100644
index 000000000..5b81fa998
--- /dev/null
+++ b/src/utils/autoUnmount.ts
@@ -0,0 +1,32 @@
+import { ComponentPublicInstance } from 'vue'
+import type { VueWrapper } from '../vueWrapper'
+
+let isEnabled = false
+const wrapperInstances: VueWrapper<ComponentPublicInstance>[] = []
+
+export function disableAutoUnmount() {
+  isEnabled = false
+  wrapperInstances.length = 0
+}
+
+export function enableAutoUnmount(hook: Function) {
+  if (isEnabled) {
+    throw new Error('enableAutoUnmount cannot be called more than once')
+  }
+
+  isEnabled = true
+
+  hook(() => {
+    wrapperInstances.forEach((wrapper: VueWrapper<ComponentPublicInstance>) => {
+      wrapper.unmount()
+    })
+
+    wrapperInstances.length = 0
+  })
+}
+
+export function trackInstance(wrapper: VueWrapper<ComponentPublicInstance>) {
+  if (!isEnabled) return
+
+  wrapperInstances.push(wrapper)
+}
diff --git a/tests/autoUnmount.spec.ts b/tests/autoUnmount.spec.ts
new file mode 100644
index 000000000..b03b404cf
--- /dev/null
+++ b/tests/autoUnmount.spec.ts
@@ -0,0 +1,37 @@
+import { mount, enableAutoUnmount, disableAutoUnmount } from '../src'
+
+describe('enableAutoUnmount', () => {
+  beforeEach(() => {
+    disableAutoUnmount()
+  })
+
+  it('calls the hook function', () => {
+    const hookMock = jest.fn()
+
+    enableAutoUnmount(hookMock)
+
+    expect(hookMock).toHaveBeenCalledWith(expect.any(Function))
+  })
+
+  it('uses the hook function to unmount wrappers', () => {
+    const hookMock = jest.fn()
+
+    enableAutoUnmount(hookMock)
+    const [unmountFn] = hookMock.mock.calls[0]
+
+    const wrapper = mount({ template: '<p>test</p>' })
+    jest.spyOn(wrapper, 'unmount')
+
+    unmountFn()
+
+    expect(wrapper.unmount).toHaveBeenCalledTimes(1)
+  })
+
+  it('cannot be called twice', () => {
+    const noop = () => {}
+
+    enableAutoUnmount(noop)
+
+    expect(() => enableAutoUnmount(noop)).toThrow()
+  })
+})