From 2f59529dc50f1d8b0dd69fa22bdf0cf8db4cbaa7 Mon Sep 17 00:00:00 2001
From: Colin Ihrig <cjihrig@gmail.com>
Date: Thu, 9 May 2024 11:48:13 +0200
Subject: [PATCH] test_runner: support test plans

Co-Authored-By: Marco Ippolito <marcoippolito54@gmail.com>
PR-URL: https://github.com/nodejs/node/pull/52860
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
Reviewed-By: Paolo Insogna <paolo@cowtech.it>
---
 doc/api/test.md                               |  58 +++++++++-
 lib/internal/test_runner/runner.js            |   3 +-
 lib/internal/test_runner/test.js              |  79 ++++++++++++-
 .../test-runner/output/test-runner-plan.js    |  79 +++++++++++++
 .../output/test-runner-plan.snapshot          | 105 ++++++++++++++++++
 test/parallel/test-runner-output.mjs          |   1 +
 6 files changed, 321 insertions(+), 4 deletions(-)
 create mode 100644 test/fixtures/test-runner/output/test-runner-plan.js
 create mode 100644 test/fixtures/test-runner/output/test-runner-plan.snapshot

diff --git a/doc/api/test.md b/doc/api/test.md
index d86c2cfc0bfb13..8d1e5a4624ce16 100644
--- a/doc/api/test.md
+++ b/doc/api/test.md
@@ -1364,6 +1364,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
@@ -2965,6 +2969,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
@@ -3095,6 +3147,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
@@ -3108,7 +3164,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');
     },
diff --git a/lib/internal/test_runner/runner.js b/lib/internal/test_runner/runner.js
index 8310e13248898a..d2eb0cb6abf7a7 100644
--- a/lib/internal/test_runner/runner.js
+++ b/lib/internal/test_runner/runner.js
@@ -462,6 +462,7 @@ function run(options = kEmptyObject) {
     watch,
     setup,
     only,
+    plan,
   } = options;
 
   if (files != null) {
@@ -534,7 +535,7 @@ function run(options = kEmptyObject) {
     });
   }
 
-  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) {
diff --git a/lib/internal/test_runner/test.js b/lib/internal/test_runner/test.js
index 737f45a4163799..da502812f52c3b 100644
--- a/lib/internal/test_runner/test.js
+++ b/lib/internal/test_runner/test.js
@@ -12,6 +12,7 @@ const {
   MathMax,
   Number,
   ObjectDefineProperty,
+  ObjectEntries,
   ObjectSeal,
   PromisePrototypeThen,
   PromiseResolve,
@@ -88,6 +89,7 @@ const {
   testOnlyFlag,
 } = parseCommandLine();
 let kResistStopPropagation;
+let assertObj;
 let findSourceMap;
 let noopTestStream;
 
@@ -101,6 +103,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);
@@ -153,7 +168,25 @@ function testMatchesPattern(test, patterns) {
   );
 }
 
