fix(vitest): independently mock each instance's methods for mocked cl…

…ass (#4564)

packages/vitest/src/runtime/mocker.ts

@@ -325,13 +325,36 @@ export class VitestMocker {
           continue
 
         if (isFunction) {
-          const spyModule = this.spyModule
-          if (!spyModule)
+          if (!this.spyModule)
             throw this.createError('[vitest] `spyModule` is not defined. This is Vitest error. Please open a new issue with reproduction.')
-          const mock = spyModule.spyOn(newContainer, property).mockImplementation(() => undefined)
+          const spyModule = this.spyModule
+          const primitives = this.primitives
+          function mockFunction(this: any) {
+            // detect constructor call and mock each instance's methods
+            // so that mock states between prototype/instances don't affect each other
+            // (jest reference https://github.com/jestjs/jest/blob/2c3d2409879952157433de215ae0eee5188a4384/packages/jest-mock/src/index.ts#L678-L691)
+            if (this instanceof newContainer[property]) {
+              for (const { key } of getAllMockableProperties(this, false, primitives)) {
+                const value = this[key]
+                const type = getType(value)
+                const isFunction = type.includes('Function') && typeof value === 'function'
+                if (isFunction) {
+                  // mock and delegate calls to original prototype method, which should be also mocked already
+                  const original = this[key]
+                  const mock = spyModule.spyOn(this, key as string).mockImplementation(original)
+                  mock.mockRestore = () => {
+                    mock.mockReset()
+                    mock.mockImplementation(original)
+                    return mock
+                  }
+                }
+              }
+            }
+          }
+          const mock = spyModule.spyOn(newContainer, property).mockImplementation(mockFunction)
           mock.mockRestore = () => {
             mock.mockReset()
-            mock.mockImplementation(() => undefined)
+            mock.mockImplementation(mockFunction)
             return mock
           }
           // tinyspy retains length, but jest doesn't.

test/core/src/mockedE.ts

@@ -0,0 +1,11 @@
+export const symbolFn = Symbol.for('symbolFn')
+
+export class MockedE {
+  public testFn(arg: string) {
+    return arg.repeat(2)
+  }
+
+  public [symbolFn](arg: string) {
+    return arg.repeat(2)
+  }
+}

test/core/test/mocked-class-restore-all.test.ts

@@ -0,0 +1,78 @@
+import { expect, test, vi } from 'vitest'
+
+import { MockedE } from '../src/mockedE'
+
+vi.mock('../src/mockedE')
+
+test(`mocked class are not affected by restoreAllMocks`, () => {
+  const instance0 = new MockedE()
+  expect(instance0.testFn('a')).toMatchInlineSnapshot(`undefined`)
+  expect(vi.mocked(instance0.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "a",
+      ],
+    ]
+  `)
+  expect(vi.mocked(MockedE.prototype.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "a",
+      ],
+    ]
+  `)
+
+  vi.restoreAllMocks()
+
+  // reset only history after restoreAllMocks
+  expect(instance0.testFn('b')).toMatchInlineSnapshot(`undefined`)
+  expect(vi.mocked(instance0.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "b",
+      ],
+    ]
+  `)
+  expect(vi.mocked(MockedE.prototype.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "b",
+      ],
+    ]
+  `)
+
+  // mocked constructor is still effective after restoreAllMocks
+  const instance1 = new MockedE()
+  const instance2 = new MockedE()
+  expect(instance1).not.toBe(instance2)
+  expect(instance1.testFn).not.toBe(instance2.testFn)
+  expect(instance1.testFn).not.toBe(MockedE.prototype.testFn)
+  expect(vi.mocked(instance1.testFn).mock).not.toBe(vi.mocked(instance2.testFn).mock)
+
+  expect(instance1.testFn('c')).toMatchInlineSnapshot(`undefined`)
+  expect(vi.mocked(instance0.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "b",
+      ],
+    ]
+  `)
+  expect(vi.mocked(instance1.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "c",
+      ],
+    ]
+  `)
+  expect(vi.mocked(instance2.testFn).mock.calls).toMatchInlineSnapshot(`[]`)
+  expect(vi.mocked(MockedE.prototype.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "b",
+      ],
+      [
+        "c",
+      ],
+    ]
+  `)
+})

test/core/test/mocked-class-restore-explicit.test.ts

@@ -0,0 +1,86 @@
+import { expect, test, vi } from 'vitest'
+
+import { MockedE } from '../src/mockedE'
+
+vi.mock('../src/mockedE')
+
+// this behavior looks odd but jest also doesn't seem to support this use case properly
+test(`mocked class methods are not restorable by explicit mockRestore calls`, () => {
+  const instance1 = new MockedE()
+
+  expect(instance1.testFn('a')).toMatchInlineSnapshot(`undefined`)
+  expect(vi.mocked(instance1.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "a",
+      ],
+    ]
+  `)
+  expect(vi.mocked(MockedE.prototype.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "a",
+      ],
+    ]
+  `)
+
+  // restoring instance method
+  vi.mocked(instance1.testFn).mockRestore()
+  expect(vi.mocked(instance1.testFn).mock.calls).toMatchInlineSnapshot(`[]`)
+  expect(vi.mocked(MockedE.prototype.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "a",
+      ],
+    ]
+  `)
+
+  expect(instance1.testFn('b')).toMatchInlineSnapshot(`undefined`)
+  expect(vi.mocked(instance1.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "b",
+      ],
+    ]
+  `)
+  expect(vi.mocked(MockedE.prototype.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "a",
+      ],
+      [
+        "b",
+      ],
+    ]
+  `)
+
+  // restoring prototype method
+  vi.mocked(MockedE.prototype.testFn).mockRestore()
+  expect(vi.mocked(instance1.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "b",
+      ],
+    ]
+  `)
+  expect(vi.mocked(MockedE.prototype.testFn).mock.calls).toMatchInlineSnapshot(`[]`)
+
+  expect(instance1.testFn('c')).toMatchInlineSnapshot(`undefined`)
+  expect(vi.mocked(instance1.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "b",
+      ],
+      [
+        "c",
+      ],
+    ]
+  `)
+  expect(vi.mocked(MockedE.prototype.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "c",
+      ],
+    ]
+  `)
+})

test/core/test/mocked-class.test.ts

@@ -0,0 +1,115 @@
+import { expect, test, vi } from 'vitest'
+
+import { MockedE, symbolFn } from '../src/mockedE'
+
+vi.mock('../src/mockedE')
+
+test(`each instance's methods of mocked class should have independent mock function state`, () => {
+  const instance1 = new MockedE()
+  const instance2 = new MockedE()
+  expect(instance1).not.toBe(instance2)
+  expect(instance1.testFn).not.toBe(instance2.testFn)
+  expect(instance1.testFn).not.toBe(MockedE.prototype.testFn)
+  expect(vi.mocked(instance1.testFn).mock).not.toBe(vi.mocked(instance2.testFn).mock)
+
+  expect(instance1.testFn('a')).toMatchInlineSnapshot(`undefined`)
+  expect(instance1.testFn).toBeCalledTimes(1)
+  expect(instance2.testFn).toBeCalledTimes(0)
+  expect(MockedE.prototype.testFn).toBeCalledTimes(1)
+  expect(vi.mocked(instance1.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "a",
+      ],
+    ]
+  `)
+  expect(vi.mocked(MockedE.prototype.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "a",
+      ],
+    ]
+  `)
+
+  expect(instance2.testFn('b')).toMatchInlineSnapshot(`undefined`)
+  expect(instance1.testFn).toBeCalledTimes(1)
+  expect(instance2.testFn).toBeCalledTimes(1)
+  expect(MockedE.prototype.testFn).toBeCalledTimes(2)
+  expect(vi.mocked(instance2.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "b",
+      ],
+    ]
+  `)
+  expect(vi.mocked(MockedE.prototype.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "a",
+      ],
+      [
+        "b",
+      ],
+    ]
+  `)
+
+  expect(instance1.testFn('c')).toMatchInlineSnapshot(`undefined`)
+  expect(instance1.testFn).toBeCalledTimes(2)
+  expect(instance2.testFn).toBeCalledTimes(1)
+  expect(MockedE.prototype.testFn).toBeCalledTimes(3)
+  expect(vi.mocked(instance1.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "a",
+      ],
+      [
+        "c",
+      ],
+    ]
+  `)
+  expect(vi.mocked(instance2.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "b",
+      ],
+    ]
+  `)
+  expect(vi.mocked(MockedE.prototype.testFn).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "a",
+      ],
+      [
+        "b",
+      ],
+      [
+        "c",
+      ],
+    ]
+  `)
+
+  // test same things for symbol key method
+  expect(instance1[symbolFn]).not.toBe(instance2[symbolFn])
+  expect(instance1[symbolFn]).not.toBe(MockedE.prototype[symbolFn])
+  expect(vi.mocked(instance1[symbolFn]).mock).not.toBe(vi.mocked(instance2[symbolFn]).mock)
+
+  expect(instance1[symbolFn]('d')).toMatchInlineSnapshot(`undefined`)
+  expect(instance1[symbolFn]).toBeCalledTimes(1)
+  expect(instance2[symbolFn]).toBeCalledTimes(0)
+  expect(MockedE.prototype[symbolFn]).toBeCalledTimes(1)
+  expect(vi.mocked(instance1[symbolFn]).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "d",
+      ],
+    ]
+  `)
+  expect(vi.mocked(instance2[symbolFn]).mock.calls).toMatchInlineSnapshot(`[]`)
+  expect(vi.mocked(MockedE.prototype[symbolFn]).mock.calls).toMatchInlineSnapshot(`
+    [
+      [
+        "d",
+      ],
+    ]
+  `)
+})