refactor(zone.js): Add a withProxyZone helper that might be used for unpatched test frameworks (angular#61626)

As an alternative to monkey patching vitest, this change adds a method that could be used
for manually running functions inside a shared proxy zone. If used inocrrectly,
this would mean that
the `fakeAsync` closure may not capture all timers and microtasks if it
invokes things created in a zone that was already forked (e.g. creating
a component in a beforeEach:
https://github.com/angular/angular/blob/2699dd65558d70fadeac1e9b420841f3dfc3a059/packages/zone.js/lib/jasmine/jasmine.ts#L363-L371)

PR Close angular#61626

Author: atscott

.github/workflows/ci.yml

@@ -174,6 +174,7 @@ jobs:
       - run: yarn --cwd packages/zone.js promisefinallytest
       - run: yarn --cwd packages/zone.js jest:test
       - run: yarn --cwd packages/zone.js jest:nodetest
+      - run: yarn --cwd packages/zone.js vitest:test
       - run: yarn --cwd packages/zone.js electrontest
       - run: yarn --cwd packages/zone.js/test/typings install --frozen-lockfile --non-interactive
       - run: yarn --cwd packages/zone.js/test/typings test

.github/workflows/pr.yml

@@ -185,6 +185,7 @@ jobs:
       - run: yarn --cwd packages/zone.js promisefinallytest
       - run: yarn --cwd packages/zone.js jest:test
       - run: yarn --cwd packages/zone.js jest:nodetest
+      - run: yarn --cwd packages/zone.js vitest:test
       - run: yarn --cwd packages/zone.js electrontest
       - run: yarn --cwd packages/zone.js/test/typings install --frozen-lockfile --non-interactive
       - run: yarn --cwd packages/zone.js/test/typings test

adev/src/content/tutorials/zoneless-migration/steps/4-make-tests-compatible-with-zoneless/src/app/test-helpers.ts

@@ -0,0 +1,17 @@
+import 'zone.js';
+import {fakeAsync as fakeAsyncInternal} from '@angular/core/testing';
+
+export function fakeAsync(fn: Function): Function {
+  return withProxyZone(fn);
+}
+export function withProxyZone(fn: Function): Function {
+  const autoProxyFn = function (this: unknown, ...args: any[]) {
+    const proxyZoneSpec = (Zone as any)['ProxyZoneSpec'];
+
+    const _sharedAutoProxyZoneSpec = new proxyZoneSpec();
+    const zone = Zone.root.fork(_sharedAutoProxyZoneSpec);
+
+    return zone.run(fakeAsyncInternal(fn), this, args);
+  };
+  return autoProxyFn;
+}

packages/zone.js/lib/zone-spec/fake-async-test.ts

@@ -7,6 +7,7 @@
  */
 
 import {ZoneType} from '../zone-impl';
+import {type ProxyZoneSpec} from './proxy';
 
 const global: any =
   (typeof window === 'object' && window) || (typeof self === 'object' && self) || globalThis.global;
@@ -789,15 +790,13 @@ class FakeAsyncTestZoneSpec implements ZoneSpec {
 
 let _fakeAsyncTestZoneSpec: FakeAsyncTestZoneSpec | null = null;
 
-type ProxyZoneSpecType = {
-  setDelegate(delegateSpec: ZoneSpec): void;
-  getDelegate(): ZoneSpec;
-  resetDelegate(): void;
-};
-function getProxyZoneSpec(): {get(): ProxyZoneSpecType; assertPresent: () => ProxyZoneSpecType} {
+function getProxyZoneSpec(): typeof ProxyZoneSpec | undefined {
   return Zone && (Zone as any)['ProxyZoneSpec'];
 }
 
+let _sharedProxyZoneSpec: ProxyZoneSpec | null = null;
+let _sharedProxyZone: Zone | null = null;
+
 /**
  * Clears out the shared fake async zone for a test.
  * To be called in a global `beforeEach`.
@@ -809,8 +808,8 @@ export function resetFakeAsyncZone() {
     _fakeAsyncTestZoneSpec.unlockDatePatch();
   }
   _fakeAsyncTestZoneSpec = null;
-  // in node.js testing we may not have ProxyZoneSpec in which case there is nothing to reset.
-  getProxyZoneSpec() && getProxyZoneSpec().assertPresent().resetDelegate();
+  getProxyZoneSpec()?.get()?.resetDelegate();
+  _sharedProxyZoneSpec?.resetDelegate();
 }
 
 /**
@@ -841,8 +840,8 @@ export function fakeAsync(fn: Function, options: {flush?: boolean} = {}): (...ar
     const ProxyZoneSpec = getProxyZoneSpec();
     if (!ProxyZoneSpec) {
       throw new Error(
-        'ProxyZoneSpec is needed for the async() test helper but could not be found. ' +
-          'Please make sure that your environment includes zone.js/plugins/proxy',
+        'ProxyZoneSpec is needed for the fakeAsync() test helper but could not be found. ' +
+          'Make sure that your environment includes zone-testing.js',
       );
     }
     const proxyZoneSpec = ProxyZoneSpec.assertPresent();
@@ -894,7 +893,7 @@ export function fakeAsync(fn: Function, options: {flush?: boolean} = {}): (...ar
       resetFakeAsyncZone();
     }
   };
-  (fakeAsyncFn as any).isFakeAsync = true;
+  fakeAsyncFn.isFakeAsync = true;
   return fakeAsyncFn;
 }
 
@@ -949,6 +948,51 @@ export function discardPeriodicTasks(): void {
   zoneSpec.pendingPeriodicTimers.length = 0;
 }
 
+/**
+ * Wraps a function to be executed in a shared ProxyZone.
+ *
+ * If no shared ProxyZone exists, one is created and reused for subsequent calls.
+ * Useful for wrapping test setup (beforeEach) and test execution (it) when test
+ * runner patching isn't available or desired for setting up the ProxyZone.
+ *
+ * @param fn The function to wrap.
+ * @returns A function that executes the original function within the shared ProxyZone.
+ *
+ * @experimental
+ */
+export function withProxyZone<T extends Function>(fn: T): T {
+  const autoProxyFn: any = function (this: unknown, ...args: any[]) {
+    const proxyZoneSpec = getProxyZoneSpec();
+    if (proxyZoneSpec === undefined) {
+      throw new Error(
+        'ProxyZoneSpec is needed for the withProxyZone() test helper but could not be found. ' +
+          'Make sure that your environment includes zone-testing.js',
+      );
+    }
+
+    const proxyZone = proxyZoneSpec.get() !== undefined ? Zone.current : getOrCreateRootProxy();
+    return proxyZone.run(fn, this, args);
+  };
+  return autoProxyFn as T;
+}
+
+function getOrCreateRootProxy() {
+  const ProxyZoneSpec = getProxyZoneSpec();
+  if (ProxyZoneSpec === undefined) {
+    throw new Error(
+      'ProxyZoneSpec is needed for withProxyZone but could not be found. ' +
+        'Make sure that your environment includes zone-testing.js',
+    );
+  }
+  // Ensure the shared ProxyZoneSpec instance exists
+  if (_sharedProxyZoneSpec === null) {
+    _sharedProxyZoneSpec = new ProxyZoneSpec() as ProxyZoneSpec;
+  }
+
+  _sharedProxyZone = Zone.root.fork(_sharedProxyZoneSpec);
+  return _sharedProxyZone;
+}
+
 /**
  * Flush any pending microtasks.
  *
@@ -973,6 +1017,7 @@ export function patchFakeAsyncTest(Zone: ZoneType): void {
         tick,
         flush,
         fakeAsync,
+        withProxyZone,
       };
     },
     true,

packages/zone.js/lib/zone-spec/proxy.ts

@@ -21,7 +21,7 @@ export class ProxyZoneSpec implements ZoneSpec {
 
   private tasks: Task[] = [];
 
-  static get(): ProxyZoneSpec {
+  static get(): ProxyZoneSpec | undefined {
     return Zone.current.get('ProxyZoneSpec');
   }
 
@@ -30,10 +30,11 @@ export class ProxyZoneSpec implements ZoneSpec {
   }
 
   static assertPresent(): ProxyZoneSpec {
-    if (!ProxyZoneSpec.isLoaded()) {
+    const spec = ProxyZoneSpec.get();
+    if (spec === undefined) {
       throw new Error(`Expected to be running in 'ProxyZone', but it was not found.`);
     }
-    return ProxyZoneSpec.get();
+    return spec;
   }
 
   constructor(private defaultSpecDelegate: ZoneSpec | null = null) {

packages/zone.js/package.json

@@ -17,12 +17,14 @@
     "jest-environment-node": "^29.0.3",
     "mocha": "^11.0.0",
     "mock-require": "3.0.3",
-    "tslib": "^2.3.0"
+    "tslib": "^2.3.0",
+    "vitest": "^3.1.3"
   },
   "scripts": {
     "electrontest": "cd test/extra && node electron.js",
     "jest:test": "jest --config ./test/jest/jest.config.js ./test/jest/jest.spec.js",
     "jest:nodetest": "jest --config ./test/jest/jest.node.config.js ./test/jest/jest.spec.js",
+    "vitest:test": "vitest ./test/vitest/vitest.spec.js",
     "promisefinallytest": "mocha ./test/promise/promise.finally.spec.mjs"
   },
   "repository": {

packages/zone.js/test/vitest/vitest.spec.js

@@ -0,0 +1,116 @@
+require('../../../../dist/bin/packages/zone.js/npm_package/bundles/zone.umd.js');
+require('../../../../dist/bin/packages/zone.js/npm_package/bundles/zone-testing.umd.js');
+
+import {expect, test, describe, beforeEach} from 'vitest';
+
+const {tick, withProxyZone, fakeAsync} = Zone[Zone.__symbol__('fakeAsyncTest')];
+
+describe('proxy zone behavior', () => {
+  const spec = new Zone['ProxyZoneSpec']();
+  const proxyZone = Zone.root.fork(spec);
+
+  function createForkedZone() {
+    const AsyncTestZoneSpec = Zone['AsyncTestZoneSpec'];
+    return Zone.current.fork(
+      new AsyncTestZoneSpec(
+        () => {},
+        () => {},
+        'asyncTest',
+      ),
+    );
+  }
+
+  test('cannot run fakeAsync outside proxy zone', () => {
+    expect(fakeAsync(() => {})).toThrow();
+  });
+
+  test('can run fakeAsync inside proxy zone', () => {
+    expect(() => {
+      proxyZone.run(fakeAsync(() => {}));
+    }).not.toThrow();
+  });
+
+  test('can flush timeouts in forked zone if created in proxy', () => {
+    let forkedZone;
+    proxyZone.run(() => {
+      forkedZone = createForkedZone();
+    });
+
+    proxyZone.run(
+      fakeAsync(() => {
+        let x = 1;
+        forkedZone.run(() => {
+          setTimeout(() => void (x = 2), 5000);
+        });
+
+        tick(5000);
+        expect(x).toBe(2);
+      }),
+    );
+  });
+
+  test('cannot flush timeouts in forked zone if created outside proxy', () => {
+    // This test is similar to creating a component in a beforeEach, which forks the zone to create the Angular NgZone
+    const forkedZone = createForkedZone();
+    proxyZone.run(() => {
+      fakeAsync(() => {
+        let x = 1;
+        forkedZone.run(() => {
+          setTimeout(() => void (x = 2), 5000);
+        });
+
+        tick(5000);
+        expect(x).toBe(1);
+      })();
+    });
+  });
+});
+
+test(
+  'withProxyZone runs inside proxy zone',
+  withProxyZone(() => {
+    expect(Zone.current.name).toEqual('ProxyZone');
+  }),
+);
+
+test(
+  'can use fakeAsync in proxy zone in test body',
+  withProxyZone(
+    fakeAsync(() => {
+      let x = 1;
+      setTimeout(() => void (x = 2), 5000);
+      tick(5000);
+      expect(x).toBe(2);
+    }),
+  ),
+);
+
+describe('can use withProxyZone and beforeEach', () => {
+  let forkedZone;
+  beforeEach(
+    withProxyZone(() => {
+      const AsyncTestZoneSpec = Zone['AsyncTestZoneSpec'];
+      forkedZone = Zone.current.fork(
+        new AsyncTestZoneSpec(
+          () => {},
+          () => {},
+          'asyncTest',
+        ),
+      );
+    }),
+  );
+
+  test(
+    'withProxyZone(fakeAsync)',
+    withProxyZone(
+      fakeAsync(() => {
+        let x = 1;
+        forkedZone.run(() => {
+          setTimeout(() => void (x = 2), 5000);
+        });
+        tick(5000);
+        expect(x).toBe(2);
+      }),
+    ),
+  );
+});