From d96d5bb8e12801312369bef17c0a9a973bfddf7c Mon Sep 17 00:00:00 2001
From: Illya Klymov <xanf@xanf.me>
Date: Mon, 10 Oct 2022 00:12:34 +0300
Subject: [PATCH] feat(stubs): Allow to stub directives

* this adds stubbing directives functionality to @vue/test-utils
---
 docs/guide/advanced/stubs-shallow-mount.md    |  76 ++++++++++++-
 src/mount.ts                                  |  17 ++-
 src/stubs.ts                                  |  32 ++++++
 src/types.ts                                  |   2 +-
 src/utils.ts                                  |  46 +++++++-
 src/utils/find.ts                             |  15 +--
 .../stubComponentsTransformer.ts              |  60 +++-------
 .../stubDirectivesTransformer.ts              |  46 ++++++++
 src/vnodeTransformers/util.ts                 |  14 ++-
 tests/findComponent.spec.ts                   |  20 ++++
 tests/mountingOptions/global.stubs.spec.ts    | 104 +++++++++++++++++-
 11 files changed, 367 insertions(+), 65 deletions(-)
 create mode 100644 src/stubs.ts
 create mode 100644 src/vnodeTransformers/stubDirectivesTransformer.ts

diff --git a/docs/guide/advanced/stubs-shallow-mount.md b/docs/guide/advanced/stubs-shallow-mount.md
index dc43fe98c..d0aef7d3d 100644
--- a/docs/guide/advanced/stubs-shallow-mount.md
+++ b/docs/guide/advanced/stubs-shallow-mount.md
@@ -1,6 +1,6 @@
 # Stubs and Shallow Mount
 
-Vue Test Utils provides some advanced features for _stubbing_ components. A _stub_ is where you replace an existing implementation of a custom component with a dummy component that doesn't do anything at all, which can simplify an otherwise complex test. Let's see an example.
+Vue Test Utils provides some advanced features for _stubbing_ components and directives. A _stub_ is where you replace an existing implementation of a custom component or directive with a dummy one that doesn't do anything at all, which can simplify an otherwise complex test. Let's see an example.
 
 ## Stubbing a single child component
 
