fix(core): restore async doMock factory and doUnmock recovery (#998)

Author: fi3ework

e2e/mock/tests/doMockAsyncFactory.test.ts

@@ -0,0 +1,18 @@
+import { expect, rs, test } from '@rstest/core';
+
+test('doMock should support async factory with importActual', async () => {
+  rs.doMock('../src/increment', async () => {
+    const actual =
+      await rs.importActual<typeof import('../src/increment')>(
+        '../src/increment',
+      );
+
+    return {
+      ...actual,
+      increment: (num: number) => num + 100,
+    };
+  });
+
+  const { increment } = await import('../src/increment');
+  expect(increment(1)).toBe(101);
+});

e2e/mock/tests/doUnmockRestoreOriginal.test.ts

@@ -0,0 +1,30 @@
+import { afterEach, expect, rs, test } from '@rstest/core';
+
+afterEach(() => {
+  rs.doUnmock('../src/increment');
+  rs.resetModules();
+});
+
+test('doUnmock should restore original after multiple doMock calls', async () => {
+  rs.doMock('../src/increment', () => ({
+    increment: (num: number) => num + 100,
+  }));
+
+  const { increment: firstIncrement } = await import('../src/increment');
+  expect(firstIncrement(1)).toBe(101);
+
+  rs.resetModules();
+
+  rs.doMock('../src/increment', () => ({
+    increment: (num: number) => num + 200,
+  }));
+
+  const { increment: secondIncrement } = await import('../src/increment');
+  expect(secondIncrement(1)).toBe(201);
+
+  rs.doUnmock('../src/increment');
+  rs.resetModules();
+
+  const { increment } = await import('../src/increment');
+  expect(increment(1)).toBe(2);
+});

packages/core/src/core/plugins/mockRuntimeCode.js

@@ -35,6 +35,10 @@ __webpack_require__ = new Proxy(
 __webpack_require__.rstest_original_modules = {};
 __webpack_require__.rstest_original_module_factories = {};
 
+const hasOwn = (target, property) => Object.hasOwn(target, property);
+
+const isPromise = (value) => value instanceof Promise;
+
 //#region rs.unmock
 __webpack_require__.rstest_unmock = (id) => {
   const originalModuleFactory =
@@ -55,12 +59,11 @@ __webpack_require__.rstest_do_unmock = __webpack_require__.rstest_unmock;
 //#region rs.requireActual
 __webpack_require__.rstest_require_actual =
   __webpack_require__.rstest_import_actual = (id) => {
-    const originalModule = __webpack_require__.rstest_original_modules[id];
-
-    if (originalModule) {
-      return originalModule;
+    if (hasOwn(__webpack_require__.rstest_original_modules, id)) {
+      return __webpack_require__.rstest_original_modules[id];
     }
-    if (id in __webpack_require__.rstest_original_module_factories) {
+
+    if (hasOwn(__webpack_require__.rstest_original_module_factories, id)) {
       const mod = __webpack_require__.rstest_original_module_factories[id];
       const moduleInstance = { exports: {} };
       mod(moduleInstance, moduleInstance.exports, __webpack_require__);
@@ -75,20 +78,30 @@ __webpack_require__.rstest_require_actual =
 const getMockImplementation = (mockType = 'mock') => {
   const isMockRequire =
     mockType === 'mockRequire' || mockType === 'doMockRequire';
+
   // The mock and mockRequire will resolve to different module ids when the module is a dual package
   return (id, modFactory) => {
     // Only load the module if it's already in cache (to avoid side effects)
-    let requiredModule = __webpack_module_cache__[id]?.exports;
-    const wasAlreadyLoaded = !!requiredModule;
-
-    if (!requiredModule) {
-      // Module hasn't been loaded yet, so we can't get the original
-      // But we still need to save the original factory if it exists
-      __webpack_require__.rstest_original_module_factories[id] =
-        __webpack_modules__[id];
-    } else {
-      // Module was already loaded, save it
+    const hasCachedModule = hasOwn(__webpack_module_cache__, id);
+    let requiredModule = hasCachedModule
+      ? __webpack_module_cache__[id].exports
+      : undefined;
+    const wasAlreadyLoaded = hasCachedModule;
+
+    const hasSavedOriginalModule = hasOwn(
+      __webpack_require__.rstest_original_modules,
+      id,
+    );
+    const hasSavedOriginalFactory = hasOwn(
+      __webpack_require__.rstest_original_module_factories,
+      id,
+    );
+
+    if (!hasSavedOriginalModule && hasCachedModule) {
       __webpack_require__.rstest_original_modules[id] = requiredModule;
+    }
+
+    if (!hasSavedOriginalFactory) {
       __webpack_require__.rstest_original_module_factories[id] =
         __webpack_modules__[id];
     }
@@ -170,6 +183,11 @@ const getMockImplementation = (mockType = 'mock') => {
       ) {
         const res = modFactory();
 
+        if (isPromise(res)) {
+          __webpack_module__.exports = res;
+          return;
+        }
+
         if (isMockRequire) {
           __webpack_module__.exports = res;
           return;

website/docs/en/api/runtime-api/rstest/mock-modules.mdx

@@ -8,7 +8,7 @@ Rstest supports mocking modules, which allows you to replace the implementation
 
 ## rs.mock
 
-- **Type**: `<T = unknown>(moduleName: string | Promise<T>, factoryOrOptions?: (() => Partial<T>) | { spy: true } | { mock: true }) => void`
+- **Type**: `<T = unknown>(moduleName: string | Promise<T>, factoryOrOptions?: (() => Promise<Partial<T>> | Partial<T>) | { spy: true } | { mock: true }) => void`
 
 Mocks and replaces the module specified in the first parameter.
 
@@ -226,7 +226,7 @@ rs.mock(import('../src/b'), () => {
 
 ## rs.doMock
 
-- **Type**: `<T = unknown>(moduleName: string | Promise<T>, factoryOrOptions?: (() => Partial<T>) | { spy: true } | { mock: true }) => void`
+- **Type**: `<T = unknown>(moduleName: string | Promise<T>, factoryOrOptions?: (() => Promise<Partial<T>> | Partial<T>) | { spy: true } | { mock: true }) => void`
 
 Similar to `rs.mock`, but it is **not hoisted** to the top of the module. It is called when it's executed, which means that if a module has already been imported before calling `rs.doMock`, that module will not be mocked, while modules imported after calling `rs.doMock` will be mocked.
 

website/docs/zh/api/runtime-api/rstest/mock-modules.mdx

@@ -8,7 +8,7 @@ Rstest 支持对模块进行 mock，这使得你可以在测试中替换模块
 
 ## rs.mock
 
-- **类型：** `<T = unknown>(moduleName: string | Promise<T>, factoryOrOptions?: (() => Partial<T>) | { spy: true } | { mock: true }) => void`
+- **类型：** `<T = unknown>(moduleName: string | Promise<T>, factoryOrOptions?: (() => Promise<Partial<T>> | Partial<T>) | { spy: true } | { mock: true }) => void`
 
 对第一个参数对应的模块进行 mock 替换。
 
@@ -226,7 +226,7 @@ rs.mock(import('../src/b'), () => {
 
 ## rs.doMock
 
-- **类型：** `<T = unknown>(moduleName: string | Promise<T>, factoryOrOptions?: (() => Partial<T>) | { spy: true } | { mock: true }) => void`
+- **类型：** `<T = unknown>(moduleName: string | Promise<T>, factoryOrOptions?: (() => Promise<Partial<T>> | Partial<T>) | { spy: true } | { mock: true }) => void`
 
 与 `rs.mock` 类似，但它**不会被提升**到模块顶部。它会在被执行到时调用，这意味着如果在调用 `rs.doMock` 之前已经导入了模块，则该模块不会被 mock，而在调用 `rs.doMock` 之后导入的模块会被 mock。