feat: support mock withImplementation (#128)

Author: 9aoy

packages/core/src/runtime/api/spy.ts

@@ -35,6 +35,37 @@ const wrapSpy = <T extends FunctionLike>(
       : implementation;
   };
 
+  function withImplementation(fn: T, cb: () => void): void;
+  function withImplementation(fn: T, cb: () => Promise<void>): Promise<void>;
+  function withImplementation(
+    fn: T,
+    cb: () => void | Promise<void>,
+  ): void | Promise<void> {
+    const originalImplementation = implementation;
+    const originalMockImplementationOnce = mockImplementationOnce;
+
+    implementation = fn;
+    mockImplementationOnce = [];
+    spyState.willCall(willCall);
+
+    const reset = () => {
+      implementation = originalImplementation;
+      mockImplementationOnce = originalMockImplementationOnce;
+    };
+
+    const result = cb();
+
+    if (result instanceof Promise) {
+      return result.then(() => {
+        reset();
+      });
+    }
+
+    reset();
+  }
+
+  spyFn.withImplementation = withImplementation;
+
   spyFn.mockImplementation = (fn) => {
     implementation = fn;
     return spyFn;

packages/core/src/types/mock.ts

@@ -89,10 +89,10 @@ export interface MockInstance<T extends FunctionLike = FunctionLike> {
   getMockImplementation(): T | undefined;
   mockImplementation(fn: T): this;
   mockImplementationOnce(fn: T): this;
-  // withImplementation<T2>(
-  //   fn: T,
-  //   callback: () => T2,
-  // ): T2 extends Promise<unknown> ? Promise<void> : void;
+  withImplementation<T2>(
+    fn: T,
+    callback: () => T2,
+  ): T2 extends Promise<unknown> ? Promise<void> : void;
   // mockReturnThis(): this;
   // mockReturnValue(value: ReturnType<T>): this;
   // mockReturnValueOnce(value: ReturnType<T>): this;

tests/spy/fixtures/withImplementation.test.ts

@@ -0,0 +1,72 @@
+import { describe, expect, it, rstest } from '@rstest/core';
+
+describe('test withImplementation', () => {
+  it('withImplementation', () => {
+    let isMockCalled = false;
+    const mockFn = () => {
+      isMockCalled = true;
+      console.log('[call original]');
+      return 'original';
+    };
+    const myMockFn = rstest.fn(mockFn);
+
+    const mockFn1 = () => {
+      isMockCalled = true;
+      console.log('[call original - 1]');
+      return 'original - 1';
+    };
+
+    myMockFn.mockImplementationOnce(mockFn1);
+
+    myMockFn.withImplementation(
+      () => {
+        console.log('[call temp]');
+        return 'temp';
+      },
+      () => {
+        console.log('[call callback]');
+        const res = myMockFn();
+        const res1 = myMockFn();
+        console.log('[callback res]', res, res1);
+      },
+    );
+
+    expect(myMockFn.getMockImplementation()).toBe(mockFn1);
+    expect(isMockCalled).toBe(false);
+
+    console.log('[call myMockFn]');
+    expect(myMockFn()).toBe('original - 1');
+    expect(isMockCalled).toBe(true);
+
+    console.log('[call myMockFn - 1]');
+
+    expect(myMockFn()).toBe('original');
+    expect(myMockFn).toHaveBeenCalledTimes(4);
+  });
+
+  it('withImplementation async', async () => {
+    const myMockFn = rstest.fn(() => {
+      console.log('[1 - call original]');
+      return 'original';
+    });
+
+    await myMockFn.withImplementation(
+      () => {
+        console.log('[1 - call temp]');
+        return 'temp';
+      },
+      async () => {
+        console.log('[1 - call callback]');
+        const res = myMockFn();
+        console.log('[1 - callback res]', res);
+      },
+    );
+
+    console.log('[1 - call myMockFn]');
+
+    expect(myMockFn()).toBe('original');
+
+    console.log('[1 - call myMockFn - 1]');
+    expect(myMockFn()).toBe('original');
+  });
+});

tests/spy/withImplementation.test.ts

@@ -0,0 +1,43 @@
+import { dirname } from 'node:path';
+import { fileURLToPath } from 'node:url';
+import { describe, expect, it } from '@rstest/core';
+import { runRstestCli } from '../scripts';
+
+const __filename = fileURLToPath(import.meta.url);
+const __dirname = dirname(__filename);
+
+it('test withImplementation', async () => {
+  const { cli } = await runRstestCli({
+    command: 'rstest',
+    args: ['run', 'fixtures/withImplementation.test'],
+    options: {
+      nodeOptions: {
+        cwd: __dirname,
+      },
+    },
+  });
+
+  await cli.exec;
+  expect(cli.exec.process?.exitCode).toBe(0);
+  const logs = cli.stdout.split('\n').filter(Boolean);
+
+  expect(logs.filter((log) => log.startsWith('['))).toMatchInlineSnapshot(`
+    [
+      "[call callback]",
+      "[call temp]",
+      "[call temp]",
+      "[callback res] temp temp",
+      "[call myMockFn]",
+      "[call original - 1]",
+      "[call myMockFn - 1]",
+      "[call original]",
+      "[1 - call callback]",
+      "[1 - call temp]",
+      "[1 - callback res] temp",
+      "[1 - call myMockFn]",
+      "[1 - call original]",
+      "[1 - call myMockFn - 1]",
+      "[1 - call original]",
+    ]
+  `);
+});