fix: wrapper.element with nested multiple roots

src/vueWrapper.ts

@@ -21,6 +21,7 @@ import {
   registerFactory,
   WrapperType
 } from './wrapperFactory'
+import { VNode } from '@vue/runtime-core'
 
 export class VueWrapper<
   T extends Omit<
@@ -66,8 +67,40 @@ export class VueWrapper<
   }
 
   private get hasMultipleRoots(): boolean {
-    // if the subtree is an array of children, we have multiple root nodes
-    return this.vm.$.subTree.shapeFlag === ShapeFlags.ARRAY_CHILDREN
+    // Recursive check subtree for nested root elements
+    // <template>
+    //   <WithMultipleRoots />
+    // </template>
+    const checkTree = (subTree: VNode): boolean => {
+      // if the subtree is an array of children, we have multiple root nodes
+      if (subTree.shapeFlag === ShapeFlags.ARRAY_CHILDREN) return true
+
+      if (subTree.shapeFlag & ShapeFlags.STATEFUL_COMPONENT) {
+        // Component has multiple children or slot with multiple children
+        if (
+          subTree.shapeFlag & ShapeFlags.ARRAY_CHILDREN ||
+          subTree.shapeFlag & ShapeFlags.SLOTS_CHILDREN
+        ) {
+          return true
+        }
+
+        if (subTree.component?.subTree) {
+          return checkTree(subTree.component.subTree)
+        }
+      } else if (subTree.shapeFlag & ShapeFlags.FUNCTIONAL_COMPONENT) {
+        if (subTree.shapeFlag & ShapeFlags.ARRAY_CHILDREN) {
+          return true
+        }
+
+        if (subTree.component?.subTree) {
+          return checkTree(subTree.component.subTree)
+        }
+      }
+
+      return false
+    }
+
+    return checkTree(this.vm.$.subTree)
   }
 
   protected getRootNodes(): VueNode[] {

tests/element.spec.ts

@@ -2,6 +2,15 @@ import { defineComponent, h } from 'vue'
 
 import { mount } from '../src'
 
+const MultiRootText = defineComponent({
+  render: () => [h('div', {}, 'foo'), h('div', {}, 'bar'), h('div', {}, 'baz')]
+})
+const ReturnSlot = defineComponent({
+  render() {
+    return this.$slots.default!({})
+  }
+})
+
 describe('element', () => {
   it('returns element when mounting single root node', () => {
     const Component = defineComponent({
@@ -16,14 +25,66 @@ describe('element', () => {
   })
 
   it('returns the VTU root element when mounting multiple root nodes', () => {
-    const Component = defineComponent({
-      render() {
-        return [h('div', {}, 'foo'), h('div', {}, 'bar'), h('div', {}, 'baz')]
-      }
+    const wrapper = mount(MultiRootText)
+
+    expect(wrapper.element.innerHTML).toBe(
+      '<div>foo</div><div>bar</div><div>baz</div>'
+    )
+  })
+
+  it('returns correct element for root component with multiple roots', () => {
+    const Parent = defineComponent({
+      components: { MultiRootText },
+      template: '<MultiRootText/>'
     })
 
-    const wrapper = mount(Component)
+    const wrapper = mount(Parent)
+
+    expect(wrapper.findComponent(MultiRootText).text()).toBe('foobarbaz')
+    expect(wrapper.text()).toBe('foobarbaz')
+  })
+
+  it('returns correct element for root slot', () => {
+    const Parent = defineComponent({
+      components: { ReturnSlot },
+      template: `
+    <ReturnSlot>
+      <div>foo</div>
+      <div>bar</div>
+      <div>baz</div>
+    </ReturnSlot>`
+    })
+
+    const wrapper = mount(Parent)
+    expect(wrapper.element.innerHTML).toBe(
+      '<div>foo</div><div>bar</div><div>baz</div>'
+    )
+  })
+
+  it('should return element for multi root functional component', () => {
+    const Foo = () => h(MultiRootText)
+    const wrapper = mount(Foo)
+
+    expect(wrapper.element.innerHTML).toBe(
+      '<div>foo</div><div>bar</div><div>baz</div>'
+    )
+  })
+
+  it('returns correct element for root slot with functional component', () => {
+    const wrapper = mount(() =>
+      h(ReturnSlot, {}, () => [
+        h('div', {}, 'foo'),
+        h('div', {}, 'bar'),
+        h('div', {}, 'baz')
+      ])
+    )
+    expect(wrapper.element.innerHTML).toBe(
+      '<div>foo</div><div>bar</div><div>baz</div>'
+    )
+  })
 
+  it('returns correct element for root slot with nested component', () => {
+    const wrapper = mount(() => h(ReturnSlot, {}, () => h(MultiRootText)))
     expect(wrapper.element.innerHTML).toBe(
       '<div>foo</div><div>bar</div><div>baz</div>'
     )

tests/text.spec.ts

@@ -2,6 +2,15 @@ import { defineComponent, h } from 'vue'
 
 import { mount } from '../src'
 
+const MultiRootText = defineComponent({
+  render: () => [h('div', {}, 'foo'), h('div', {}, 'bar'), h('div', {}, 'baz')]
+})
+const ReturnSlot = defineComponent({
+  render() {
+    return this.$slots.default!({})
+  }
+})
+
 describe('text', () => {
   it('returns text when mounting single root node', () => {
     const Component = defineComponent({
@@ -16,13 +25,7 @@ describe('text', () => {
   })
 
   it('returns text when mounting multiple root nodes', () => {
-    const Component = defineComponent({
-      render() {
-        return [h('div', {}, 'foo'), h('div', {}, 'bar'), h('div', {}, 'baz')]
-      }
-    })
-
-    const wrapper = mount(Component)
+    const wrapper = mount(MultiRootText)
 
     expect(wrapper.text()).toBe('foobarbaz')
   })
@@ -39,4 +42,54 @@ describe('text', () => {
 
     expect(wrapper.text()).toBe('')
   })
+
+  it('returns correct text for root component with multiple roots', () => {
+    const Parent = defineComponent({
+      components: { MultiRootText },
+      template: '<MultiRootText/>'
+    })
+
+    const wrapper = mount(Parent)
+
+    expect(wrapper.findComponent(MultiRootText).text()).toBe('foobarbaz')
+    expect(wrapper.text()).toBe('foobarbaz')
+  })
+
+  it('returns correct text for root slot', () => {
+    const Parent = defineComponent({
+      components: { ReturnSlot },
+      template: `
+    <ReturnSlot>
+      <div>foo</div>
+      <div>bar</div>
+      <div>baz</div>
+    </ReturnSlot>`
+    })
+
+    const wrapper = mount(Parent)
+    expect(wrapper.text()).toBe('foobarbaz')
+  })
+
+  it('should return text for multi root functional component', () => {
+    const Foo = () => h(MultiRootText)
+    const wrapper = mount(Foo)
+
+    expect(wrapper.text()).toBe('foobarbaz')
+  })
+
+  it('returns correct text for root slot with functional component', () => {
+    const wrapper = mount(() =>
+      h(ReturnSlot, {}, () => [
+        h('div', {}, 'foo'),
+        h('div', {}, 'bar'),
+        h('div', {}, 'baz')
+      ])
+    )
+    expect(wrapper.text()).toBe('foobarbaz')
+  })
+
+  it('returns correct text for root slot with nested component', () => {
+    const wrapper = mount(() => h(ReturnSlot, {}, () => h(MultiRootText)))
+    expect(wrapper.text()).toBe('foobarbaz')
+  })
 })