From e80b7b28046362fb1e5a281353637064722fbc6b 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/DefineExpose.vue       | 23 ++++++++++++++
 tests/components/ScriptSetup_Expose.vue | 19 ++++++++++++
 tests/expose.spec.ts                    | 40 +++++++++++++++++++++++++
 tests/vm.spec.ts                        |  1 -
 tsconfig.volar.json                     |  3 +-
 7 files changed, 93 insertions(+), 23 deletions(-)
 create mode 100644 tests/components/DefineExpose.vue
 create mode 100644 tests/components/ScriptSetup_Expose.vue
 create mode 100644 tests/expose.spec.ts

diff --git a/src/mount.ts b/src/mount.ts
index 4fa1fe3fa..8deca1982 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 a7906c56c..78cf72757 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/DefineExpose.vue b/tests/components/DefineExpose.vue
new file mode 100644
index 000000000..d9319ee89
--- /dev/null
+++ b/tests/components/DefineExpose.vue
@@ -0,0 +1,23 @@
+<template>
+  <div id="root">
+    <div id="msg">{{ msg }}</div>
+    <div>{{ other }}</div>
+  </div>
+</template>
+
+<script lang="ts">
+import { defineComponent, ref } from 'vue'
+
+export default defineComponent({
+  name: 'Hello',
+
+  setup(props, { expose }) {
+    const other = ref('other')
+    expose({ other })
+    return {
+      msg: ref('Hello world'),
+      other
+    }
+  }
+})
+</script>
diff --git a/tests/components/ScriptSetup_Expose.vue b/tests/components/ScriptSetup_Expose.vue
new file mode 100644
index 000000000..d2dc97d4b
--- /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/expose.spec.ts b/tests/expose.spec.ts
new file mode 100644
index 000000000..c4bcc4036
--- /dev/null
+++ b/tests/expose.spec.ts
@@ -0,0 +1,40 @@
+import { mount } from '../src'
+import Hello from './components/Hello.vue'
+import DefineExpose from './components/DefineExpose.vue'
+import ScriptSetupExpose from './components/ScriptSetup_Expose.vue'
+import ScriptSetup from './components/ScriptSetup.vue'
+
+describe('expose', () => {
+  it('access vm on simple components', async () => {
+    const wrapper = mount(Hello)
+
+    expect(wrapper.vm.msg).toBe('Hello world')
+  })
+
+  it('access vm on simple components with custom `expose`', async () => {
+    const wrapper = mount(DefineExpose)
+
+    // other is exposed vie `expose`
+    expect(wrapper.vm.other).toBe('other')
+    // can access `msg` even if not exposed
+    expect(wrapper.vm.msg).toBe('Hello world')
+  })
+
+  it('access vm with <script setup> and defineExpose()', async () => {
+    const wrapper = mount(ScriptSetupExpose)
+
+    await wrapper.find('button').trigger('click')
+    expect(wrapper.html()).toContain('1')
+    // can access `count` as it is exposed via `defineExpose()`
+    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')
+    // can access `count` even if it is _not_ exposed
+    expect(wrapper.vm.count).toBe(1)
+  })
+})
diff --git a/tests/vm.spec.ts b/tests/vm.spec.ts
index 136d74029..a278a8524 100644
--- a/tests/vm.spec.ts
+++ b/tests/vm.spec.ts
@@ -1,5 +1,4 @@
 import { defineComponent, ref } from 'vue'
-
 import { mount } from '../src'
 
 describe('vm', () => {
diff --git a/tsconfig.volar.json b/tsconfig.volar.json
index fe1ec6d81..7fd9fdd91 100644
--- a/tsconfig.volar.json
+++ b/tsconfig.volar.json
@@ -3,5 +3,6 @@
   "compilerOptions": {
     "lib": ["DOM", "ES2020"],
     "skipLibCheck": true
-  }
+  },
+  "exclude": ["tests/expose.spec.ts"]
 }