Fix setData/data through shallowMount/mount for components using setup()

src/createInstance.ts

@@ -9,7 +9,8 @@ import {
   ComponentOptions,
   ConcreteComponent,
   DefineComponent,
-  transformVNodeArgs
+  transformVNodeArgs,
+  proxyRefs
 } from 'vue'
 
 import { MountingOptions, Slot } from './types'
@@ -20,6 +21,7 @@ import {
   isObject,
   isObjectComponent,
   isScriptSetup,
+  mergeDeep,
   mergeGlobalProperties
 } from './utils'
 import { processSlot } from './utils/compileSlots'
@@ -153,11 +155,22 @@ export function createInstance(
     if (isObjectComponent(originalComponent)) {
       // component is guaranteed to be the same type as originalComponent
       const objectComponent = component as ComponentOptions
-      const originalDataFn = originalComponent.data || (() => ({}))
-      objectComponent.data = (vm) => ({
-        ...originalDataFn.call(vm, vm),
-        ...providedData
-      })
+      if (originalComponent.data) {
+        const originalDataFn = originalComponent.data
+        objectComponent.data = (vm) => ({
+          ...originalDataFn.call(vm, vm),
+          ...providedData
+        })
+      } else if (objectComponent.setup) {
+        const originalSetupFn = objectComponent.setup
+        objectComponent.setup = function (a, b) {
+          const data = originalSetupFn(a, b)
+          mergeDeep(proxyRefs(data), providedData)
+          return data
+        }
+      } else {
+        objectComponent.data = () => ({ ...providedData })
+      }
     } else {
       throw new Error(
         'data() option is not supported on functional and class components'

src/vueWrapper.ts

@@ -1,4 +1,4 @@
-import { nextTick, App, ComponentPublicInstance, VNode } from 'vue'
+import { nextTick, App, ComponentPublicInstance, VNode, proxyRefs } from 'vue'
 
 import { config } from './config'
 import domEvents from './constants/dom-events'
@@ -247,7 +247,29 @@ export class VueWrapper<
   }
 
   setData(data: Record<string, unknown>): Promise<void> {
-    mergeDeep(this.componentVM.$data, data)
+    /*
+    Depending on how the component was defined, data can live in different places 
+    Vue sets some default placeholder in all the locations however, so we cannot just check
+    if the data object exists or not.
+    When using <script setup>, data lives in the setupState object, which is then marked with __isScriptSetup
+    When using the setup() function, data lives in the setupState object, but is not marked with __isScriptSetup
+    When using the object api, data lives in the data object, proxied through $data, HOWEVER
+    the setupState object will also exist, and be frozen.
+    */
+    // @ts-ignore
+    if (this.componentVM.$.setupState.__isScriptSetup) {
+      // data from <script setup>
+      // @ts-ignore
+      mergeDeep(this.componentVM.$.setupState, data)
+      // @ts-ignore
+    } else if (!Object.isFrozen(this.componentVM.$.setupState)) {
+      // data from setup() function when using the object api
+      // @ts-ignore
+      mergeDeep(proxyRefs(this.componentVM.$.setupState), data)
+    } else {
+      // data when using data: {...} in the object api
+      mergeDeep(this.componentVM.$data, data)
+    }
     return nextTick()
   }
 

tests/setData.spec.ts

@@ -1,7 +1,8 @@
 import { describe, expect, it, vi } from 'vitest'
 import { defineComponent, ref } from 'vue'
 
-import { mount } from '../src'
+import { mount, shallowMount } from '../src'
+import ScriptSetup from './components/ScriptSetup.vue'
 
 describe('setData', () => {
   it('sets component data', async () => {
@@ -111,20 +112,6 @@ describe('setData', () => {
     )
   })
 
-  it('does not modify composition API setup data', async () => {
-    const Component = defineComponent({
-      template: `<div>Count is: {{ count }}</div>`,
-      setup: () => ({ count: ref(1) })
-    })
-    const wrapper = mount(Component)
-
-    expect(wrapper.html()).toContain('Count is: 1')
-
-    expect(() => wrapper.setData({ count: 2 })).toThrowError(
-      'Cannot add property count'
-    )
-  })
-
   // https://github.com/vuejs/test-utils/issues/538
   it('updates data set via data mounting option using setData', async () => {
     const Comp = defineComponent({
@@ -266,4 +253,85 @@ describe('setData', () => {
       secondArray: [3, 4]
     })
   })
+
+  it('updates initial data on a component using <script setup>', async () => {
+    const wrapper = shallowMount(ScriptSetup, {
+      data() {
+        return {
+          count: 10
+        }
+      }
+    })
+    expect(wrapper.html()).toContain('10')
+  })
+
+  it('updates nested data returned from setup()', async () => {
+    const Component = {
+      template: `<div>{{ nested.property }}</div>`,
+      setup: () => ({
+        nested: ref({
+          property: 'initial value'
+        })
+      })
+    }
+
+    const wrapper = mount(Component)
+
+    expect(wrapper.html()).toContain('initial value')
+
+    await wrapper.setData({ nested: { property: 'updated value' } })
+
+    expect(wrapper.html()).toContain('updated value')
+  })
+
+  it('updates data on a component using <script setup>', async () => {
+    const wrapper = shallowMount(ScriptSetup)
+    await wrapper.setData({ count: 20 })
+    expect(wrapper.html()).toContain('20')
+  })
+
+  it('updates data on a component using setup()', async () => {
+    const Component = {
+      template: `<div><div v-if="show" id="show">Show</div></div>`,
+      setup() {
+        return {
+          show: ref(false)
+        }
+      }
+    }
+
+    const wrapper = mount(Component)
+
+    expect(wrapper.find('#show').exists()).toBe(false)
+
+    await wrapper.setData({ show: true })
+
+    expect(wrapper.find('#show').exists()).toBe(true)
+  })
+
+  it('correctly sets initial array refs data when using <script setup>', async () => {
+    const Component = {
+      template: `<div><div v-for="item in items" :key="item">{{ item }}</div></div>`,
+      setup() {
+        const items = ref([])
+        return { items }
+      }
+    }
+
+    const wrapper = shallowMount(Component, {
+      data() {
+        return {
+          items: ['item1', 'item2']
+        }
+      }
+    })
+
+    expect(wrapper.html()).toContain('item1')
+    expect(wrapper.html()).toContain('item2')
+
+    await wrapper.setData({ items: ['item3', 'item4'] })
+
+    expect(wrapper.html()).toContain('item3')
+    expect(wrapper.html()).toContain('item4')
+  })
 })