From dac59509986674d3598cd548407bb98f873fc85c Mon Sep 17 00:00:00 2001
From: Alec Gibson <12036746+alecgibson@users.noreply.github.com>
Date: Tue, 25 Jul 2023 11:58:02 +0100
Subject: [PATCH] fix: stricter props types

At the moment, `mount()` offers [strong typing][1] of props, but this
strong typing is lost when dealing with the `VueWrapper` through either
the `props()` or `setProps()` methods.

This change strengthens the typing of these methods to help raise
compile-time errors when trying to get or set incorrect props.

[1]: https://github.com/vuejs/test-utils/blob/11b34745e8e66fc747881dfb1ce94cef537c455e/src/types.ts#L44
---
 src/mount.ts               |  8 +++++++-
 src/vueWrapper.ts          | 14 +++++++++-----
 test-dts/wrapper.d-test.ts | 27 ++++++++++++++++++++++++++-
 tests/getComponent.spec.ts | 10 +++++++---
 tests/props.spec.ts        |  1 +
 5 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/src/mount.ts b/src/mount.ts
index e4450c25a..f43833a04 100644
--- a/src/mount.ts
+++ b/src/mount.ts
@@ -47,7 +47,13 @@ export function mount<
 >(
   originalComponent: T,
   options?: ComponentMountingOptions<C>
-): VueWrapper<ComponentExposed<C> & ComponentProps<C> & ComponentData<C>>
+): VueWrapper<
+  ComponentProps<C> & ComponentData<C> & ComponentExposed<C>,
+  ComponentPublicInstance<
+    ComponentProps<C>,
+    ComponentData<C> & ComponentExposed<C>
+  >
+>
 
 // implementation
 export function mount(
diff --git a/src/vueWrapper.ts b/src/vueWrapper.ts
index 27d5f9744..e1d1e774e 100644
--- a/src/vueWrapper.ts
+++ b/src/vueWrapper.ts
@@ -215,10 +215,14 @@ export class VueWrapper<
     return this.componentVM
   }
 
-  props(): { [key: string]: any }
-  props(selector: string): any
-  props(selector?: string): { [key: string]: any } | any {
-    const props = this.componentVM.$props as { [key: string]: any }
+  props(): T['$props']
+  props<Selector extends keyof T['$props']>(
+    selector: Selector
+  ): T['$props'][Selector]
+  props<Selector extends keyof T['$props']>(
+    selector?: Selector
+  ): T['$props'] | T['$props'][Selector] {
+    const props = this.componentVM.$props as T['$props']
     return selector ? props[selector] : props
   }
 
@@ -240,7 +244,7 @@ export class VueWrapper<
     return nextTick()
   }
 
-  setProps(props: Record<string, unknown>): Promise<void> {
+  setProps(props: T['$props']): Promise<void> {
     // if this VM's parent is not the root or if setProps does not exist, error out
     if (this.vm.$parent !== this.rootVM || !this.__setProps) {
       throw Error('You can only use setProps on your mounted component')
diff --git a/test-dts/wrapper.d-test.ts b/test-dts/wrapper.d-test.ts
index fcd3fc21c..811376ab7 100644
--- a/test-dts/wrapper.d-test.ts
+++ b/test-dts/wrapper.d-test.ts
@@ -116,4 +116,29 @@ expectType<boolean>(domWrapper.classes('class'))
 
 // props
 expectType<{ [key: string]: any }>(wrapper.props())
-expectType<any>(wrapper.props('prop'))
+
+const ComponentWithProps = defineComponent({
+  props: {
+    foo: String,
+    bar: Number,
+  },
+})
+
+const propsWrapper = mount(ComponentWithProps);
+
+propsWrapper.setProps({foo: 'abc'})
+propsWrapper.setProps({foo: 'abc', bar: 123})
+// @ts-expect-error :: should require string
+propsWrapper.setProps({foo: 123})
+// @ts-expect-error :: unknown prop
+propsWrapper.setProps({badProp: true})
+
+expectType<string | undefined>(propsWrapper.props().foo)
+expectType<number | undefined>(propsWrapper.props().bar)
+// @ts-expect-error :: unknown prop
+propsWrapper.props().badProp;
+
+expectType<string | undefined>(propsWrapper.props('foo'))
+expectType<number | undefined>(propsWrapper.props('bar'))
+// @ts-expect-error :: unknown prop
+propsWrapper.props('badProp')
diff --git a/tests/getComponent.spec.ts b/tests/getComponent.spec.ts
index d58057de4..117678278 100644
--- a/tests/getComponent.spec.ts
+++ b/tests/getComponent.spec.ts
@@ -1,5 +1,5 @@
 import { describe, expect, it, vi } from 'vitest'
-import { DefineComponent, defineComponent } from 'vue'
+import { defineComponent } from 'vue'
 import { mount, RouterLinkStub, shallowMount } from '../src'
 import Issue425 from './components/Issue425.vue'
 
@@ -70,7 +70,9 @@ describe('getComponent', () => {
   // https://github.com/vuejs/test-utils/issues/425
   it('works with router-link and mount', () => {
     const wrapper = mount(Issue425, options)
-    expect(wrapper.getComponent<DefineComponent>('.link').props('to')).toEqual({
+    expect(
+      wrapper.getComponent<typeof RouterLinkStub>('.link').props('to')
+    ).toEqual({
       name
     })
   })
@@ -78,7 +80,9 @@ describe('getComponent', () => {
   // https://github.com/vuejs/test-utils/issues/425
   it('works with router-link and shallowMount', () => {
     const wrapper = shallowMount(Issue425, options)
-    expect(wrapper.getComponent<DefineComponent>('.link').props('to')).toEqual({
+    expect(
+      wrapper.getComponent<typeof RouterLinkStub>('.link').props('to')
+    ).toEqual({
       name
     })
   })
diff --git a/tests/props.spec.ts b/tests/props.spec.ts
index 690b7700e..503a959f9 100644
--- a/tests/props.spec.ts
+++ b/tests/props.spec.ts
@@ -20,6 +20,7 @@ describe('props', () => {
 
   it('returns undefined if props does not exist', () => {
     const wrapper = mount(WithProps, { props: { msg: 'ABC' } })
+    // @ts-expect-error :: non-existent prop
     expect(wrapper.props('foo')).toEqual(undefined)
   })