test_runner: refactor testPlan counter increse

PR-URL: https://github.com/nodejs/node/pull/56765
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
Reviewed-By: Colin Ihrig <cjihrig@gmail.com>
Reviewed-By: Chemi Atlow <chemi@atlow.co.il>

doc/api/test.md

@@ -3381,13 +3381,17 @@ 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
     - v22.13.0
@@ -3396,6 +3400,16 @@ changes:
 -->
 
 * `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
@@ -3434,6 +3448,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,7 +294,7 @@ 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',
@@ -236,7 +302,7 @@ class TestContext {
       );
     }
 
-    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(
@@ -968,11 +1034,12 @@ 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)) {
@@ -980,23 +1047,36 @@ class Test extends AsyncResource {
             '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]();
 
       if (this.subtestsPromise !== null) {
         await SafePromiseRace([this.subtestsPromise.promise, stopPromise]);
       }
 
-      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();
+  })
+});

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

@@ -0,0 +1,97 @@
+TAP version 13
+# Subtest: planning with wait
+    # Subtest: planning with wait and passing
+    ok 1 - planning with wait and passing
+      ---
+      duration_ms: *
+      type: 'test'
+      ...
+    # Subtest: planning with wait and failing
+    not ok 2 - planning with wait and failing
+      ---
+      duration_ms: *
+      type: 'test'
+      location: '/test/fixtures/test-runner/output/test-runner-plan-timeout.js:(LINE):3'
+      failureType: 'uncaughtException'
+      error: |-
+        The expression evaluated to a falsy value:
+        
+          t.assert.ok(false)
+        
+      code: 'ERR_ASSERTION'
+      name: 'AssertionError'
+      expected: true
+      actual: false
+      operator: '=='
+      stack: |-
+        *
+        *
+        *
+      ...
+    # Subtest: planning wait time expires before plan is met
+    not ok 3 - planning wait time expires before plan is met
+      ---
+      duration_ms: *
+      type: 'test'
+      location: '/test/fixtures/test-runner/output/test-runner-plan-timeout.js:(LINE):3'
+      failureType: 'testTimeoutFailure'
+      error: 'plan timed out after 500ms with 0 assertions when expecting 2'
+      code: 'ERR_TEST_FAILURE'
+      ...
+    # Subtest: planning with wait "options.wait : true" and passing
+    ok 4 - planning with wait "options.wait : true" and passing
+      ---
+      duration_ms: *
+      type: 'test'
+      ...
+    # Subtest: planning with wait "options.wait : true" and failing
+    not ok 5 - planning with wait "options.wait : true" and failing
+      ---
+      duration_ms: *
+      type: 'test'
+      location: '/test/fixtures/test-runner/output/test-runner-plan-timeout.js:(LINE):3'
+      failureType: 'uncaughtException'
+      error: |-
+        The expression evaluated to a falsy value:
+        
+          t.assert.ok(false)
+        
+      code: 'ERR_ASSERTION'
+      name: 'AssertionError'
+      expected: true
+      actual: false
+      operator: '=='
+      stack: |-
+        *
+        *
+        *
+      ...
+    # Subtest: planning with wait "options.wait : false" should not wait
+    not ok 6 - planning with wait "options.wait : false" should not wait
+      ---
+      duration_ms: *
+      type: 'test'
+      location: '/test/fixtures/test-runner/output/test-runner-plan-timeout.js:(LINE):3'
+      failureType: 'testCodeFailure'
+      error: 'plan expected 1 assertions but received 0'
+      code: 'ERR_TEST_FAILURE'
+      ...
+    1..6
+not ok 1 - planning with wait
+  ---
+  duration_ms: *
+  type: 'suite'
+  location: '/test/fixtures/test-runner/output/test-runner-plan-timeout.js:(LINE):1'
+  failureType: 'subtestsFailed'
+  error: '4 subtests failed'
+  code: 'ERR_TEST_FAILURE'
+  ...
+1..1
+# tests 6
+# suites 1
+# pass 2
+# fail 3
+# cancelled 1
+# skipped 0
+# todo 0
+# duration_ms *

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

@@ -1,7 +1,36 @@
 'use strict';
-const { test } = require('node:test');
+const { test, suite } = require('node:test');
 const { Readable } = require('node:stream');
 
+suite('input validation', () => {
+  test('throws if options is not an object', (t) => {
+    t.assert.throws(() => {
+      t.plan(1, null);
+    }, {
+      code: 'ERR_INVALID_ARG_TYPE',
+      message: /The "options" argument must be of type object/,
+    });
+  });
+
+  test('throws if options.wait is not a number or a boolean', (t) => {
+    t.assert.throws(() => {
+      t.plan(1, { wait: 'foo' });
+    }, {
+      code: 'ERR_INVALID_ARG_TYPE',
+      message: /The "options\.wait" property must be one of type boolean or number\. Received type string/,
+    });
+  });
+
+  test('throws if count is not a number', (t) => {
+    t.assert.throws(() => {
+      t.plan('foo');
+    }, {
+      code: 'ERR_INVALID_ARG_TYPE',
+      message: /The "count" argument must be of type number/,
+    });
+  });
+});
+
 test('test planning basic', (t) => {
   t.plan(2);
   t.assert.ok(true);
@@ -76,4 +105,4 @@ test('planning with streams', (t, done) => {
   stream.on('end', () => {
     done();
   });
-})
+});

test/fixtures/test-runner/output/test-runner-plan.snapshot

@@ -1,12 +1,37 @@
 TAP version 13
+# Subtest: input validation
+    # Subtest: throws if options is not an object
+    ok 1 - throws if options is not an object
+      ---
+      duration_ms: *
+      type: 'test'
+      ...
+    # Subtest: throws if options.wait is not a number or a boolean
+    ok 2 - throws if options.wait is not a number or a boolean
+      ---
+      duration_ms: *
+      type: 'test'
+      ...
+    # Subtest: throws if count is not a number
+    ok 3 - throws if count is not a number
+      ---
+      duration_ms: *
+      type: 'test'
+      ...
+    1..3
+ok 1 - input validation
+  ---
+  duration_ms: *
+  type: 'suite'
+  ...
 # Subtest: test planning basic
-ok 1 - test planning basic
+ok 2 - test planning basic
   ---
   duration_ms: *
   type: 'test'
   ...
 # Subtest: less assertions than planned
-not ok 2 - less assertions than planned
+not ok 3 - less assertions than planned
   ---
   duration_ms: *
   type: 'test'
@@ -16,7 +41,7 @@ not ok 2 - less assertions than planned
   code: 'ERR_TEST_FAILURE'
   ...
 # Subtest: more assertions than planned
-not ok 3 - more assertions than planned
+not ok 4 - more assertions than planned
   ---
   duration_ms: *
   type: 'test'
@@ -33,7 +58,7 @@ not ok 3 - more assertions than planned
       type: 'test'
       ...
     1..1
-ok 4 - subtesting
+ok 5 - subtesting
   ---
   duration_ms: *
   type: 'test'
@@ -46,7 +71,7 @@ ok 4 - subtesting
       type: 'test'
       ...
     1..1
-ok 5 - subtesting correctly
+ok 6 - subtesting correctly
   ---
   duration_ms: *
   type: 'test'
@@ -59,13 +84,13 @@ ok 5 - subtesting correctly
       type: 'test'
       ...
     1..1
-ok 6 - correctly ignoring subtesting plan
+ok 7 - correctly ignoring subtesting plan
   ---
   duration_ms: *
   type: 'test'
   ...
 # Subtest: failing planning by options
-not ok 7 - failing planning by options
+not ok 8 - failing planning by options
   ---
   duration_ms: *
   type: 'test'
@@ -75,7 +100,7 @@ not ok 7 - failing planning by options
   code: 'ERR_TEST_FAILURE'
   ...
 # Subtest: not failing planning by options
-ok 8 - not failing planning by options
+ok 9 - not failing planning by options
   ---
   duration_ms: *
   type: 'test'
@@ -88,13 +113,13 @@ ok 8 - not failing planning by options
       type: 'test'
       ...
     1..1
-ok 9 - subtest planning by options
+ok 10 - subtest planning by options
   ---
   duration_ms: *
   type: 'test'
   ...
 # Subtest: failing more assertions than planned
-not ok 10 - failing more assertions than planned
+not ok 11 - failing more assertions than planned
   ---
   duration_ms: *
   type: 'test'
@@ -104,15 +129,15 @@ not ok 10 - failing more assertions than planned
   code: 'ERR_TEST_FAILURE'
   ...
 # Subtest: planning with streams
-ok 11 - planning with streams
+ok 12 - planning with streams
   ---
   duration_ms: *
   type: 'test'
   ...
-1..11
-# tests 15
-# suites 0
-# pass 11
+1..12
+# tests 18
+# suites 1
+# pass 14
 # fail 4
 # cancelled 0
 # skipped 0

test/parallel/test-runner-output.mjs

@@ -239,6 +239,10 @@ const tests = [
     name: 'test-runner/output/test-runner-watch-spec.mjs',
     transform: specTransform,
   },
+  {
+    name: 'test-runner/output/test-runner-plan-timeout.js',
+    flags: ['--test-reporter=tap', '--test-force-exit'],
+  },
   process.features.inspector ? {
     name: 'test-runner/output/coverage_failure.js',
     flags: ['--test-reporter=tap', '--test-coverage-exclude=!test/**'],