refactor: Refactor the way VueWrapper is created, to allow for wrapping nested instances.

Author: dobromir-hristov

src/emitMixin.ts

@@ -1,10 +1,11 @@
-import { getCurrentInstance } from 'vue'
-
-export const createEmitMixin = () => {
-  const events: Record<string, unknown[]> = {}
+import { getCurrentInstance, App } from 'vue'
 
+export const attachEventListener = (vm: App) => {
   const emitMixin = {
     beforeCreate() {
+      let events: Record<string, unknown[]> = {}
+      this.__emitted = events
+
       getCurrentInstance().emit = (event: string, ...args: unknown[]) => {
         events[event]
           ? (events[event] = [...events[event], [...args]])
@@ -15,8 +16,5 @@ export const createEmitMixin = () => {
     }
   }
 
-  return {
-    events,
-    emitMixin
-  }
+  vm.mixin(emitMixin)
 }

src/error-wrapper.ts

@@ -1,9 +1,11 @@
+import { FindComponentSelector } from './types'
+
 interface Options {
-  selector: string
+  selector: FindComponentSelector
 }
 
 export class ErrorWrapper {
-  selector: string
+  selector: FindComponentSelector
   element: null
 
   constructor({ selector }: Options) {
@@ -14,6 +16,10 @@ export class ErrorWrapper {
     return Error(`Cannot call ${method} on an empty wrapper.`)
   }
 
+  vm(): Error {
+    throw this.wrapperError('vm')
+  }
+
   attributes() {
     throw this.wrapperError('attributes')
   }
@@ -34,8 +40,8 @@ export class ErrorWrapper {
     throw this.wrapperError('findAll')
   }
 
-  setChecked() {
-    throw this.wrapperError('setChecked')
+  setProps() {
+    throw this.wrapperError('setProps')
   }
 
   setValue() {

src/mount.ts

@@ -9,11 +9,12 @@ import {
   Plugin,
   Directive,
   Component,
-  reactive
+  reactive,
+  ComponentPublicInstance
 } from 'vue'
 
 import { createWrapper } from './vue-wrapper'
-import { createEmitMixin } from './emitMixin'
+import { attachEventListener } from './emitMixin'
 import { createDataMixin } from './dataMixin'
 import { MOUNT_ELEMENT_ID } from './constants'
 import { stubComponents } from './stubs'
@@ -134,8 +135,7 @@ export function mount(originalComponent: any, options?: MountingOptions) {
   }
 
   // add tracking for emitted events
-  const { emitMixin, events } = createEmitMixin()
-  vm.mixin(emitMixin)
+  attachEventListener(vm)
 
   // stubs
   if (options?.global?.stubs) {
@@ -146,6 +146,6 @@ export function mount(originalComponent: any, options?: MountingOptions) {
 
   // mount the app!
   const app = vm.mount(el)
-
-  return createWrapper(app, events, setProps)
+  const App = app.$refs['VTU_COMPONENT'] as ComponentPublicInstance
+  return createWrapper(App, setProps)
 }

src/vue-wrapper.ts

@@ -8,44 +8,35 @@ import {
   WrapperAPI
 } from './types'
 import { ErrorWrapper } from './error-wrapper'
-import { MOUNT_ELEMENT_ID } from './constants'
 import { find } from './utils/find'
 
 export class VueWrapper implements WrapperAPI {
   private componentVM: ComponentPublicInstance
-  private __emitted: Record<string, unknown[]> = {}
-  private __vm: ComponentPublicInstance
+  private rootVM: ComponentPublicInstance
   private __setProps: (props: Record<string, any>) => void
 
   constructor(
     vm: ComponentPublicInstance,
-    events: Record<string, unknown[]>,
-    setProps: (props: Record<string, any>) => void
+    setProps?: (props: Record<string, any>) => void
   ) {
-    this.__vm = vm
+    // TODO Remove cast after Vue releases the fix
+    this.rootVM = (vm.$root as any) as ComponentPublicInstance
+    this.componentVM = vm
     this.__setProps = setProps
-    this.componentVM = this.__vm.$refs[
-      'VTU_COMPONENT'
-    ] as ComponentPublicInstance
-    this.__emitted = events
-  }
-
-  private get appRootNode() {
-    return document.getElementById(MOUNT_ELEMENT_ID) as HTMLDivElement
   }
 
   private get hasMultipleRoots(): boolean {
     // if the subtree is an array of children, we have multiple root nodes
-    return this.componentVM.$.subTree.shapeFlag === ShapeFlags.ARRAY_CHILDREN
+    return this.vm.$.subTree.shapeFlag === ShapeFlags.ARRAY_CHILDREN
   }
 
   private get parentElement(): Element {
-    return this.componentVM.$el.parentElement
+    return this.vm.$el.parentElement
   }
 
   get element(): Element {
     // if the component has multiple root elements, we use the parent's element
-    return this.hasMultipleRoots ? this.parentElement : this.componentVM.$el
+    return this.hasMultipleRoots ? this.parentElement : this.vm.$el
   }
 
   get vm(): ComponentPublicInstance {
@@ -64,8 +55,10 @@ export class VueWrapper implements WrapperAPI {
     return true
   }
 
-  emitted() {
-    return this.__emitted
+  emitted(): Record<string, unknown[]> {
+    // TODO Should we define this?
+    // @ts-ignore
+    return this.vm.__emitted
   }
 
   html() {
@@ -95,26 +88,32 @@ export class VueWrapper implements WrapperAPI {
     return result
   }
 
-  findComponent(selector: FindComponentSelector): ComponentPublicInstance {
+  findComponent(selector: FindComponentSelector): VueWrapper | ErrorWrapper {
     if (typeof selector === 'object' && 'ref' in selector) {
-      return this.componentVM.$refs[selector.ref] as ComponentPublicInstance
+      return createWrapper(
+        this.vm.$refs[selector.ref] as ComponentPublicInstance
+      )
     }
-    const result = find(this.componentVM.$.subTree, selector)
-    return result.length ? result[0] : undefined
+    const result = find(this.vm.$.subTree, selector)
+    if (!result.length) return new ErrorWrapper({ selector })
+    return createWrapper(result[0])
   }
 
-  findAllComponents(
-    selector: FindAllComponentsSelector
-  ): ComponentPublicInstance[] {
-    return find(this.componentVM.$.subTree, selector)
+  findAllComponents(selector: FindAllComponentsSelector): VueWrapper[] {
+    return find(this.vm.$.subTree, selector).map((c) => createWrapper(c))
   }
 
   findAll<T extends Element>(selector: string): DOMWrapper<T>[] {
     const results = this.parentElement.querySelectorAll<T>(selector)
     return Array.from(results).map((x) => new DOMWrapper(x))
   }
 
-  setProps(props: Record<string, any>) {
+  setProps(props: Record<string, any>): Promise<void> {
+    // if this VM's parent is not the root, error out
+    // TODO: Remove ignore after Vue releases fix
+    // @ts-ignore
+    if (this.vm.$parent !== this.rootVM)
+      throw Error('You can only use setProps on your mounted component')
     this.__setProps(props)
     return nextTick()
   }
@@ -127,8 +126,7 @@ export class VueWrapper implements WrapperAPI {
 
 export function createWrapper(
   vm: ComponentPublicInstance,
-  events: Record<string, unknown[]>,
-  setProps: (props: Record<string, any>) => void
+  setProps?: (props: Record<string, any>) => void
 ): VueWrapper {
-  return new VueWrapper(vm, events, setProps)
+  return new VueWrapper(vm, setProps)
 }

tests/findAllComponents.spec.ts

@@ -19,11 +19,9 @@ describe('findAllComponents', () => {
     const wrapper = mount(compA)
     // find by DOM selector
     expect(wrapper.findAllComponents('.C')).toHaveLength(2)
-    expect(
-      wrapper.findAllComponents({ name: 'Hello' })[0].$el.textContent
-    ).toBe('Hello world')
-    expect(wrapper.findAllComponents(Hello)[0].$el.textContent).toBe(
+    expect(wrapper.findAllComponents({ name: 'Hello' })[0].text()).toBe(
       'Hello world'
     )
+    expect(wrapper.findAllComponents(Hello)[0].text()).toBe('Hello world')
   })
 })

tests/findComponent.spec.ts

@@ -37,7 +37,7 @@ const compA = {
 describe('findComponent', () => {
   it('does not find plain dom elements', () => {
     const wrapper = mount(compA)
-    expect(wrapper.findComponent('.domElement')).toBeFalsy()
+    expect(wrapper.findComponent('.domElement').exists()).toBeFalsy()
   })
 
   it('finds component by ref', () => {
@@ -49,34 +49,39 @@ describe('findComponent', () => {
   it('finds component by dom selector', () => {
     const wrapper = mount(compA)
     // find by DOM selector
-    expect(wrapper.findComponent('.C').$options.name).toEqual('ComponentC')
+    expect(wrapper.findComponent('.C').vm).toHaveProperty(
+      '$options.name',
+      'ComponentC'
+    )
   })
 
   it('does allows using complicated DOM selector query', () => {
     const wrapper = mount(compA)
-    expect(wrapper.findComponent('.B > .C').$options.name).toEqual('ComponentC')
+    expect(wrapper.findComponent('.B > .C').vm).toHaveProperty(
+      '$options.name',
+      'ComponentC'
+    )
   })
 
   it('finds a component when root of mounted component', async () => {
     const wrapper = mount(compD)
     // make sure it finds the component, not its root
-    expect(wrapper.findComponent('.c-as-root-on-d').$options.name).toEqual(
+    expect(wrapper.findComponent('.c-as-root-on-d').vm).toHaveProperty(
+      '$options.name',
       'ComponentC'
     )
   })
 
   it('finds component by name', () => {
     const wrapper = mount(compA)
-    expect(wrapper.findComponent({ name: 'Hello' }).$el.textContent).toBe(
-      'Hello world'
-    )
-    expect(wrapper.findComponent({ name: 'ComponentB' })).toBeTruthy()
-    expect(wrapper.findComponent({ name: 'component-c' })).toBeTruthy()
+    expect(wrapper.findComponent({ name: 'Hello' }).text()).toBe('Hello world')
+    expect(wrapper.findComponent({ name: 'ComponentB' }).exists()).toBeTruthy()
+    expect(wrapper.findComponent({ name: 'component-c' }).exists()).toBeTruthy()
   })
 
   it('finds component by imported SFC file', () => {
     const wrapper = mount(compA)
-    expect(wrapper.findComponent(Hello).$el.textContent).toBe('Hello world')
-    expect(wrapper.findComponent(compC).$el.textContent).toBe('C')
+    expect(wrapper.findComponent(Hello).text()).toBe('Hello world')
+    expect(wrapper.findComponent(compC).text()).toBe('C')
   })
 })

tests/setProps.spec.ts

@@ -10,10 +10,10 @@ describe('setProps', () => {
     }
     const wrapper = mount(Foo, {
       props: {
-        foo: 'foo'
+        foo: 'bar'
       }
     })
-    expect(wrapper.html()).toContain('foo')
+    expect(wrapper.html()).toContain('bar')
 
     await wrapper.setProps({ foo: 'qux' })
     expect(wrapper.html()).toContain('qux')
@@ -44,7 +44,8 @@ describe('setProps', () => {
   it('sets component props, and updates DOM when props were not initially passed', async () => {
     const Foo = {
       props: ['foo'],
-      template: `<div>{{ foo }}</div>`
+      template: `
+        <div>{{ foo }}</div>`
     }
     const wrapper = mount(Foo)
     expect(wrapper.html()).not.toContain('foo')
@@ -67,7 +68,8 @@ describe('setProps', () => {
           this.bar = val
         }
       },
-      template: `<div>{{ bar }}</div>`
+      template: `
+        <div>{{ bar }}</div>`
     }
     const wrapper = mount(Foo)
     expect(wrapper.html()).toContain('original-bar')
@@ -118,4 +120,27 @@ describe('setProps', () => {
     expect(wrapper.attributes()).toEqual(nonExistentProp)
     expect(wrapper.html()).toBe('<div bar="qux">foo</div>')
   })
+
+  it('allows using only on mounted component', async () => {
+    const Foo = {
+      name: 'Foo',
+      props: ['foo'],
+      template: '<div>{{ foo }}</div>'
+    }
+    const Baz = {
+      props: ['baz'],
+      template: '<div><Foo :foo="baz"/></div>',
+      components: { Foo }
+    }
+
+    const wrapper = mount(Baz, {
+      props: {
+        baz: 'baz'
+      }
+    })
+    const FooResult = wrapper.findComponent({ name: 'Foo' })
+    expect(() => FooResult.setProps({ baz: 'bin' })).toThrowError(
+      'You can only use setProps on your mounted component'
+    )
+  })
 })

tests/vm.spec.ts

@@ -5,6 +5,7 @@ import { mount } from '../src'
 describe('vm', () => {
   it('returns the component vm', () => {
     const Component = defineComponent({
+      name: 'VTUComponent',
       template: '<div>{{ msg }}</div>',
       setup() {
         const msg = 'hello'