test_runner: support test plans

Co-Authored-By: Marco Ippolito <marcoippolito54@gmail.com>
PR-URL: #52860
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
Reviewed-By: Paolo Insogna <paolo@cowtech.it>

Author: cjihrig

doc/api/test.md

@@ -1325,6 +1325,10 @@ changes:
   * `timeout` {number} A number of milliseconds the test will fail after.
     If unspecified, subtests inherit this value from their parent.
     **Default:** `Infinity`.
+  * `plan` {number} The number of assertions and subtests expected to be run in the test.
+    If the number of assertions run in the test does not match the number
+    specified in the plan, the test will fail.
+    **Default:** `undefined`.
 * `fn` {Function|AsyncFunction} The function under test. The first argument
   to this function is a [`TestContext`][] object. If the test uses callbacks,
   the callback function is passed as the second argument. **Default:** A no-op
@@ -2912,6 +2916,54 @@ added:
 
 The name of the test.
 
+### `context.plan(count)`
+
+<!-- YAML
+added:
+  - REPLACEME
+-->
+
+> Stability: 1 - Experimental
+
+* `count` {number} The number of assertions and subtests that are expected to run.
+
+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
+expected count, the test will fail.
+
+> Note: To make sure assertions are tracked, `t.assert` must be used instead of `assert` directly.
+
+```js
+test('top level test', (t) => {
+  t.plan(2);
+  t.assert.ok('some relevant assertion here');
+  t.subtest('subtest', () => {});
+});
+```
+
+When working with asynchronous code, the `plan` function can be used to ensure that the
+correct number of assertions are run:
+
+```js
+test('planning with streams', (t, done) => {
+  function* generate() {
+    yield 'a';
+    yield 'b';
+    yield 'c';
+  }
+  const expected = ['a', 'b', 'c'];
+  t.plan(expected.length);
+  const stream = Readable.from(generate());
+  stream.on('data', (chunk) => {
+    t.assert.strictEqual(chunk, expected.shift());
+  });
+
+  stream.on('end', () => {
+    done();
+  });
+});
+```
+
 ### `context.runOnly(shouldRunOnlyTests)`
 
 <!-- YAML
@@ -3042,6 +3094,10 @@ changes:
   * `timeout` {number} A number of milliseconds the test will fail after.
     If unspecified, subtests inherit this value from their parent.
     **Default:** `Infinity`.
+  * `plan` {number} The number of assertions and subtests expected to be run in the test.
+    If the number of assertions run in the test does not match the number
+    specified in the plan, the test will fail.
+    **Default:** `undefined`.
 * `fn` {Function|AsyncFunction} The function under test. The first argument
   to this function is a [`TestContext`][] object. If the test uses callbacks,
   the callback function is passed as the second argument. **Default:** A no-op
