fix(find): find element inside suspense with multi root elements (#1397)

Author: freakzlike

src/utils/getRootNodes.ts

@@ -3,11 +3,13 @@ import { isNotNullOrUndefined } from '../utils'
 import { VNode, VNodeArrayChildren } from 'vue'
 
 export function getRootNodes(vnode: VNode): Node[] {
-  if (vnode.shapeFlag & (ShapeFlags.ELEMENT | ShapeFlags.SUSPENSE)) {
+  if (vnode.shapeFlag & ShapeFlags.ELEMENT) {
     return [vnode.el as Node]
   } else if (vnode.shapeFlag & ShapeFlags.COMPONENT) {
     const { subTree } = vnode.component!
     return getRootNodes(subTree)
+  } else if (vnode.shapeFlag & ShapeFlags.SUSPENSE) {
+    return getRootNodes(vnode.suspense!.activeBranch!)
   } else if (
     vnode.shapeFlag &
     (ShapeFlags.TEXT_CHILDREN | ShapeFlags.TELEPORT)

tests/features/suspense.spec.ts

@@ -1,5 +1,6 @@
 import SuspenseComponent from '../components/Suspense.vue'
 import { mount, flushPromises } from '../../src'
+import { defineComponent } from 'vue'
 
 let mockShouldError = false
 jest.mock('../utils', () => ({
@@ -32,4 +33,30 @@ describe('suspense', () => {
 
     expect(wrapper.html()).toContain('Error!')
   })
+
+  test('returns the element if it is a multi root element inside Suspense', () => {
+    const Async = defineComponent({
+      template: '<h1>Hello</h1><span id="my-span">There</span>'
+    })
+    const Component = defineComponent({
+      components: { Async },
+      template: '<Suspense><Async/></Suspense>'
+    })
+
+    const wrapper = mount(Component)
+    expect(wrapper.get('#my-span')).not.toBeNull()
+  })
+
+  test('returns the element if it is a root element inside Suspense', () => {
+    const Async = defineComponent({
+      template: '<div><h1>Hello</h1><span id="my-span">There</span></div>'
+    })
+    const Component = defineComponent({
+      components: { Async },
+      template: '<Suspense><Async/></Suspense>'
+    })
+
+    const wrapper = mount(Component)
+    expect(wrapper.get('#my-span')).not.toBeNull()
+  })
 })