feat(zone.js): support vitest patching in zone.js/testing

To support `fakeAsync` usage while using `vitest` as a test runner, Zone.js
now provides patching when using the `zone.js/testing` package import.
This patching is similar to that of the existing jasmine, mocha, and jest
functionality.

Author: clydin

packages/zone.js/lib/testing/zone-testing.ts

@@ -9,6 +9,7 @@
 import {patchJasmine} from '../jasmine/jasmine';
 import {patchJest} from '../jest/jest';
 import {patchMocha} from '../mocha/mocha';
+import {patchVitest} from '../vitest/vitest';
 import {ZoneType} from '../zone-impl';
 import {patchAsyncTest} from '../zone-spec/async-test';
 import {patchFakeAsyncTest} from '../zone-spec/fake-async-test';
@@ -25,6 +26,7 @@ export function rollupTesting(Zone: ZoneType): void {
   patchJasmine(Zone);
   patchJest(Zone);
   patchMocha(Zone);
+  patchVitest(Zone);
   patchAsyncTest(Zone);
   patchFakeAsyncTest(Zone);
   patchPromiseTesting(Zone);

packages/zone.js/lib/vitest/vitest.ts

@@ -0,0 +1,170 @@
+/**
+ * @license
+ * Copyright Google LLC All Rights Reserved.
+ *
+ * Use of this source code is governed by an MIT-style license that can be
+ * found in the LICENSE file at https://angular.dev/license
+ */
+
+import {ZoneType} from '../zone-impl';
+
+/**
+ * The ZoneType with additional testing related members added.
+ * The additional members must be patched onto the type via the
+ * `zone.js/testing` package entry point.
+ */
+interface TestingZoneType extends ZoneType {
+  ProxyZoneSpec?: typeof import('../zone-spec/proxy').ProxyZoneSpec;
+  SyncTestZoneSpec?: {new (namePrefix: string): ZoneSpec};
+}
+
+export function patchVitest(Zone: ZoneType): void {
+  Zone.__load_patch('vitest', (context: any, Zone: TestingZoneType) => {
+    // Vitest global variable set by the Vitest runner during test execution
+    const vitestGlobal = context['vitest'] as {['__zone_patch__']?: boolean} | undefined;
+
+    // Skip patching if vitest is not present or has already been patched
+    if (typeof vitestGlobal === 'undefined' || vitestGlobal['__zone_patch__']) {
+      return;
+    }
+    vitestGlobal['__zone_patch__'] = true;
+
+    // Ensure other testing related Zone.js patches have been applied
+    if (!Zone.ProxyZoneSpec) {
+      throw new Error('Missing ProxyZoneSpec');
+    }
+    if (!Zone.SyncTestZoneSpec) {
+      throw new Error('Missing SyncTestZoneSpec');
+    }
+
+    // Setup testing related Zone instances
+    const rootZone = Zone.current;
+    const syncZone = rootZone.fork(new Zone.SyncTestZoneSpec('vitest.describe'));
+    const proxyZone = rootZone.fork(new Zone.ProxyZoneSpec());
+
+    function wrapDescribeFactoryInZone(originalVitestFn: Function) {
+      return function (this: unknown, ...factoryArgs: unknown[]) {
+        const originalDescribeFn = originalVitestFn.apply(this, factoryArgs);
+        return function (this: unknown, ...args: [unknown, Function, ...unknown[]]) {
+          args[1] = wrapDescribeInZone(args[1]);
+          return originalDescribeFn.apply(this, args);
+        };
+      };
+    }
+
+    function wrapTestFactoryInZone(originalVitestFn: Function) {
+      return function (this: unknown, ...factoryArgs: unknown[]) {
+        return function (this: unknown, ...args: [unknown, Function, ...unknown[]]) {
+          args[1] = wrapTestInZone(args[1]);
+          return originalVitestFn.apply(this, factoryArgs).apply(this, args);
+        };
+      };
+    }
+
+    /**
+     * Gets a function wrapping the body of a vitest `describe` block to execute in a
+     * synchronous-only zone.
+     */
+    function wrapDescribeInZone(describeBody: Function): Function {
+      return function (this: unknown, ...args: unknown[]) {
+        return syncZone.run(describeBody, this, args);
+      };
+    }
+
+    /**
+     * Gets a function wrapping the body of a vitest `it/beforeEach/afterEach` block to
+     * execute in a ProxyZone zone.
+     * This will run in the `proxyZone`.
+     */
+    function wrapTestInZone(testBody: Function): Function {
+      if (typeof testBody !== 'function') {
+        return testBody;
+      }
+      const wrappedFunc = function () {
+        return proxyZone.run(testBody, null, arguments as any);
+      };
+      // Update the length of wrappedFunc to be the same as the length of the testBody
+      // So vitest core can handle whether the test function has `done()` or not correctly
+      Object.defineProperty(wrappedFunc, 'length', {
+        configurable: true,
+        writable: true,
+        enumerable: false,
+      });
+      wrappedFunc.length = testBody.length;
+      return wrappedFunc;
+    }
+
+    ['suite', 'describe'].forEach((methodName) => {
+      let originalVitestFn: Function = context[methodName];
+      if (context[Zone.__symbol__(methodName)]) {
+        return;
+      }
+
+      context[Zone.__symbol__(methodName)] = originalVitestFn;
+      context[methodName] = function (this: unknown, ...args: any[]) {
+        args[1] = wrapDescribeInZone(args[1]);
+        return originalVitestFn.apply(this, args);
+      };
+
+      [
+        'skip',
+        'skipIf',
+        'runIf',
+        'only',
+        'concurrent',
+        'sequential',
+        'shuffle',
+        'todo',
+        'each',
+        'for',
+      ].forEach((subMethodName) => {
+        context[methodName][subMethodName] = wrapDescribeFactoryInZone(
+          (originalVitestFn as any)[subMethodName],
+        );
+      });
+    });
+
+    ['it', 'test'].forEach((methodName) => {
+      let originalVitestFn: Function = context[methodName];
+      if (context[Zone.__symbol__(methodName)]) {
+        return;
+      }
+
+      context[Zone.__symbol__(methodName)] = originalVitestFn;
+      context[methodName] = function (this: unknown, ...args: any[]) {
+        args[1] = wrapTestInZone(args[1]);
+        return originalVitestFn.apply(this, args);
+      };
+
+      [
+        'skip',
+        'skipIf',
+        'runIf',
+        'only',
+        'concurrent',
+        'sequential',
+        'todo',
+        'fails',
+        'each',
+        'for',
+      ].forEach((subMethodName) => {
+        context[methodName][subMethodName] = wrapTestFactoryInZone(
+          (originalVitestFn as any)[subMethodName],
+        );
+      });
+    });
+
+    ['beforeEach', 'afterEach', 'beforeAll', 'afterAll'].forEach((methodName) => {
+      const originalVitestFn: Function = context[methodName];
+      if (context[Zone.__symbol__(methodName)]) {
+        return;
+      }
+
+      context[Zone.__symbol__(methodName)] = originalVitestFn;
+      context[methodName] = function (this: unknown, ...args: [Function, ...unknown[]]) {
+        args[0] = wrapTestInZone(args[0]);
+        return originalVitestFn.apply(this, args);
+      };
+    });
+  });
+}