From 8bd4c3036aba5f1ebaa2420ffd27ad9fbc9138a0 Mon Sep 17 00:00:00 2001
From: Marco Ippolito <marcoippolito54@gmail.com>
Date: Wed, 13 Mar 2024 19:19:32 -0400
Subject: [PATCH] test_runner: support test plans

Co-Authored-By: Colin Ihrig <cjihrig@gmail.com>
---
 doc/api/test.md                               | 33 +++++++-
 lib/internal/test_runner/runner.js            |  3 +-
 lib/internal/test_runner/test.js              | 78 ++++++++++++++++-
 .../test-runner/output/test-runner-plan.js    | 49 +++++++++++
 .../output/test-runner-plan.snapshot          | 84 +++++++++++++++++++
 test/parallel/test-runner-output.mjs          |  1 +
 6 files changed, 242 insertions(+), 6 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..f7ce098125e4e2 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 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,29 @@ added:
 
 The name of the test.
 
+### `context.plan(count)`
+
+<!-- YAML
+added:
+  - REPLACEME
+-->
+
+* `count` {number} The number of assertions that are expected to run.
+
+This function is used to set the number of assertions that are expected to run
+within the test. If the number of assertions that run does not match the
+expected count, the test will fail.
+
+> Note: To make sure assertion are tracked, it must be used `t.assert` instead of `assert` directly.
+
+```js
+test('top level test', (t) => {
+  t.plan(2);
+  t.assert.ok('some relevant assertion here');
+  t.assert.ok('another relevant assertion here');
+});
+```
+
 ### `context.runOnly(shouldRunOnlyTests)`
 
 <!-- YAML
@@ -3095,6 +3122,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 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 +3139,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..6952fba4d4e97f 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;
@@ -257,7 +320,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 +436,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 +768,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 +824,7 @@ class Test extends AsyncResource {
         this.postRun();
         return;
       }
-
+      this.plan?.check();
       this.pass();
       await afterEach();
       await after();
@@ -910,7 +980,7 @@ class Test extends AsyncResource {
     this.finished = true;
 
     if (this.parent === this.root &&
-        this.root.waitingOn > this.root.subtests.length) {
+      this.root.waitingOn > this.root.subtests.length) {
       // At this point all of the tests have finished running. However, there
       // might be ref'ed handles keeping the event loop alive. This gives the
       // global after() hook a chance to clean them up. The user may also
@@ -1008,7 +1078,7 @@ class TestHook extends Test {
 
     // Report failures in the root test's after() hook.
     if (error && parent !== null &&
-        parent === parent.root && this.hookType === 'after') {
+      parent === parent.root && this.hookType === 'after') {
 
       if (isTestFailureError(error)) {
         error.failureType = kHookFailure;
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..66d6fc7ec7eb7c
--- /dev/null
+++ b/test/fixtures/test-runner/output/test-runner-plan.js
@@ -0,0 +1,49 @@
+'use strict';
+const { test } = require('node:test');
+
+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 correctly', (t) => {
+  t.plan(1);
+  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);
+  });
+});
+
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..4d827b05e63e55
--- /dev/null
+++ b/test/fixtures/test-runner/output/test-runner-plan.snapshot
@@ -0,0 +1,84 @@
+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 correctly
+    # Subtest: subtest
+    ok 1 - subtest
+      ---
+      duration_ms: *
+      ...
+    1..1
+ok 4 - subtesting correctly
+  ---
+  duration_ms: *
+  ...
+# Subtest: correctly ignoring subtesting plan
+    # Subtest: subtest
+    ok 1 - subtest
+      ---
+      duration_ms: *
+      ...
+    1..1
+not ok 5 - correctly ignoring subtesting plan
+  ---
+  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: failing planning by options
+not ok 6 - 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 7 - not failing planning by options
+  ---
+  duration_ms: *
+  ...
+# Subtest: subtest planning by options
+    # Subtest: subtest
+    ok 1 - subtest
+      ---
+      duration_ms: *
+      ...
+    1..1
+ok 8 - subtest planning by options
+  ---
+  duration_ms: *
+  ...
+1..8
+# tests 11
+# suites 0
+# pass 7
+# 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 1c5b078b23b297..db4913feaa6adf 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)