test_runner: refactor testPlan counter increse

PR-URL: #56765
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
Reviewed-By: Colin Ihrig <cjihrig@gmail.com>
Reviewed-By: Chemi Atlow <chemi@atlow.co.il>

Author: pmarchini

doc/api/test.md

@@ -3398,19 +3398,33 @@ added:
 
 The name of the test.
 
-### `context.plan(count)`
+### `context.plan(count[,options])`
 
 <!-- YAML
 added:
   - v22.2.0
   - v20.15.0
 changes:
+  - version:
+    - REPLACEME
+    pr-url: https://github.com/nodejs/node/pull/56765
+    description: Add the `options` parameter.
   - version: v23.4.0
     pr-url: https://github.com/nodejs/node/pull/55895
     description: This function is no longer experimental.
 -->
 
 * `count` {number} The number of assertions and subtests that are expected to run.
+* `options` {Object} Additional options for the plan.
+  * `wait` {boolean|number} The wait time for the plan:
+    * If `true`, the plan waits indefinitely for all assertions and subtests to run.
+    * If `false`, the plan performs an immediate check after the test function completes,
+      without waiting for any pending assertions or subtests.
+      Any assertions or subtests that complete after this check will not be counted towards the plan.
+    * If a number, it specifies the maximum wait time in milliseconds
+      before timing out while waiting for expected assertions and subtests to be matched.
+      If the timeout is reached, the test will fail.
+      **Default:** `false`.
 
 This function is used to set the number of assertions and subtests that are expected to run
 within the test. If the number of assertions and subtests that run does not match the
