From 3912554b50e6111722961a4315366de636bd66a8 Mon Sep 17 00:00:00 2001
From: chaptergy <git@chaptergy.xyz>
Date: Tue, 7 Jan 2025 19:36:42 +0100
Subject: [PATCH] fix: allow `getMockImplementation` to return "once"
 implementation (#7033)

---
 packages/spy/src/index.ts        | 10 +++++-----
 test/core/test/jest-mock.test.ts | 20 ++++++++++++++++++++
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/packages/spy/src/index.ts b/packages/spy/src/index.ts
index 4f3180224621..5cc6a787deb9 100644
--- a/packages/spy/src/index.ts
+++ b/packages/spy/src/index.ts
@@ -215,12 +215,11 @@ export interface MockInstance<T extends Procedure = Procedure> {
    */
   mockRestore(): void
   /**
-   * Performs the same actions as `mockReset` and restores the inner implementation to the original function.
+   * Returns current permanent mock implementation if there is one.
    *
-   * Note that restoring a mock created with `vi.fn()` will set the implementation to an empty function that returns `undefined`. Restoring a mock created with `vi.fn(impl)` will restore the implementation to `impl`.
+   * If mock was created with `vi.fn`, it will consider passed down method as a mock implementation.
    *
-   * To automatically call this method before each test, enable the [`restoreMocks`](https://vitest.dev/config/#restoremocks) setting in the configuration.
-   * @see https://vitest.dev/api/mock#getmockimplementation
+   * If mock was created with `vi.spyOn`, it will return `undefined` unless a custom implementation was provided.
    */
   getMockImplementation(): NormalizedPrecedure<T> | undefined
   /**
@@ -575,7 +574,8 @@ function enhanceSpy<T extends Procedure>(
     return stub
   }
 
-  stub.getMockImplementation = () => implementation
+  stub.getMockImplementation = () =>
+    implementationChangedTemporarily ? implementation : (onceImplementations.at(0) || implementation)
   stub.mockImplementation = (fn: T) => {
     implementation = fn
     state.willCall(mockCall)
diff --git a/test/core/test/jest-mock.test.ts b/test/core/test/jest-mock.test.ts
index 32c699601c82..03026c80c292 100644
--- a/test/core/test/jest-mock.test.ts
+++ b/test/core/test/jest-mock.test.ts
@@ -518,4 +518,24 @@ describe('jest mock compat layer', () => {
     vi.mocked(dogMax.speak).mockReturnValue('woof woof')
     expect(dogMax.speak()).toBe('woof woof')
   })
+
+  it('returns temporary implementations from getMockImplementation()', () => {
+    const fn = vi.fn()
+
+    const temporaryMockImplementation = () => 'mocked value'
+    fn.mockImplementationOnce(temporaryMockImplementation)
+    expect(fn.getMockImplementation()).toBe(temporaryMockImplementation)
+
+    // After calling it, it should be back to undefined
+    fn()
+    expect(fn.getMockImplementation()).toBe(undefined)
+
+    const mockImplementation = () => 'other mocked value'
+    fn.mockImplementation(mockImplementation)
+    expect(fn.getMockImplementation()).toBe(mockImplementation)
+
+    // It should also overwrite permanent implementations
+    fn.mockImplementationOnce(temporaryMockImplementation)
+    expect(fn.getMockImplementation()).toBe(temporaryMockImplementation)
+  })
 })