From 6334414d8be5ed0006d959fc4aca511ef2aed9e5 Mon Sep 17 00:00:00 2001
From: cexbrayat <cedric@ninja-squad.com>
Date: Thu, 9 Sep 2021 17:39:12 +0200
Subject: [PATCH] feat: expose everything on wrapper.vm

This commit changes `wrapper.vm` to actually point to `vm.$.proxy`.
This shouldn't change the behaviour of existing tests, but allows components written with `script setup` to be tested as well,
without the need to expose everything just for testing purposes.

For example a component like:

```vue
<script setup lang="ts">
import { ref } from 'vue'

const count = ref(0)
</script>
```

can now be tested like `expect(wrapper.vm.count).toBe(0)`, whereas you previously had to add `defineExpose({ count })` for this to work.

The downside is that you now can't test that something is _not_ exposed by a component, but I don't think this is a problem.

This also removes the previous hacks for script setup, as it looks like they are no longer necessary.
---
 src/mount.ts                            | 23 +++--------------------
 src/vueWrapper.ts                       |  7 ++++++-
 tests/components/ScriptSetup_Expose.vue | 19 +++++++++++++++++++
 tests/vm.spec.ts                        | 19 ++++++++++++++++++-
 4 files changed, 46 insertions(+), 22 deletions(-)
 create mode 100644 tests/components/ScriptSetup_Expose.vue

diff --git a/src/mount.ts b/src/mount.ts
index f3c383d961..7bccf36c4f 100644
--- a/src/mount.ts
+++ b/src/mount.ts
@@ -337,17 +337,6 @@ export function mount(
   const Parent = defineComponent({
     name: 'VTU_ROOT',
     render() {
-      // https://github.com/vuejs/vue-test-utils-next/issues/651
-      // script setup components include an empty `expose` array as part of the
-      // code generated by the SFC compiler. Eg a component might look like
-      // { expose: [], setup: [Function], render: [Function] }
-      // not sure why (yet), but the empty expose array causes events to not be
-      // correctly captured.
-      // TODO: figure out why this is happening and understand the implications of
-      // the expose rfc for Test Utils.
-      if (isObjectComponent(component)) {
-        delete component.expose
-      }
       return h(component, props, slots)
     }
   })
@@ -457,19 +446,13 @@ export function mount(
   const warnSave = console.warn
   console.warn = () => {}
 
-  // get `vm`.
-  // for some unknown reason, getting the `vm` for components using `<script setup>`
-  // as of Vue 3.0.3 works differently.
-  // if `appRef` has keys, use that (vm always has keys like $el, $props etc).
-  // if not, use the return value from app.mount.
   const appRef = vm.$refs[MOUNT_COMPONENT_REF] as ComponentPublicInstance
-  const $vm = Reflect.ownKeys(appRef).length ? appRef : vm
   // we add `hasOwnProperty` so jest can spy on the proxied vm without throwing
-  $vm.hasOwnProperty = (property) => {
-    return Reflect.has($vm, property)
+  appRef.hasOwnProperty = (property) => {
+    return Reflect.has(appRef, property)
   }
   console.warn = warnSave
-  return createWrapper(app, $vm, setProps)
+  return createWrapper(app, appRef, setProps)
 }
 
 export const shallowMount: typeof mount = (component: any, options?: any) => {
diff --git a/src/vueWrapper.ts b/src/vueWrapper.ts
index a7906c56cc..78c577c753 100644
--- a/src/vueWrapper.ts
+++ b/src/vueWrapper.ts
@@ -36,7 +36,12 @@ export class VueWrapper<T extends ComponentPublicInstance>
     this.__app = app
     // root is null on functional components
     this.rootVM = vm?.$root
-    this.componentVM = vm as T
+    // vm.$.proxy is what the template has access to
+    // so even if the component is closed (as they are by default for `script setup`)
+    // a test will still be able to do something like
+    // `expect(wrapper.vm.count).toBe(1)`
+    // (note that vm can be null for functional components, hence the condition)
+    this.componentVM = vm ? vm.$.proxy as T : vm as T
     this.__setProps = setProps
 
     this.attachNativeEventListener()
diff --git a/tests/components/ScriptSetup_Expose.vue b/tests/components/ScriptSetup_Expose.vue
new file mode 100644
index 0000000000..d2dc97d4b9
--- /dev/null
+++ b/tests/components/ScriptSetup_Expose.vue
@@ -0,0 +1,19 @@
+<script setup lang="ts">
+// imported components are also directly usable in template
+import { ref } from 'vue'
+import Hello from './Hello.vue'
+
+const count = ref(0)
+const inc = () => {
+  count.value++
+}
+
+defineExpose({
+  count
+})
+</script>
+
+<template>
+  <button @click="inc">{{ count }}</button>
+  <Hello />
+</template>
diff --git a/tests/vm.spec.ts b/tests/vm.spec.ts
index 136d740297..445772dea9 100644
--- a/tests/vm.spec.ts
+++ b/tests/vm.spec.ts
@@ -1,6 +1,7 @@
 import { defineComponent, ref } from 'vue'
-
 import { mount } from '../src'
+import ScriptSetup from './components/ScriptSetup.vue'
+import ScriptSetupExpose from './components/ScriptSetup_Expose.vue'
 
 describe('vm', () => {
   it('returns the component vm', () => {
@@ -45,4 +46,20 @@ describe('vm', () => {
 
     expect(wrapper.vm.toggle).toHaveBeenCalled()
   })
+
+  it('access vm with <script setup> with defineExpose()', async () => {
+    const wrapper = mount(ScriptSetupExpose)
+
+    await wrapper.find('button').trigger('click')
+    expect(wrapper.html()).toContain('1')
+    expect(wrapper.vm.count).toBe(1)
+  })
+
+  it('access vm with <script setup> even without defineExpose()', async () => {
+    const wrapper = mount(ScriptSetup)
+
+    await wrapper.find('button').trigger('click')
+    expect(wrapper.html()).toContain('1')
+    expect(wrapper.vm.count).toBe(1)
+  })
 })