@@ -3449,6 +3463,26 @@ test('planning with streams', (t, done) => {
 });
 ```
 
+When using the `wait` option, you can control how long the test will wait for the expected assertions.
+For example, setting a maximum wait time ensures that the test will wait for asynchronous assertions
+to complete within the specified timeframe:
+
+```js
+test('plan with wait: 2000 waits for async assertions', (t) => {
+  t.plan(1, { wait: 2000 }); // Waits for up to 2 seconds for the assertion to complete.
+
+  const asyncActivity = () => {
+    setTimeout(() => {
+      t.assert.ok(true, 'Async assertion completed within the wait time');
+    }, 1000); // Completes after 1 second, within the 2-second wait time.
+  };
+
+  asyncActivity(); // The test will pass because the assertion is completed in time.
+});
+```
+
+Note: If a `wait` timeout is specified, it begins counting down only after the test function finishes executing.
+
 ### `context.runOnly(shouldRunOnlyTests)`
 
 <!-- YAML

lib/internal/test_runner/test.js

@@ -176,22 +176,88 @@ function testMatchesPattern(test, patterns) {
 }
 
 class TestPlan {
-  constructor(count) {
+  #waitIndefinitely = false;
+  #planPromise = null;
+  #timeoutId = null;
+
+  constructor(count, options = kEmptyObject) {
     validateUint32(count, 'count');
+    validateObject(options, 'options');
     this.expected = count;
     this.actual = 0;
+
+    const { wait } = options;
+    if (typeof wait === 'boolean') {
+      this.wait = wait;
+      this.#waitIndefinitely = wait;
+    } else if (typeof wait === 'number') {
+      validateNumber(wait, 'options.wait', 0, TIMEOUT_MAX);
+      this.wait = wait;
+    } else if (wait !== undefined) {
+      throw new ERR_INVALID_ARG_TYPE('options.wait', ['boolean', 'number'], wait);
+    }
+  }
+
+  #planMet() {
+    return this.actual === this.expected;
+  }
+
+  #createTimeout(reject) {
+    return setTimeout(() => {
+      const err = new ERR_TEST_FAILURE(
+        `plan timed out after ${this.wait}ms with ${this.actual} assertions when expecting ${this.expected}`,
+        kTestTimeoutFailure,
+      );
+      reject(err);
+    }, this.wait);
   }
 
   check() {
-    if (this.actual !== this.expected) {
+    if (this.#planMet()) {
+      if (this.#timeoutId) {
+        clearTimeout(this.#timeoutId);
+        this.#timeoutId = null;
+      }
+      if (this.#planPromise) {
+        const { resolve } = this.#planPromise;
+        resolve();
+        this.#planPromise = null;
+      }
+      return;
+    }
+
+    if (!this.#shouldWait()) {
       throw new ERR_TEST_FAILURE(
         `plan expected ${this.expected} assertions but received ${this.actual}`,
         kTestCodeFailure,
       );
     }
+
+    if (!this.#planPromise) {
+      const { promise, resolve, reject } = PromiseWithResolvers();
+      this.#planPromise = { __proto__: null, promise, resolve, reject };
+
+      if (!this.#waitIndefinitely) {
+        this.#timeoutId = this.#createTimeout(reject);
+      }
+    }
+
+    return this.#planPromise.promise;
+  }
+
+  count() {
+    this.actual++;
+    if (this.#planPromise) {
+      this.check();
+    }
+  }
+
+  #shouldWait() {
+    return this.wait !== undefined && this.wait !== false;
   }
 }
 
+
 class TestContext {
   #assert;
   #test;
@@ -228,15 +294,15 @@ class TestContext {
     this.#test.diagnostic(message);
   }
 
-  plan(count) {
+  plan(count, options = kEmptyObject) {
     if (this.#test.plan !== null) {
       throw new ERR_TEST_FAILURE(
         'cannot set plan more than once',
         kTestCodeFailure,
       );
     }
 
-    this.#test.plan = new TestPlan(count);
+    this.#test.plan = new TestPlan(count, options);
   }
 
   get assert() {
@@ -249,7 +315,7 @@ class TestContext {
       map.forEach((method, name) => {
         assert[name] = (...args) => {
           if (plan !== null) {
-            plan.actual++;
+            plan.count();
           }
           return ReflectApply(method, this, args);
         };
@@ -260,7 +326,7 @@ class TestContext {
         // stacktrace from the correct starting point.
         function ok(...args) {
           if (plan !== null) {
-            plan.actual++;
+            plan.count();
           }
           innerOk(ok, args.length, ...args);
         }
@@ -296,7 +362,7 @@ class TestContext {
 
     const { plan } = this.#test;
     if (plan !== null) {
-      plan.actual++;
+      plan.count();
     }
 
     const subtest = this.#test.createSubtest(
@@ -959,30 +1025,46 @@ class Test extends AsyncResource {
       const runArgs = ArrayPrototypeSlice(args);
       ArrayPrototypeUnshift(runArgs, this.fn, ctx);
 
+      const promises = [];
       if (this.fn.length === runArgs.length - 1) {
-        // This test is using legacy Node.js error first callbacks.
+        // This test is using legacy Node.js error-first callbacks.
         const { promise, cb } = createDeferredCallback();
-
         ArrayPrototypePush(runArgs, cb);
+
         const ret = ReflectApply(this.runInAsyncScope, this, runArgs);
 
         if (isPromise(ret)) {
           this.fail(new ERR_TEST_FAILURE(
             'passed a callback but also returned a Promise',
             kCallbackAndPromisePresent,
           ));
-          await SafePromiseRace([ret, stopPromise]);
+          ArrayPrototypePush(promises, ret);
         } else {
-          await SafePromiseRace([PromiseResolve(promise), stopPromise]);
+          ArrayPrototypePush(promises, PromiseResolve(promise));
         }
       } else {
         // This test is synchronous or using Promises.
         const promise = ReflectApply(this.runInAsyncScope, this, runArgs);
-        await SafePromiseRace([PromiseResolve(promise), stopPromise]);
+        ArrayPrototypePush(promises, PromiseResolve(promise));
       }
 
+      ArrayPrototypePush(promises, stopPromise);
+
+      // Wait for the race to finish
+      await SafePromiseRace(promises);
+
       this[kShouldAbort]();
       this.plan?.check();
+
+      if (this.plan !== null) {
+        const planPromise = this.plan?.check();
+        // If the plan returns a promise, it means that it is waiting for more assertions to be made before
+        // continuing.
+        if (planPromise) {
+          await SafePromiseRace([planPromise, stopPromise]);
+        }
+      }
+
       this.pass();
       await afterEach();
       await after();

test/fixtures/test-runner/output/test-runner-plan-timeout.js

@@ -0,0 +1,76 @@
+'use strict';
+const { describe, it } = require('node:test');
+
+describe('planning with wait', () => {
+  it('planning with wait and passing', async (t) => {
+    t.plan(1, { wait: 5000 });
+
+    const asyncActivity = () => {
+      setTimeout(() => {
+        t.assert.ok(true);
+      }, 250);
+    };
+
+    asyncActivity();
+  });
+
+  it('planning with wait and failing', async (t) => {
+    t.plan(1, { wait: 5000 });
+
+    const asyncActivity = () => {
+      setTimeout(() => {
+        t.assert.ok(false);
+      }, 250);
+    };
+
+    asyncActivity();
+  });
+
+  it('planning wait time expires before plan is met', async (t) => {
+    t.plan(2, { wait: 500 });
+
+    const asyncActivity = () => {
+      setTimeout(() => {
+        t.assert.ok(true);
+      }, 50_000_000);
+    };
+
+    asyncActivity();
+  });
+
+  it(`planning with wait "options.wait : true" and passing`, async (t) => {
+    t.plan(1, { wait: true });
+
+    const asyncActivity = () => {
+      setTimeout(() => {
+        t.assert.ok(true);
+      }, 250);
+    };
+
+    asyncActivity();
+  });
+
+  it(`planning with wait "options.wait : true" and failing`, async (t) => {
+    t.plan(1, { wait: true });
+
+    const asyncActivity = () => {
+      setTimeout(() => {
+        t.assert.ok(false);
+      }, 250);
+    };
+
+    asyncActivity();
+  });
+
+  it(`planning with wait "options.wait : false" should not wait`, async (t) => {
+    t.plan(1, { wait: false });
+
+    const asyncActivity = () => {
+      setTimeout(() => {
+        t.assert.ok(true);
+      }, 500_000);
+    };
+
+    asyncActivity();
+  })
+});