fix(spy): update to set initial implementation through normal logic (f…

…ixes #3260) (#3263)
vitest-dev · Apr 27, 2023 · c759a9a · c759a9a
1 parent 481b1fd
commit c759a9a
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/packages/spy/src/index.ts b/packages/spy/src/index.ts
@@ -324,5 +324,9 @@ export function fn<TArgs extends any[] = any[], R = any>(
 export function fn<TArgs extends any[] = any[], R = any>(
   implementation?: (...args: TArgs) => R,
 ): Mock<TArgs, R> {
-  return enhanceSpy(tinyspy.internalSpyOn({ fn: implementation || (() => {}) }, 'fn')) as unknown as Mock
+  const enhancedSpy = enhanceSpy(tinyspy.internalSpyOn({ spy: () => {} }, 'spy'))
+  if (implementation)
+    enhancedSpy.mockImplementation(implementation)
+
+  return enhancedSpy as Mock
 }
diff --git a/test/core/test/jest-mock.test.ts b/test/core/test/jest-mock.test.ts
@@ -41,6 +41,17 @@ describe('jest mock compat layer', () => {
     expect(Spy.mock.instances).toHaveLength(0)
   })
 
+  it('implementation is set correctly on init', () => {
+    const impl = () => 1
+    const mock1 = vi.fn(impl)
+
+    expect(mock1.getMockImplementation()).toEqual(impl)
+
+    const mock2 = vi.fn()
+
+    expect(mock2.getMockImplementation()).toBeUndefined()
+  })
+
   it('implementation sync fn', () => {
     const originalFn = function () {
       return 'original'
@@ -49,7 +60,7 @@ describe('jest mock compat layer', () => {
 
     spy() // returns 'original'
 
-    expect(spy.getMockImplementation()).toBe(undefined)
+    expect(spy.getMockImplementation()).toBe(originalFn)
 
     spy.mockReturnValueOnce('2-once').mockReturnValueOnce('3-once')