+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) {
@@ -180,6 +213,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;
@@ -203,6 +266,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,
@@ -257,7 +325,7 @@ class Test extends AsyncResource {
     super('Test');
 
     let { fn, name, parent } = options;
-    const { concurrency, loc, only, timeout, todo, skip, signal } = options;
+    const { concurrency, loc, only, timeout, todo, skip, signal, plan } = options;
 
     if (typeof fn !== 'function') {
       fn = noop;
@@ -373,6 +441,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;
@@ -703,6 +773,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);
@@ -754,7 +829,7 @@ class Test extends AsyncResource {
         this.postRun();
         return;
       }
-
+      this.plan?.check();
       this.pass();
       await afterEach();
       await after();
diff --git a/test/fixtures/test-runner/output/test-runner-plan.js b/test/fixtures/test-runner/output/test-runner-plan.js
new file mode 100644
index 00000000000000..0d6cc8ed05cdd8
--- /dev/null
+++ b/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();
+  });
+})
diff --git a/test/fixtures/test-runner/output/test-runner-plan.snapshot b/test/fixtures/test-runner/output/test-runner-plan.snapshot
new file mode 100644
index 00000000000000..b172c2da05a884
--- /dev/null
+++ b/test/fixtures/test-runner/output/test-runner-plan.snapshot
@@ -0,0 +1,105 @@
+TAP version 13
+# Subtest: test planning basic
+ok 1 - test planning basic
+  ---
+  duration_ms: *
+  ...
+# Subtest: less assertions than planned
+not ok 2 - less assertions than planned
+  ---
+  duration_ms: *
+  location: '/test/fixtures/test-runner/output/test-runner-plan.js:(LINE):1'
+  failureType: 'testCodeFailure'
+  error: 'plan expected 1 assertions but received 0'
+  code: 'ERR_TEST_FAILURE'
+  ...
+# Subtest: more assertions than planned
+not ok 3 - more assertions than planned
+  ---
+  duration_ms: *
+  location: '/test/fixtures/test-runner/output/test-runner-plan.js:(LINE):1'
+  failureType: 'testCodeFailure'
+  error: 'plan expected 1 assertions but received 2'
+  code: 'ERR_TEST_FAILURE'
+  ...
+# Subtest: subtesting
+    # Subtest: subtest
+    ok 1 - subtest
+      ---
+      duration_ms: *
+      ...
+    1..1
+ok 4 - subtesting
+  ---
+  duration_ms: *
+  ...
+# Subtest: subtesting correctly
+    # Subtest: subtest
+    ok 1 - subtest
+      ---
+      duration_ms: *
+      ...
+    1..1
+ok 5 - subtesting correctly
+  ---
+  duration_ms: *
+  ...
+# Subtest: correctly ignoring subtesting plan
+    # Subtest: subtest
+    ok 1 - subtest
+      ---
+      duration_ms: *
+      ...
+    1..1
+ok 6 - correctly ignoring subtesting plan
+  ---
+  duration_ms: *
+  ...
+# Subtest: failing planning by options
+not ok 7 - failing planning by options
+  ---
+  duration_ms: *
+  location: '/test/fixtures/test-runner/output/test-runner-plan.js:(LINE):1'
+  failureType: 'testCodeFailure'
+  error: 'plan expected 1 assertions but received 0'
+  code: 'ERR_TEST_FAILURE'
+  ...
+# Subtest: not failing planning by options
+ok 8 - not failing planning by options
+  ---
+  duration_ms: *
+  ...
+# Subtest: subtest planning by options
+    # Subtest: subtest
+    ok 1 - subtest
+      ---
+      duration_ms: *
+      ...
+    1..1
+ok 9 - subtest planning by options
+  ---
+  duration_ms: *
+  ...
+# Subtest: failing more assertions than planned
+not ok 10 - failing more assertions than planned
+  ---
+  duration_ms: *
+  location: '/test/fixtures/test-runner/output/test-runner-plan.js:(LINE):1'
+  failureType: 'testCodeFailure'
+  error: 'plan expected 2 assertions but received 3'
+  code: 'ERR_TEST_FAILURE'
+  ...
+# Subtest: planning with streams
+ok 11 - planning with streams
+  ---
+  duration_ms: *
+  ...
+1..11
+# tests 15
+# suites 0
+# pass 11
+# fail 4
+# cancelled 0
+# skipped 0
+# todo 0
+# duration_ms *
diff --git a/test/parallel/test-runner-output.mjs b/test/parallel/test-runner-output.mjs
index d65e14de1fb0cb..a3f7354f24b8ab 100644
--- a/test/parallel/test-runner-output.mjs
+++ b/test/parallel/test-runner-output.mjs
@@ -140,6 +140,7 @@ const tests = [
       replaceTestDuration,
     ),
   },
+  { name: 'test-runner/output/test-runner-plan.js' },
   process.features.inspector ? { name: 'test-runner/output/coverage_failure.js' } : false,
 ]
 .filter(Boolean)