From ad506d8d5ae7559f265d8567bf53ba363f4c0cf6 Mon Sep 17 00:00:00 2001
From: Illya Klymov <xanf@xanf.me>
Date: Mon, 16 May 2022 17:15:34 +0300
Subject: [PATCH] fix(wrapper): fix wrapper.element for component with slot

---
 src/vueWrapper.ts     | 16 ++++------------
 tests/element.spec.ts | 15 +++++++++++++++
 2 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/src/vueWrapper.ts b/src/vueWrapper.ts
index b6a6d86e5..46090ea05 100644
--- a/src/vueWrapper.ts
+++ b/src/vueWrapper.ts
@@ -75,19 +75,11 @@ export class VueWrapper<
       // 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) {
+      if (
+        subTree.shapeFlag & ShapeFlags.STATEFUL_COMPONENT ||
+        subTree.shapeFlag & ShapeFlags.FUNCTIONAL_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
         }
diff --git a/tests/element.spec.ts b/tests/element.spec.ts
index e1c6e71b2..2e765afc5 100644
--- a/tests/element.spec.ts
+++ b/tests/element.spec.ts
@@ -24,6 +24,21 @@ describe('element', () => {
     expect(wrapper.element.nodeName).toBe('DIV')
   })
 
+  it('returns the VTU root element when mounting component with slot', () => {
+    const NestedComponent = defineComponent({
+      template: `<h1><slot></slot></h1>`
+    })
+
+    const RootComponent = defineComponent({
+      components: { NestedComponent },
+      template: `<nested-component>test</nested-component>`
+    })
+    const wrapper = mount(RootComponent)
+
+    expect(wrapper.element.tagName).toBe('H1')
+    expect(wrapper.html()).toBe('<h1>test</h1>')
+  })
+
   it('returns the VTU root element when mounting multiple root nodes', () => {
     const wrapper = mount(MultiRootText)