From 97926a27895e8352912cf6e15f0c59eb565a366f Mon Sep 17 00:00:00 2001
From: Illya Klymov <xanf@xanf.me>
Date: Thu, 5 Aug 2021 06:22:17 +0300
Subject: [PATCH] fix(stubs): Do not create stubs across multiple renders
 (#813)

---
 src/stubs.ts                               | 50 +++++++++++++++-------
 tests/mountingOptions/global.stubs.spec.ts | 30 ++++++++++++-
 2 files changed, 64 insertions(+), 16 deletions(-)

diff --git a/src/stubs.ts b/src/stubs.ts
index 1291ff0fa..e91680651 100644
--- a/src/stubs.ts
+++ b/src/stubs.ts
@@ -118,12 +118,35 @@ const getComponentName = (type: VNodeTypes): string => {
   return ''
 }
 
+function createStubOnceForType(
+  type: {} & VNodeTypes,
+  factoryFn: () => ConcreteComponent,
+  cache: WeakMap<{} & VNodeTypes, ConcreteComponent>
+): ConcreteComponent {
+  const cachedStub = cache.get(type)
+  if (cachedStub) {
+    return cachedStub
+  }
+
+  const stub = factoryFn()
+  stubsMap.set(stub, type)
+  cache.set(type, stub)
+  return stub
+}
+
 export function stubComponents(
   stubs: Stubs = {},
   shallow: boolean = false,
   renderStubDefaultSlot: boolean = false
 ) {
-  const components: Record<string, ComponentOptions> = {}
+  const createdStubsMap: WeakMap<{} & VNodeTypes, ConcreteComponent> =
+    new WeakMap()
+
+  const createStubOnce = (
+    type: {} & VNodeTypes,
+    factoryFn: () => ConcreteComponent
+  ) => createStubOnceForType(type, factoryFn, createdStubsMap)
+
   transformVNodeArgs((args, instance: ComponentInternalInstance | null) => {
     const [nodeType, props, children, patchFlag, dynamicProps] = args
     const type = nodeType as VNodeTypes
@@ -184,13 +207,15 @@ export function stubComponents(
       // case 2: custom implementation
       if (isComponent(stub)) {
         const stubFn = isFunctionalComponent(stub) ? stub : null
-
-        const specializedStub: ConcreteComponent = stubFn
+        const specializedStubComponent: ConcreteComponent = stubFn
           ? (...args) => stubFn(...args)
           : { ...stub }
+        specializedStubComponent.props = stub.props
 
-        specializedStub.props = stub.props
-        stubsMap.set(specializedStub, type)
+        const specializedStub = createStubOnce(
+          type,
+          () => specializedStubComponent
+        )
         // pass the props and children, for advanced stubbing
         return [specializedStub, props, children, patchFlag, dynamicProps]
       }
@@ -205,25 +230,20 @@ export function stubComponents(
       // case 1: default stub
       if (stub === true || shallow) {
         // Set name when using shallow without stub
-        if (!name) {
-          name = registeredName || componentName
-        }
+        const stubName = name || registeredName || componentName
 
         if (!isComponent(type)) {
           throw new Error('Attempted to stub a non-component')
         }
 
         const propsDeclaration = type.props || {}
-        let newStub = components[name]
-        if (!newStub) {
-          newStub = createStub({
-            name,
+        const newStub = createStubOnce(type, () =>
+          createStub({
+            name: stubName,
             propsDeclaration,
             renderStubDefaultSlot
           })
-          components[name] = newStub
-          stubsMap.set(newStub, type)
-        }
+        )
         return [newStub, props, children, patchFlag, dynamicProps]
       }
     }
diff --git a/tests/mountingOptions/global.stubs.spec.ts b/tests/mountingOptions/global.stubs.spec.ts
index 124279850..deb0dfff4 100644
--- a/tests/mountingOptions/global.stubs.spec.ts
+++ b/tests/mountingOptions/global.stubs.spec.ts
@@ -580,7 +580,6 @@ describe('mounting options: stubs', () => {
           }
         }
       })
-
       expect(wrapper.html()).toBe('<span>StubComponent</span>')
     })
   })
@@ -694,4 +693,33 @@ describe('mounting options: stubs', () => {
 
     expect(wrapper.html()).toBe('<anonymous-stub>test</anonymous-stub>')
   })
+
+  it('should not recreate stub across multiple renders', async () => {
+    const FooBar = {
+      name: 'FooBar',
+      render: () => h('span', 'real foobar')
+    }
+
+    const Comp = defineComponent({
+      data: () => ({ value: 1 }),
+      render() {
+        return h('div', {}, [this.value, h(FooBar)])
+      }
+    })
+
+    const wrapper = mount(Comp, {
+      global: {
+        stubs: {
+          'foo-bar': { name: 'FooBar', template: 'stub' }
+        }
+      }
+    })
+
+    const stub = wrapper.findComponent({ name: 'FooBar' })
+    await wrapper.setData({ value: 2 })
+
+    const stubAfterSecondRender = wrapper.findComponent({ name: 'FooBar' })
+
+    expect(stub.vm).toBe(stubAfterSecondRender.vm)
+  })
 })