fix(find): implement proper find for multi-root components (#1116)

src/baseWrapper.ts

@@ -1,4 +1,4 @@
-import { textContent } from './utils'
+import { isNotNullOrUndefined, textContent } from './utils'
 import type { TriggerOptions } from './createDomEvent'
 import {
   ComponentInternalInstance,
@@ -14,7 +14,8 @@ import {
   FindAllComponentsSelector,
   FindComponentSelector,
   NameSelector,
-  RefSelector
+  RefSelector,
+  VueNode
 } from './types'
 import WrapperLike from './interfaces/wrapperLike'
 import { find, matches } from './utils/find'
@@ -27,9 +28,8 @@ import { createDOMWrapper, createVueWrapper } from './wrapperFactory'
 export default abstract class BaseWrapper<ElementType extends Node>
   implements WrapperLike
 {
-  private readonly wrapperElement: ElementType & {
-    __vueParentComponent?: ComponentInternalInstance
-  }
+  protected readonly wrapperElement: VueNode<ElementType>
+  protected abstract getRootNodes(): VueNode[]
 
   get element() {
     return this.wrapperElement
@@ -48,11 +48,6 @@ export default abstract class BaseWrapper<ElementType extends Node>
   find<T extends Element = Element>(selector: string): DOMWrapper<T>
   find<T extends Node = Node>(selector: string | RefSelector): DOMWrapper<T>
   find(selector: string | RefSelector): DOMWrapper<Node> {
-    // allow finding the root element
-    if (!isElement(this.element)) {
-      return createWrapperError('DOMWrapper')
-    }
-
     if (typeof selector === 'object' && 'ref' in selector) {
       const currentComponent = this.getCurrentComponent()
       if (!currentComponent) {
@@ -68,12 +63,24 @@ export default abstract class BaseWrapper<ElementType extends Node>
       }
     }
 
-    if (this.element.matches(selector)) {
-      return createDOMWrapper(this.element)
+    const elementRootNodes = this.getRootNodes().filter(
+      (node): node is Element => node instanceof Element
+    )
+    if (elementRootNodes.length === 0) {
+      return createWrapperError('DOMWrapper')
     }
-    const result = this.element.querySelector(selector)
-    if (result) {
-      return createDOMWrapper(result)
+    const matchingRootNode = elementRootNodes.find((node) =>
+      node.matches(selector)
+    )
+    if (matchingRootNode) {
+      return createDOMWrapper(matchingRootNode)
+    }
+
+    const result = elementRootNodes
+      .map((node) => node.querySelector(selector))
+      .filter(isNotNullOrUndefined)
+    if (result.length > 0) {
+      return createDOMWrapper(result[0])
     }
 
     return createWrapperError('DOMWrapper')

src/domWrapper.ts

@@ -14,6 +14,10 @@ export class DOMWrapper<NodeType extends Node> extends BaseWrapper<NodeType> {
     config.plugins.DOMWrapper.extend(this)
   }
 
+  getRootNodes() {
+    return [this.wrapperElement]
+  }
+
   getCurrentComponent() {
     return this.element.__vueParentComponent
   }

src/types.ts

@@ -6,7 +6,8 @@ import {
   AppConfig,
   VNode,
   VNodeProps,
-  FunctionalComponent
+  FunctionalComponent,
+  ComponentInternalInstance
 } from 'vue'
 
 export interface RefSelector {
@@ -139,6 +140,11 @@ export type GlobalMountOptions = {
   renderStubDefaultSlot?: boolean
 }
 
-export type VueElement = Element & { __vue_app__?: any }
+export type VueNode<T extends Node = Node> = T & {
+  __vue_app__?: any
+  __vueParentComponent?: ComponentInternalInstance
+}
+
+export type VueElement = VueNode<Element>
 
 export type DefinedComponent = new (...args: any[]) => any

src/utils.ts

@@ -108,6 +108,12 @@ export function hasOwnProperty<O extends {}, P extends PropertyKey>(
   return obj.hasOwnProperty(prop)
 }
 
+export function isNotNullOrUndefined<T extends {}>(
+  obj: T | null | undefined
+): obj is T {
+  return Boolean(obj)
+}
+
 export function isRefSelector(
   selector: string | RefSelector
 ): selector is RefSelector {

src/utils/getRootNodes.ts

@@ -0,0 +1,35 @@
+import { ShapeFlags } from '@vue/shared'
+import { isNotNullOrUndefined } from '../utils'
+import { VNode, VNodeArrayChildren } from 'vue'
+
+export function getRootNodes(vnode: VNode): Node[] {
+  if (vnode.shapeFlag & (ShapeFlags.ELEMENT | ShapeFlags.SUSPENSE)) {
+    return [vnode.el as Node]
+  } else if (vnode.shapeFlag & ShapeFlags.COMPONENT) {
+    const { subTree } = vnode.component!
+    return getRootNodes(subTree)
+  } else if (vnode.shapeFlag & ShapeFlags.TEXT_CHILDREN) {
+    // static node optimization, subTree.children will be static string and will not help us
+    const result = [vnode.el as Node]
+    if (vnode.anchor) {
+      let currentNode: Node | null = result[0].nextSibling
+      while (currentNode && currentNode.previousSibling !== vnode.anchor) {
+        result.push(currentNode)
+        currentNode = currentNode.nextSibling
+      }
+    }
+    return result
+  } else if (vnode.shapeFlag & ShapeFlags.ARRAY_CHILDREN) {
+    const children = (vnode.children as unknown as VNodeArrayChildren).flat(
+      Infinity
+    ) as VNode[]
+
+    return children
+      .flatMap((vnode) => getRootNodes(vnode))
+      .filter(isNotNullOrUndefined)
+  }
+  // Missing cases which do not need special handling:
+  // ShapeFlags.SLOTS_CHILDREN comes with ShapeFlags.ELEMENT
+  // ShapeFlags.TELEPORT comes with ShapeFlags.TEXT_CHILDREN
+  return []
+}

src/vueWrapper.ts

@@ -10,8 +10,9 @@ import pretty from 'pretty'
 
 import { config } from './config'
 import domEvents from './constants/dom-events'
-import { VueElement } from './types'
+import { VueElement, VueNode } from './types'
 import { mergeDeep } from './utils'
+import { getRootNodes } from './utils/getRootNodes'
 import { emitted, recordEvent } from './emit'
 import BaseWrapper from './baseWrapper'
 import {
@@ -68,6 +69,10 @@ export class VueWrapper<
     return this.vm.$.subTree.shapeFlag === ShapeFlags.ARRAY_CHILDREN
   }
 
+  protected getRootNodes(): VueNode[] {
+    return getRootNodes(this.vm.$.vnode)
+  }
+
   private get parentElement(): VueElement {
     return this.vm.$el.parentElement
   }

tests/find.spec.ts

@@ -55,18 +55,7 @@ describe('find', () => {
     })
   })
 
-  it('find using multiple root nodes', () => {
-    const Component = defineComponent({
-      render() {
-        return [h('div', 'text'), h('span', { id: 'my-span' })]
-      }
-    })
-
-    const wrapper = mount(Component)
-    expect(wrapper.find('#my-span').exists()).toBe(true)
-  })
-
-  test('find using current node after findAllComponents', () => {
+  it('find using current node after findAllComponents', () => {
     const ComponentB = defineComponent({
       name: 'ComponentB',
       template: '<div><slot></slot></div>'
@@ -126,14 +115,14 @@ describe('find', () => {
     expect(wrapper.find('.foo').find('.bar').exists()).toBe(true)
   })
 
-  test('works with suspense', async () => {
+  it('works with suspense', async () => {
     const wrapper = mount(SuspenseComponent)
 
     expect(wrapper.html()).toContain('Fallback content')
     expect(wrapper.find('div').exists()).toBeTruthy()
   })
 
-  test('can wrap `find` in an async function', async () => {
+  it('can wrap `find` in an async function', async () => {
     async function findAfterNextTick(
       wrapper: VueWrapper<any>,
       selector: string
@@ -152,7 +141,7 @@ describe('find', () => {
     expect(foundElement.exists()).toBeFalsy()
   })
 
-  test('handle empty root node', () => {
+  it('handle empty root node', () => {
     const EmptyTestComponent = {
       name: 'EmptyTestComponent',
       render: () => null
@@ -170,7 +159,7 @@ describe('find', () => {
 })
 
 describe('findAll', () => {
-  test('findAll using single root node', () => {
+  it('findAll using single root node', () => {
     const Component = defineComponent({
       render() {
         return h('div', {}, [
@@ -184,7 +173,7 @@ describe('findAll', () => {
     expect(wrapper.findAll('.span')).toHaveLength(2)
   })
 
-  test('findAll using multiple root nodes', () => {
+  it('findAll using multiple root nodes', () => {
     const Component = defineComponent({
       render() {
         return [
@@ -198,7 +187,7 @@ describe('findAll', () => {
     expect(wrapper.findAll('.span')).toHaveLength(2)
   })
 
-  test('findAll using current node after findAllComponents', () => {
+  it('findAll using current node after findAllComponents', () => {
     const ComponentB = defineComponent({
       name: 'ComponentB',
       template: '<div><slot></slot></div>'
@@ -224,14 +213,14 @@ describe('findAll', () => {
     expect(lastCompB.findAll('input')[0].element.value).toBe('2')
   })
 
-  test('works with suspense', async () => {
+  it('works with suspense', async () => {
     const wrapper = mount(SuspenseComponent)
 
     expect(wrapper.html()).toContain('Fallback content')
     expect(wrapper.findAll('div')).toBeTruthy()
   })
 
-  test('chaining finds compiles successfully', () => {
+  it('chaining finds compiles successfully', () => {
     const Bar = {
       render() {
         return h('span', { id: 'bar' })
@@ -247,7 +236,7 @@ describe('findAll', () => {
     expect(wrapper.find('#foo').find('#bar').exists()).toBe(true)
   })
 
-  test('handle empty/comment root node', () => {
+  it('handle empty/comment root node', () => {
     const EmptyTestComponent = {
       name: 'EmptyTestComponent',
       render: () => null
@@ -262,4 +251,89 @@ describe('findAll', () => {
     const etc = wrapper.findComponent({ name: 'EmptyTestComponent' })
     expect(etc.findAll('p')).toHaveLength(0)
   })
+
+  describe('with multiple root nodes', () => {
+    const MultipleRootComponentWithRenderFn = defineComponent({
+      render() {
+        return [
+          h(
+            'div',
+            { class: 'root1' },
+            h('div', { class: 'target1' }, 'target1')
+          ),
+          h(
+            'div',
+            { class: 'root2' },
+            h('div', { class: 'target2' }, 'target2')
+          ),
+          h(
+            'div',
+            { class: 'root3' },
+            h('div', { class: 'target3' }, 'target3')
+          )
+        ]
+      }
+    })
+
+    const MultipleRootComponentWithTemplate = defineComponent({
+      template: [
+        '<div class="root1"><div class="target1">target1</div></div>',
+        '<div class="root2"><div class="target2">target2</div></div>',
+        '<div class="root3"><div class="target3">target3</div></div>'
+      ].join('\n')
+    })
+
+    const WrapperComponent = defineComponent({
+      components: {
+        MultipleRootComponentWithTemplate,
+        MultipleRootComponentWithRenderFn
+      },
+      template: [
+        '<div><multiple-root-component-with-template /></div>',
+        '<div><multiple-root-component-with-render-fn /></div>'
+      ].join('\n')
+    })
+    it('find one of root nodes', () => {
+      const Component = defineComponent({
+        render() {
+          return [h('div', 'text'), h('span', { id: 'my-span' })]
+        }
+      })
+
+      const wrapper = mount(Component)
+      expect(wrapper.find('#my-span').exists()).toBe(true)
+    })
+
+    it('finds second root node when component is not mount root', () => {
+      const wrapper = mount(WrapperComponent)
+      expect(
+        wrapper
+          .findComponent(MultipleRootComponentWithRenderFn)
+          .find('.root2')
+          .exists()
+      ).toBe(true)
+      expect(
+        wrapper
+          .findComponent(MultipleRootComponentWithTemplate)
+          .find('.root2')
+          .exists()
+      ).toBe(true)
+    })
+
+    it('finds contents of second root node when component is not mount root', () => {
+      const wrapper = mount(WrapperComponent)
+      expect(
+        wrapper
+          .findComponent(MultipleRootComponentWithTemplate)
+          .find('.target2')
+          .exists()
+      ).toBe(true)
+      expect(
+        wrapper
+          .findComponent(MultipleRootComponentWithRenderFn)
+          .find('.target2')
+          .exists()
+      ).toBe(true)
+    })
+  })
 })