fix(find): allow finding stubs by stub definition

src/mount.ts

@@ -277,7 +277,7 @@ export function mount(
   }
 
   addToDoNotStubComponents(component)
-  registerStub(originalComponent, component)
+  registerStub({ source: originalComponent, stub: component})
   const el = document.createElement('div')
 
   if (options?.attachTo) {

src/stubs.ts

@@ -26,17 +26,21 @@ interface StubOptions {
   renderStubDefaultSlot?: boolean
 }
 
-const stubsMap: WeakMap<ConcreteComponent, ConcreteComponent> = new WeakMap()
-export const registerStub = (
-  source: ConcreteComponent,
-  stub: ConcreteComponent
+const stubsMap: WeakMap<ConcreteComponent, { source: ConcreteComponent, originalStub?: ConcreteComponent }> = new WeakMap()
+export const registerStub = ({ source, stub, originalStub} :
+  { source: ConcreteComponent, stub: ConcreteComponent, originalStub?: ConcreteComponent }
 ) => {
-  stubsMap.set(stub, source)
+  stubsMap.set(stub, { source, originalStub })
 }
 
 export const getOriginalVNodeTypeFromStub = (
   type: ConcreteComponent
-): VNodeTypes | undefined => stubsMap.get(type)
+): VNodeTypes | undefined => stubsMap.get(type)?.source
+
+export const getOriginalStubFromSpecializedStub = (
+  type: ConcreteComponent
+): VNodeTypes | undefined => stubsMap.get(type)?.originalStub
+
 
 const doNotStubComponents: WeakSet<ConcreteComponent> = new WeakSet()
 const shouldNotStub = (type: ConcreteComponent) => doNotStubComponents.has(type)
@@ -151,7 +155,6 @@ function createStubOnceForType(
   }
 
   const stub = factoryFn()
-  registerStub(type, stub)
   cache.set(type, stub)
   return stub
 }
@@ -253,7 +256,7 @@ export function stubComponents(
           () => specializedStubComponent
         )
         specializedStub.props = unwrappedStub.props
-        registerStub(type, specializedStub)
+        registerStub({ source: type, stub: specializedStub, originalStub: stub })
         // pass the props and children, for advanced stubbing
         return [specializedStub, props, children, patchFlag, dynamicProps]
       }
@@ -274,14 +277,15 @@ export function stubComponents(
           throw new Error('Attempted to stub a non-component')
         }
 
-        const propsDeclaration = type.props || {}
+        const propsDeclaration = unwrapLegacyVueExtendComponent(type).props || {}
         const newStub = createStubOnce(type, () =>
           createStub({
             name: stubName,
             propsDeclaration,
             renderStubDefaultSlot
           })
         )
+        registerStub({ source: type, stub: newStub })
         return [newStub, props, children, patchFlag, dynamicProps]
       }
     }

src/utils/find.ts

@@ -2,10 +2,11 @@ import {
   ComponentInternalInstance,
   VNode,
   VNodeArrayChildren,
-  VNodeNormalizedChildren
+  VNodeNormalizedChildren,
+  VNodeTypes
 } from 'vue'
 import { FindAllComponentsSelector } from '../types'
-import { getOriginalVNodeTypeFromStub } from '../stubs'
+import { getOriginalStubFromSpecializedStub, getOriginalVNodeTypeFromStub } from '../stubs'
 import { isComponent } from '../utils'
 import { matchName } from './matchName'
 import { unwrapLegacyVueExtendComponent } from './vueCompatSupport'
@@ -28,33 +29,29 @@ export function matches(
   const nodeType = node.type
   if (!isComponent(nodeType)) return false
 
-  if (node.type === selector) {
-    return true
-  }
-
   if (typeof selector === 'string') {
     return node.el?.matches?.(selector)
   }
 
-  if (
-    typeof selector === 'object' &&
-    getOriginalVNodeTypeFromStub(nodeType) === selector
-  ) {
-    // we are looking at stub of this exact component
-    return true
-  }
+  // When we're using stubs we want user to be able to
+  // find stubbed components both by original component
+  // or stub definition. That's why we are trying to
+  // extract original component and also stub, which was
+  // used to create specialized stub for render
 
-  let componentName: string | undefined
-  if ('displayName' in nodeType) {
-    // match functional components
-    componentName = nodeType.displayName
-  }
+  const nodeTypeCandidates: VNodeTypes[] = [
+    nodeType,
+    getOriginalVNodeTypeFromStub(nodeType),
+    getOriginalStubFromSpecializedStub(nodeType)
+  ].filter(Boolean) as VNodeTypes[]
 
-  if (!componentName && 'name' in nodeType) {
-    // match normal component definitions
-    componentName = nodeType.name
+  if (nodeTypeCandidates.includes(selector)) {
+    return true;
   }
 
+  let componentName: string | undefined
+  componentName = nodeType.displayName || nodeType.name
+
   let selectorName = selector.name
 
   // the component and selector both have a name