@@ -3055,7 +3111,7 @@ behaves in the same fashion as the top level [`test()`][] function.
 test('top level test', async (t) => {
   await t.test(
     'This is a subtest',
-    { only: false, skip: false, concurrency: 1, todo: false },
+    { only: false, skip: false, concurrency: 1, todo: false, plan: 4 },
     (t) => {
       assert.ok('some relevant assertion here');
     },

lib/internal/test_runner/runner.js

@@ -509,6 +509,7 @@ function run(options) {
     watch,
     setup,
     only,
+    plan,
   } = options;
 
   if (files != null) {
@@ -565,7 +566,7 @@ function run(options) {
     });
   }
 
-  const root = createTestTree({ __proto__: null, concurrency, timeout, signal });
+  const root = createTestTree({ __proto__: null, concurrency, timeout, signal, plan });
   root.harness.shouldColorizeTestFiles ||= shouldColorizeTestFiles(root);
 
   if (process.env.NODE_TEST_CONTEXT !== undefined) {

lib/internal/test_runner/test.js

@@ -11,6 +11,7 @@ const {
   FunctionPrototype,
   MathMax,
   Number,
+  ObjectEntries,
   ObjectSeal,
   PromisePrototypeThen,
   PromiseResolve,
@@ -85,6 +86,7 @@ const {
   testOnlyFlag,
 } = parseCommandLine();
 let kResistStopPropagation;
+let assertObj;
 let findSourceMap;
 
 const kRunOnceOptions = { __proto__: null, preserveReturnValue: true };
@@ -97,6 +99,19 @@ function lazyFindSourceMap(file) {
   return findSourceMap(file);
 }
 
+function lazyAssertObject() {
+  if (assertObj === undefined) {
+    assertObj = new SafeMap();
+    const assert = require('assert');
+    for (const { 0: key, 1: value } of ObjectEntries(assert)) {
+      if (typeof value === 'function') {
+        assertObj.set(value, key);
+      }
+    }
+  }
+  return assertObj;
+}
+
 function stopTest(timeout, signal) {
   const deferred = createDeferredPromise();
   const abortListener = addAbortListener(signal, deferred.resolve);
@@ -136,7 +151,25 @@ function stopTest(timeout, signal) {
   return deferred.promise;
 }
 
+class TestPlan {
+  constructor(count) {
+    validateUint32(count, 'count', 0);
+    this.expected = count;
+    this.actual = 0;
+  }
+
+  check() {
+    if (this.actual !== this.expected) {
+      throw new ERR_TEST_FAILURE(
+        `plan expected ${this.expected} assertions but received ${this.actual}`,
+        kTestCodeFailure,
+      );
+    }
+  }
+}
+
 class TestContext {
+  #assert;
   #test;
 
   constructor(test) {
@@ -163,6 +196,36 @@ class TestContext {
     this.#test.diagnostic(message);
   }
 
+  plan(count) {
+    if (this.#test.plan !== null) {
+      throw new ERR_TEST_FAILURE(
+        'cannot set plan more than once',
+        kTestCodeFailure,
+      );
+    }
+
+    this.#test.plan = new TestPlan(count);
+  }
+
+  get assert() {
+    if (this.#assert === undefined) {
+      const { plan } = this.#test;
+      const assertions = lazyAssertObject();
+      const assert = { __proto__: null };
+
+      this.#assert = assert;
+      for (const { 0: method, 1: name } of assertions.entries()) {
+        assert[name] = (...args) => {
+          if (plan !== null) {
+            plan.actual++;
+          }
+          return ReflectApply(method, assert, args);
+        };
+      }
+    }
+    return this.#assert;
+  }
+
   get mock() {
     this.#test.mock ??= new MockTracker();
     return this.#test.mock;
@@ -186,6 +249,11 @@ class TestContext {
       loc: getCallerLocation(),
     };
 
+    const { plan } = this.#test;
+    if (plan !== null) {
+      plan.actual++;
+    }
+
     const subtest = this.#test.createSubtest(
       // eslint-disable-next-line no-use-before-define
       Test, name, options, fn, overrides,
@@ -240,7 +308,7 @@ class Test extends AsyncResource {
     super('Test');
 
     let { fn, name, parent, skip } = options;
-    const { concurrency, loc, only, timeout, todo, signal } = options;
+    const { concurrency, loc, only, timeout, todo, signal, plan } = options;
 
     if (typeof fn !== 'function') {
       fn = noop;
@@ -351,6 +419,8 @@ class Test extends AsyncResource {
     this.fn = fn;
     this.harness = null; // Configured on the root test by the test harness.
     this.mock = null;
+    this.plan = null;
+    this.expectedAssertions = plan;
     this.cancelled = false;
     this.skipped = skip !== undefined && skip !== false;
     this.isTodo = todo !== undefined && todo !== false;
@@ -643,6 +713,11 @@ class Test extends AsyncResource {
 
     const hookArgs = this.getRunArgs();
     const { args, ctx } = hookArgs;
+
+    if (this.plan === null && this.expectedAssertions) {
+      ctx.plan(this.expectedAssertions);
+    }
+
     const after = async () => {
       if (this.hooks.after.length > 0) {
         await this.runHook('after', hookArgs);
@@ -694,7 +769,7 @@ class Test extends AsyncResource {
         this.postRun();
         return;
       }
-
+      this.plan?.check();
       this.pass();
       await afterEach();
       await after();

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

@@ -0,0 +1,79 @@
+'use strict';
+const { test } = require('node:test');
+const { Readable } = require('node:stream');
+
+test('test planning basic', (t) => {
+  t.plan(2);
+  t.assert.ok(true);
+  t.assert.ok(true);
+});
+
+test('less assertions than planned', (t) => {
+  t.plan(1);
+});
+
+test('more assertions than planned', (t) => {
+  t.plan(1);
+  t.assert.ok(true);
+  t.assert.ok(true);
+});
+
+test('subtesting', (t) => {
+  t.plan(1);
+  t.test('subtest', () => { });
+});
+
+test('subtesting correctly', (t) => {
+  t.plan(2);
+  t.assert.ok(true);
+  t.test('subtest', (st) => {
+    st.plan(1);
+    st.assert.ok(true);
+  });
+});
+
+test('correctly ignoring subtesting plan', (t) => {
+  t.plan(1);
+  t.test('subtest', (st) => {
+    st.plan(1);
+    st.assert.ok(true);
+  });
+});
+
+test('failing planning by options', { plan: 1 }, () => {
+});
+
+test('not failing planning by options', { plan: 1 }, (t) => {
+  t.assert.ok(true);
+});
+
+test('subtest planning by options', (t) => {
+  t.test('subtest', { plan: 1 }, (st) => {
+    st.assert.ok(true);
+  });
+});
+
+test('failing more assertions than planned', (t) => {
+  t.plan(2);
+  t.assert.ok(true);
+  t.assert.ok(true);
+  t.assert.ok(true);
+});
+
+test('planning with streams', (t, done) => {
+  function* generate() {
+    yield 'a';
+    yield 'b';
+    yield 'c';
+  }
+  const expected = ['a', 'b', 'c'];
+  t.plan(expected.length);
+  const stream = Readable.from(generate());
+  stream.on('data', (chunk) => {
+    t.assert.strictEqual(chunk, expected.shift());
+  });
+
+  stream.on('end', () => {
+   done();
+  });
+})