From 8f0d12a1851854a07b4bb47c5b28507cd65df822 Mon Sep 17 00:00:00 2001
From: cexbrayat <cedric@ninja-squad.com>
Date: Tue, 14 Apr 2020 13:59:56 +0200
Subject: [PATCH] fix: vm should be properly typed

The types have been tweaked to avoid the need of casting like `(wrapper.vm as any).msg`.
---
 src/mount.ts       | 21 ++++++++++++++++-----
 src/vue-wrapper.ts | 17 ++++++++---------
 tests/vm.spec.ts   | 11 ++++++++---
 3 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/src/mount.ts b/src/mount.ts
index 97b84104f..5b29504d3 100644
--- a/src/mount.ts
+++ b/src/mount.ts
@@ -8,10 +8,11 @@ import {
   Plugin,
   Directive,
   Component,
-  reactive
+  reactive,
+  ComponentPublicInstance
 } from 'vue'
 
-import { createWrapper } from './vue-wrapper'
+import { createWrapper, VueWrapper } from './vue-wrapper'
 import { createEmitMixin } from './emitMixin'
 import { createDataMixin } from './dataMixin'
 import { MOUNT_ELEMENT_ID } from './constants'
@@ -37,7 +38,14 @@ interface MountingOptions {
   stubs?: Record<string, any>
 }
 
-export function mount(originalComponent: any, options?: MountingOptions) {
+export function mount<T extends any>(
+  originalComponent: any,
+  options?: MountingOptions
+): VueWrapper<any>
+export function mount<T extends ComponentPublicInstance>(
+  originalComponent: new () => T,
+  options?: MountingOptions
+): VueWrapper<T> {
   const component = { ...originalComponent }
 
   // Reset the document.body
@@ -63,7 +71,10 @@ export function mount(originalComponent: any, options?: MountingOptions) {
   // override component data with mounting options data
   if (options?.data) {
     const dataMixin = createDataMixin(options.data())
-    component.mixins = [...(component.mixins || []), dataMixin]
+    ;(component as any).mixins = [
+      ...((component as any).mixins || []),
+      dataMixin
+    ]
   }
 
   // we define props as reactive so that way when we update them with `setProps`
@@ -136,5 +147,5 @@ export function mount(originalComponent: any, options?: MountingOptions) {
   // mount the app!
   const app = vm.mount(el)
 
-  return createWrapper(app, events, setProps)
+  return createWrapper<T>(app, events, setProps)
 }
diff --git a/src/vue-wrapper.ts b/src/vue-wrapper.ts
index 73960d350..4b22a33ac 100644
--- a/src/vue-wrapper.ts
+++ b/src/vue-wrapper.ts
@@ -6,8 +6,9 @@ import { WrapperAPI } from './types'
 import { ErrorWrapper } from './error-wrapper'
 import { MOUNT_ELEMENT_ID } from './constants'
 
-export class VueWrapper implements WrapperAPI {
-  private componentVM: ComponentPublicInstance
+export class VueWrapper<T extends ComponentPublicInstance>
+  implements WrapperAPI {
+  private componentVM: T
   private __emitted: Record<string, unknown[]> = {}
   private __vm: ComponentPublicInstance
   private __setProps: (props: Record<string, any>) => void
@@ -19,9 +20,7 @@ export class VueWrapper implements WrapperAPI {
   ) {
     this.__vm = vm
     this.__setProps = setProps
-    this.componentVM = this.__vm.$refs[
-      'VTU_COMPONENT'
-    ] as ComponentPublicInstance
+    this.componentVM = this.__vm.$refs['VTU_COMPONENT'] as T
     this.__emitted = events
   }
 
@@ -43,7 +42,7 @@ export class VueWrapper implements WrapperAPI {
     return this.hasMultipleRoots ? this.parentElement : this.componentVM.$el
   }
 
-  get vm(): ComponentPublicInstance {
+  get vm(): T {
     return this.componentVM
   }
 
@@ -106,10 +105,10 @@ export class VueWrapper implements WrapperAPI {
   }
 }
 
-export function createWrapper(
+export function createWrapper<T extends ComponentPublicInstance>(
   vm: ComponentPublicInstance,
   events: Record<string, unknown[]>,
   setProps: (props: Record<string, any>) => void
-): VueWrapper {
-  return new VueWrapper(vm, events, setProps)
+): VueWrapper<T> {
+  return new VueWrapper<T>(vm, events, setProps)
 }
diff --git a/tests/vm.spec.ts b/tests/vm.spec.ts
index 0d909d5c6..2613b3f0b 100644
--- a/tests/vm.spec.ts
+++ b/tests/vm.spec.ts
@@ -9,13 +9,18 @@ describe('vm', () => {
       setup() {
         const msg = 'hello'
         const isEnabled = ref(true)
-        return { msg, isEnabled }
+        const toggle = () => (isEnabled.value = !isEnabled.value)
+        return { msg, isEnabled, toggle }
       }
     })
 
     const wrapper = mount(Component)
 
-    expect((wrapper.vm as any).msg).toBe('hello')
-    expect((wrapper.vm as any).isEnabled).toBe(true)
+    expect(wrapper.vm.msg).toBe('hello')
+    expect(wrapper.vm.isEnabled).toBe(true)
+
+    wrapper.vm.toggle()
+
+    expect(wrapper.vm.isEnabled).toBe(false)
   })
 })