@@ -238,6 +238,78 @@ test('stubs async component with resolving', async () => {
 })
 ```
 
+## Stubbing a directive
+
+Sometimes directives do quite complex things, like perform a lot of DOM manipulation which might result in errors in your tests (due to JSDOM not resembling entire DOM behavior). A common example is tooltip directives from various libraries, which usually rely heavily on measuring DOM nodes position/sizes.
+
+In this example, we have another `<App>` that renders a message with tooltip
+
+```js
+// tooltip directive declared somewhere, named `Tooltip`
+
+const App = {
+  directives: {
+    Tooltip
+  },
+  template: '<h1 v-tooltip title="Welcome tooltip">Welcome to Vue.js 3</h1>'
+}
+```
+
+We do not want `Tooltip` directive code to be executed in this test, we just want to assert the message is rendered. In this case, we could use the `stubs`, which appears in the `global` mounting option passing `vTooltip`.
+
+```js
+test('stubs component with custom template', () => {
+  const wrapper = mount(App, {
+    global: {
+      stubs: {
+        vTooltip: true
+      }
+    }
+  })
+
+  console.log(wrapper.html())
+  // <h1>Welcome to Vue.js 3</h1><span></span>
+
+  expect(wrapper.html()).toContain('Welcome to Vue.js 3')
+})
+```
+
+::: tip
+Usage of `vCustomDirective` naming scheme to differentiate between components and directives is inspired by [same approach](https://vuejs.org/api/sfc-script-setup.html#using-custom-directives) used in `<script setup>`
+:::
+
+Sometimes, we need a part of directive functionality (usually because some code relies on it). Let's assume our directive adds `with-tooltip` CSS class when executed and this is important behavior for our code. In this case we can swap `true` with our mock directive implementation
+
+```js
+test('stubs component with custom template', () => {
+  const wrapper = mount(App, {
+    global: {
+      stubs: {
+        vTooltip: {
+          beforeMount(el: Element) {
+            console.log('directive called')
+            el.classList.add('with-tooltip')
+          }
+        }
+      }
+    }
+  })
+
+  // 'directive called' logged to console
+
+  console.log(wrapper.html())
+  // <h1>Welcome to Vue.js 3</h1><span></span>
+
+  expect(wrapper.classes('with-tooltip')).toBe(true)
+})
+```
+
+We've just swapped our directive implementation with own one!
+
+::: warning
+Stubbing directives won't work on functional components due to lack of directive name inside of [withDirectives](https://vuejs.org/api/render-function.html#withdirectives) function. Consider mocking directive module via your testing framework if you need to mock directive used in functional component
+:::
+
 ## Default Slots and `shallow`
 
 Since `shallow` stubs out all the content of a components, any `<slot>` won't get rendered when using `shallow`. While this is not a problem in most cases, there are some scenarios where this isn't ideal.
@@ -314,6 +386,6 @@ So regardless of which mounting method you choose, we suggest keeping these guid
 
 ## Conclusion
 
-- use `global.stubs` to replace a component with a dummy one to simplify your tests
+- use `global.stubs` to replace a component or directive with a dummy one to simplify your tests
 - use `shallow: true` (or `shallowMount`) to stub out all child components
 - use `config.renderStubDefaultSlot` to render the default `<slot>` for a stubbed component
diff --git a/src/mount.ts b/src/mount.ts
index 31ddda39d..f8f04a1c6 100644
--- a/src/mount.ts
+++ b/src/mount.ts
@@ -30,6 +30,8 @@ import {
 
 import { MountingOptions, Slot } from './types'
 import {
+  getComponentsFromStubs,
+  getDirectivesFromStubs,
   isFunctionalComponent,
   isObject,
   isObjectComponent,
@@ -41,15 +43,16 @@ import { attachEmitListener } from './emit'
 import { createVNodeTransformer } from './vnodeTransformers/util'
 import {
   createStubComponentsTransformer,
-  addToDoNotStubComponents,
-  registerStub
+  addToDoNotStubComponents
 } from './vnodeTransformers/stubComponentsTransformer'
+import { createStubDirectivesTransformer } from './vnodeTransformers/stubDirectivesTransformer'
 import {
   isLegacyFunctionalComponent,
   unwrapLegacyVueExtendComponent
 } from './utils/vueCompatSupport'
 import { trackInstance } from './utils/autoUnmount'
 import { createVueWrapper } from './wrapperFactory'
+import { registerStub } from './stubs'
 
 // NOTE this should come from `vue`
 const MOUNT_OPTIONS: Array<keyof MountingOptions<any>> = [
@@ -338,7 +341,10 @@ export function mount(
   }
 
   addToDoNotStubComponents(component)
+  // We've just replaced our component with it's copy
+  // Let's register it as a stub so user can find it
   registerStub({ source: originalComponent, stub: component })
+
   const el = document.createElement('div')
 
   if (options?.attachTo) {
@@ -532,9 +538,12 @@ export function mount(
     createVNodeTransformer({
       transformers: [
         createStubComponentsTransformer({
-          stubs: global.stubs,
+          stubs: getComponentsFromStubs(global.stubs),
           shallow: options?.shallow,
           renderStubDefaultSlot: global.renderStubDefaultSlot
+        }),
+        createStubDirectivesTransformer({
+          directives: getDirectivesFromStubs(global.stubs)
         })
       ]
     })
@@ -551,7 +560,7 @@ export function mount(
   // ref: https://github.com/vuejs/test-utils/issues/249
   // ref: https://github.com/vuejs/test-utils/issues/425
   if (global?.stubs) {
-    for (const name of Object.keys(global.stubs)) {
+    for (const name of Object.keys(getComponentsFromStubs(global.stubs))) {
       if (!app.component(name)) {
         app.component(name, { name })
       }
diff --git a/src/stubs.ts b/src/stubs.ts
new file mode 100644
index 000000000..a0b8d7cb2
--- /dev/null
+++ b/src/stubs.ts
@@ -0,0 +1,32 @@
+import { Component } from 'vue'
+
+// Stubbing occurs when in vnode transformer we're swapping
+// component vnode type due to stubbing either component
+// or directive on component
+
+// In order to be able to find components we need to track pairs
+// stub --> original component
+
+// Having this as global might feel unsafe at first point
+// One can assume that sharing stub map across mounts might
+// lead to false matches, however our vnode mappers always
+// produce new nodeTypes for each mount even if you're reusing
+// same stub, so we're safe and do not need to pass these stubs
+// for each mount operation
+const stubs: WeakMap<Component, Component> = new WeakMap()
+
+export function registerStub({
+  source,
+  stub
+}: {
+  source: Component
+  stub: Component
+}) {
+  stubs.set(stub, source)
+}
+
+export function getOriginalComponentFromStub(
+  stub: Component
+): Component | undefined {
+  return stubs.get(stub)
+}
diff --git a/src/types.ts b/src/types.ts
index 4bda84b1c..e332e7742 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -86,7 +86,7 @@ export interface MountingOptions<Props, Data = {}> {
   shallow?: boolean
 }
 
-export type Stub = boolean | Component
+export type Stub = boolean | Component | Directive
 export type Stubs = Record<string, Stub> | Array<string>
 export type GlobalMountOptions = {
   /**
diff --git a/src/utils.ts b/src/utils.ts
index 23c10ea61..c4a3af213 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -1,5 +1,11 @@
-import { GlobalMountOptions, RefSelector } from './types'
-import { ComponentOptions, ConcreteComponent, FunctionalComponent } from 'vue'
+import { GlobalMountOptions, RefSelector, Stub, Stubs } from './types'
+import {
+  Component,
+  ComponentOptions,
+  ConcreteComponent,
+  Directive,
+  FunctionalComponent
+} from 'vue'
 import { config } from './config'
 
 function mergeStubs(target: Record<string, any>, source: GlobalMountOptions) {
@@ -137,3 +143,39 @@ export function isRefSelector(
 ): selector is RefSelector {
   return typeof selector === 'object' && 'ref' in selector
 }
+
+export function convertStubsToRecord(stubs: Stubs) {
+  if (Array.isArray(stubs)) {
+    // ['Foo', 'Bar'] => { Foo: true, Bar: true }
+    return stubs.reduce((acc, current) => {
+      acc[current] = true
+      return acc
+    }, {} as Record<string, Stub>)
+  }
+
+  return stubs
+}
+
+const isDirectiveKey = (key: string) => key.match(/^v[A-Z].*/)
+
+export function getComponentsFromStubs(
+  stubs: Stubs
+): Record<string, Component | boolean> {
+  const normalizedStubs = convertStubsToRecord(stubs)
+
+  return Object.fromEntries(
+    Object.entries(normalizedStubs).filter(([key]) => !isDirectiveKey(key))
+  ) as Record<string, Component | boolean>
+}
+
+export function getDirectivesFromStubs(
+  stubs: Stubs
+): Record<string, Directive> {
+  const normalizedStubs = convertStubsToRecord(stubs)
+
+  return Object.fromEntries(
+    Object.entries(normalizedStubs)
+      .filter(([key]) => isDirectiveKey(key))
+      .map(([key, value]) => [key.substring(1), value])
+  ) as Record<string, Directive>
+}
diff --git a/src/utils/find.ts b/src/utils/find.ts
index ad7785cb1..b11a1035e 100644
--- a/src/utils/find.ts
+++ b/src/utils/find.ts
@@ -5,11 +5,8 @@ import {
   VNodeNormalizedChildren,
   VNodeTypes
 } from 'vue'
+import { getOriginalComponentFromStub } from '../stubs'
 import { FindAllComponentsSelector } from '../types'
-import {
-  getOriginalStubFromSpecializedStub,
-  getOriginalVNodeTypeFromStub
-} from '../vnodeTransformers/stubComponentsTransformer'
 import { isComponent } from '../utils'
 import { matchName } from './matchName'
 import { unwrapLegacyVueExtendComponent } from './vueCompatSupport'
@@ -45,11 +42,15 @@ export function matches(
 
   const nodeTypeCandidates: VNodeTypes[] = [
     nodeType,
-    getOriginalVNodeTypeFromStub(nodeType),
-    getOriginalStubFromSpecializedStub(nodeType)
+    getOriginalComponentFromStub(nodeType)
+    // getOriginalVNodeTypeFromStub(nodeType),
+    // getOriginalStubFromSpecializedStub(nodeType)
   ].filter(Boolean) as VNodeTypes[]
 
-  if (nodeTypeCandidates.includes(selector)) {
+  // our selector might be a stub itself
+  const target = getOriginalComponentFromStub(selector) ?? selector
+
+  if (nodeTypeCandidates.includes(target)) {
     return true
   }
 
diff --git a/src/vnodeTransformers/stubComponentsTransformer.ts b/src/vnodeTransformers/stubComponentsTransformer.ts
index 62affe12d..e80c3960d 100644
--- a/src/vnodeTransformers/stubComponentsTransformer.ts
+++ b/src/vnodeTransformers/stubComponentsTransformer.ts
@@ -10,18 +10,19 @@ import {
   ConcreteComponent,
   ComponentPropsOptions,
   ComponentObjectPropsOptions,
-  DefineComponent
+  DefineComponent,
+  Component
 } from 'vue'
 import { hyphenate } from '../utils/vueShared'
 import { matchName } from '../utils/matchName'
 import { isComponent, isFunctionalComponent } from '../utils'
 import { unwrapLegacyVueExtendComponent } from '../utils/vueCompatSupport'
-import { Stub, Stubs } from '../types'
 import {
   getComponentName,
   getComponentRegisteredName
 } from '../utils/componentName'
 import { config } from '../config'
+import { registerStub } from '../stubs'
 
 export type CustomCreateStub = (params: {
   name: string
@@ -34,30 +35,6 @@ interface StubOptions {
   renderStubDefaultSlot?: boolean
 }
 
-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, originalStub })
-}
-
-export const getOriginalVNodeTypeFromStub = (
-  type: ConcreteComponent
-): 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)
 export const addToDoNotStubComponents = (type: ConcreteComponent) =>
@@ -112,15 +89,10 @@ export const createStub = ({
   })
 }
 
-const resolveComponentStubByName = (componentName: string, stubs: Stubs) => {
-  if (Array.isArray(stubs) && stubs.length) {
-    // ['Foo', 'Bar'] => { Foo: true, Bar: true }
-    stubs = stubs.reduce((acc, current) => {
-      acc[current] = true
-      return acc
-    }, {} as Record<string, Stub>)
-  }
-
+const resolveComponentStubByName = (
+  componentName: string,
+  stubs: Record<string, Component | boolean>
+) => {
   for (const [stubKey, value] of Object.entries(stubs)) {
     if (matchName(componentName, stubKey)) {
       return value
@@ -129,7 +101,7 @@ const resolveComponentStubByName = (componentName: string, stubs: Stubs) => {
 }
 
 interface CreateStubComponentsTransformerConfig {
-  stubs?: Stubs
+  stubs?: Record<string, Component | boolean>
   shallow?: boolean
   renderStubDefaultSlot: boolean
 }
@@ -153,7 +125,7 @@ export function createStubComponentsTransformer({
 
     // stub transition by default via config.global.stubs
     if (
-      (type === Transition || type === BaseTransition) &&
+      (type === Transition || (type as any) === BaseTransition) &&
       'transition' in stubs
     ) {
       if (stubs.transition === false) return type
@@ -166,7 +138,7 @@ export function createStubComponentsTransformer({
     }
 
     // stub transition-group by default via config.global.stubs
-    if (type === TransitionGroup && 'transition-group' in stubs) {
+    if ((type as any) === TransitionGroup && 'transition-group' in stubs) {
       if (stubs['transition-group'] === false) return type
 
       return createStub({
@@ -206,16 +178,17 @@ export function createStubComponentsTransformer({
     if (isComponent(stub)) {
       const unwrappedStub = unwrapLegacyVueExtendComponent(stub)
       const stubFn = isFunctionalComponent(unwrappedStub) ? unwrappedStub : null
+
+      // Edge case: stub is component, we will not render stub but instead will create
+      // a new "copy" of stub component definition, but we want user still to be able
+      // to find our component by stub definition, so we register it manually
+      registerStub({ source: type, stub })
+
       const specializedStubComponent: ConcreteComponent = stubFn
         ? (...args) => stubFn(...args)
         : { ...unwrappedStub }
       specializedStubComponent.props = unwrappedStub.props
 
-      registerStub({
-        source: type,
-        stub: specializedStubComponent,
-        originalStub: stub
-      })
       return specializedStubComponent
     }
 
@@ -241,7 +214,6 @@ export function createStubComponentsTransformer({
           type,
           renderStubDefaultSlot
         })
-      registerStub({ source: type, stub: newStub })
       return newStub
     }
 
diff --git a/src/vnodeTransformers/stubDirectivesTransformer.ts b/src/vnodeTransformers/stubDirectivesTransformer.ts
new file mode 100644
index 000000000..2e26d32ba
--- /dev/null
+++ b/src/vnodeTransformers/stubDirectivesTransformer.ts
@@ -0,0 +1,46 @@
+import { Directive } from 'vue'
+import { isObjectComponent } from '../utils'
+import type { VTUVNodeTypeTransformer } from './util'
+
+interface CreateDirectivesTransformerConfig {
+  directives: Record<string, Directive | boolean>
+}
+
+const noop = () => {}
+export function createStubDirectivesTransformer({
+  directives = {}
+}: CreateDirectivesTransformerConfig): VTUVNodeTypeTransformer {
+  if (Object.keys(directives).length === 0) {
+    return (type) => type
+  }
+
+  return function directivesTransformer(type) {
+    if (isObjectComponent(type) && type.directives) {
+      const directivesToPatch = Object.keys(type.directives).filter(
+        (key) => key in directives
+      )
+
+      if (!directivesToPatch.length) {
+        return type
+      }
+
+      const replacementDirectives = Object.fromEntries(
+        directivesToPatch.map((name) => {
+          const directive = directives[name]
+
+          return [name, typeof directive === 'boolean' ? noop : directive]
+        })
+      )
+
+      return {
+        ...type,
+        directives: {
+          ...type.directives,
+          ...replacementDirectives
+        }
+      }
+    }
+
+    return type
+  }
+}
diff --git a/src/vnodeTransformers/util.ts b/src/vnodeTransformers/util.ts
index 563143ab6..3a3987b5a 100644
--- a/src/vnodeTransformers/util.ts
+++ b/src/vnodeTransformers/util.ts
@@ -1,4 +1,5 @@
 import { isComponent } from '../utils'
+import { registerStub } from '../stubs'
 import { ConcreteComponent, transformVNodeArgs } from 'vue'
 
 type VNodeArgsTransformerFn = NonNullable<
@@ -9,8 +10,10 @@ type VNodeTransformerArgsType = TransformVNodeArgs[0]
 type InstanceArgsType = TransformVNodeArgs[1]
 type VNodeTransformerInputType = VNodeTransformerArgsType[0]
 
-type VNodeTransformerInputComponentType = VNodeTransformerInputType &
-  ConcreteComponent
+type ExtractComponentTypes<T> = T extends ConcreteComponent ? T : never
+
+type VNodeTransformerInputComponentType =
+  ExtractComponentTypes<VNodeTransformerInputType>
 
 export type VTUVNodeTypeTransformer = (
   inputType: VNodeTransformerInputComponentType,
@@ -45,7 +48,12 @@ export const createVNodeTransformer = ({
       (type, transformer) => transformer(type, instance),
       componentType
     )
-    transformationCache.set(originalType, transformedType)
+
+    if (originalType !== transformedType) {
+      transformationCache.set(originalType, transformedType)
+
+      registerStub({ source: originalType, stub: transformedType })
+    }
 
     return [transformedType, ...restVNodeArgs]
   }
diff --git a/tests/findComponent.spec.ts b/tests/findComponent.spec.ts
index e30e06608..2882ce5f9 100644
--- a/tests/findComponent.spec.ts
+++ b/tests/findComponent.spec.ts
@@ -541,4 +541,24 @@ describe('findComponent', () => {
       ).toBe(true)
     })
   })
+
+  it('finds a component without name by its definition when stubbing directives', () => {
+    const ChildComponent = { template: 'child' }
+
+    const Component = defineComponent({
+      template: '<div><other-name /></div>',
+      components: {
+        OtherName: ChildComponent
+      }
+    })
+
+    const wrapper = mount(Component, {
+      global: {
+        stubs: {
+          vFakeDirective: {}
+        }
+      }
+    })
+    expect(wrapper.findComponent(ChildComponent).exists()).toBe(true)
+  })
 })
diff --git a/tests/mountingOptions/global.stubs.spec.ts b/tests/mountingOptions/global.stubs.spec.ts
index f7534b0f6..4c47185fc 100644
--- a/tests/mountingOptions/global.stubs.spec.ts
+++ b/tests/mountingOptions/global.stubs.spec.ts
@@ -1,6 +1,5 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
-import { h, defineComponent, defineAsyncComponent } from 'vue'
-
+import { h, defineComponent, defineAsyncComponent, Directive } from 'vue'
 import { config, flushPromises, mount, RouterLinkStub } from '../../src'
 import Hello from '../components/Hello.vue'
 import ComponentWithoutName from '../components/ComponentWithoutName.vue'
@@ -853,4 +852,105 @@ describe('mounting options: stubs', () => {
 
     expect(stub.vm).toBe(stubAfterSecondRender.vm)
   })
+
+  describe('directives', () => {
+    const MyDirective: Directive = {
+      beforeMount(el: Element) {
+        el.classList.add('DirectiveAdded')
+      }
+    }
+
+    const MyDirectiveStub: Directive = {
+      beforeMount(el: Element) {
+        el.classList.add('DirectiveStubAdded')
+      }
+    }
+
+    it('stubs directive on root component', () => {
+      const Component = {
+        template: '<div v-my-directive>text</div>',
+        directives: { MyDirective }
+      }
+
+      const wrapper = mount(Component, {
+        global: {
+          stubs: {
+            vMyDirective: MyDirectiveStub
+          }
+        }
+      })
+
+      expect(wrapper.classes()).toContain('DirectiveStubAdded')
+    })
+
+    it('stubs directive as noop', () => {
+      const SomeDirective = () => {
+        throw new Error('I will blow up!')
+      }
+
+      const Component = {
+        template: '<div v-some-directive>text</div>',
+        directives: { SomeDirective }
+      }
+
+      const wrapper = mount(Component, {
+        global: {
+          stubs: {
+            vSomeDirective: true
+          }
+        }
+      })
+
+      expect(wrapper.html()).toBe('<div>text</div>')
+    })
+
+    it('stubs directive on child component', () => {
+      const Component = {
+        template: '<div v-my-directive>text</div>',
+        directives: { MyDirective }
+      }
+
+      const RootComponent = {
+        components: { ChildComponent: Component },
+        template: '<div><child-component class="child" /></div>'
+      }
+
+      const wrapper = mount(RootComponent, {
+        global: {
+          stubs: {
+            vMyDirective: MyDirectiveStub
+          }
+        }
+      })
+
+      expect(wrapper.find('.child').classes()).toContain('DirectiveStubAdded')
+    })
+
+    it('stubs directive on root component with script setup', () => {
+      const Component = {
+        setup() {
+          // @ts-ignore-error (directive used by script setup)
+          const vMyDirective = MyDirective
+        },
+        template: '<div v-my-directive>text</div>',
+        directives: { MyDirective }
+      }
+
+      const MyDirectiveStub: Directive = {
+        beforeMount(el: Element) {
+          el.classList.add('DirectiveStubAdded')
+        }
+      }
+
+      const wrapper = mount(Component, {
+        global: {
+          stubs: {
+            vMyDirective: MyDirectiveStub
+          }
+        }
+      })
+
+      expect(wrapper.classes()).toContain('DirectiveStubAdded')
+    })
+  })
 })