From 4a7b8157b5585cb62e77e9e3c66d6af3aaf3a448 Mon Sep 17 00:00:00 2001
From: Colin Ihrig <cjihrig@gmail.com>
Date: Sat, 4 Jan 2025 13:30:04 -0500
Subject: [PATCH] test_runner: add assert.register() API

This commit adds a top level assert.register() API to the test
runner. This function allows users to define their own custom
assertion functions on the TestContext.

Fixes: https://github.com/nodejs/node/issues/52033
PR-URL: https://github.com/nodejs/node/pull/56434
Reviewed-By: Jacob Smith <jacob@frende.me>
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
Reviewed-By: Pietro Marchini <pietro.marchini94@gmail.com>
---
 doc/api/test.md                               | 23 +++++++
 lib/internal/test_runner/assert.js            | 50 +++++++++++++++
 lib/internal/test_runner/test.js              | 52 ++++++---------
 lib/test.js                                   | 12 ++++
 .../parallel/test-runner-custom-assertions.js | 63 +++++++++++++++++++
 5 files changed, 166 insertions(+), 34 deletions(-)
 create mode 100644 lib/internal/test_runner/assert.js
 create mode 100644 test/parallel/test-runner-custom-assertions.js

diff --git a/doc/api/test.md b/doc/api/test.md
index 7deac4fae1efdf..c3ea68fa4173bf 100644
--- a/doc/api/test.md
+++ b/doc/api/test.md
@@ -1750,6 +1750,29 @@ describe('tests', async () => {
 });
 ```
 
+## `assert`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+An object whose methods are used to configure available assertions on the
+`TestContext` objects in the current process. The methods from `node:assert`
+and snapshot testing functions are available by default.
+
+It is possible to apply the same configuration to all files by placing common
+configuration code in a module
+preloaded with `--require` or `--import`.
+
+### `assert.register(name, fn)`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+Defines a new assertion function with the provided name and function. If an
+assertion already exists with the same name, it is overwritten.
+
 ## `snapshot`
 
 <!-- YAML
diff --git a/lib/internal/test_runner/assert.js b/lib/internal/test_runner/assert.js
new file mode 100644
index 00000000000000..776c1e25cbfb61
--- /dev/null
+++ b/lib/internal/test_runner/assert.js
@@ -0,0 +1,50 @@
+'use strict';
+const {
+  SafeMap,
+} = primordials;
+const {
+  validateFunction,
+  validateString,
+} = require('internal/validators');
+const assert = require('assert');
+const methodsToCopy = [
+  'deepEqual',
+  'deepStrictEqual',
+  'doesNotMatch',
+  'doesNotReject',
+  'doesNotThrow',
+  'equal',
+  'fail',
+  'ifError',
+  'match',
+  'notDeepEqual',
+  'notDeepStrictEqual',
+  'notEqual',
+  'notStrictEqual',
+  'partialDeepStrictEqual',
+  'rejects',
+  'strictEqual',
+  'throws',
+];
+let assertMap;
+
+function getAssertionMap() {
+  if (assertMap === undefined) {
+    assertMap = new SafeMap();
+
+    for (let i = 0; i < methodsToCopy.length; i++) {
+      assertMap.set(methodsToCopy[i], assert[methodsToCopy[i]]);
+    }
+  }
+
+  return assertMap;
+}
+
+function register(name, fn) {
+  validateString(name, 'name');
+  validateFunction(fn, 'fn');
+  const map = getAssertionMap();
+  map.set(name, fn);
+}
+
+module.exports = { getAssertionMap, register };
diff --git a/lib/internal/test_runner/test.js b/lib/internal/test_runner/test.js
index 6a92ec335ddf34..b4181635f9a784 100644
--- a/lib/internal/test_runner/test.js
+++ b/lib/internal/test_runner/test.js
@@ -100,34 +100,15 @@ function lazyFindSourceMap(file) {
 
 function lazyAssertObject(harness) {
   if (assertObj === undefined) {
-    assertObj = new SafeMap();
-    const assert = require('assert');
-    const { SnapshotManager } = require('internal/test_runner/snapshot');
-    const methodsToCopy = [
-      'deepEqual',
-      'deepStrictEqual',
-      'doesNotMatch',
-      'doesNotReject',
-      'doesNotThrow',
-      'equal',
-      'fail',
-      'ifError',
-      'match',
-      'notDeepEqual',
-      'notDeepStrictEqual',
-      'notEqual',
-      'notStrictEqual',
-      'partialDeepStrictEqual',
-      'rejects',
-      'strictEqual',
-      'throws',
-    ];
-    for (let i = 0; i < methodsToCopy.length; i++) {
-      assertObj.set(methodsToCopy[i], assert[methodsToCopy[i]]);
-    }
+    const { getAssertionMap } = require('internal/test_runner/assert');
+
+    assertObj = getAssertionMap();
+    if (!assertObj.has('snapshot')) {
+      const { SnapshotManager } = require('internal/test_runner/snapshot');
 
-    harness.snapshotManager = new SnapshotManager(harness.config.updateSnapshots);
-    assertObj.set('snapshot', harness.snapshotManager.createAssert());
+      harness.snapshotManager = new SnapshotManager(harness.config.updateSnapshots);
+      assertObj.set('snapshot', harness.snapshotManager.createAssert());
+    }
   }
   return assertObj;
 }
@@ -264,15 +245,18 @@ class TestContext {
         };
       });
 
-      // This is a hack. It allows the innerOk function to collect the stacktrace from the correct starting point.
-      function ok(...args) {
-        if (plan !== null) {
-          plan.actual++;
+      if (!map.has('ok')) {
+        // This is a hack. It allows the innerOk function to collect the
+        // stacktrace from the correct starting point.
+        function ok(...args) {
+          if (plan !== null) {
+            plan.actual++;
+          }
+          innerOk(ok, args.length, ...args);
         }
-        innerOk(ok, args.length, ...args);
-      }
 
-      assert.ok = ok;
+        assert.ok = ok;
+      }
     }
     return this.#assert;
   }
diff --git a/lib/test.js b/lib/test.js
index 8e7c6295a37225..d6a313cd0763eb 100644
--- a/lib/test.js
+++ b/lib/test.js
@@ -61,3 +61,15 @@ ObjectDefineProperty(module.exports, 'snapshot', {
     return lazySnapshot;
   },
 });
+
+ObjectDefineProperty(module.exports, 'assert', {
+  __proto__: null,
+  configurable: true,
+  enumerable: true,
+  get() {
+    const { register } = require('internal/test_runner/assert');
+    const assert = { __proto__: null, register };
+    ObjectDefineProperty(module.exports, 'assert', assert);
+    return assert;
+  },
+});
diff --git a/test/parallel/test-runner-custom-assertions.js b/test/parallel/test-runner-custom-assertions.js
new file mode 100644
index 00000000000000..a4bdf0f548be80
--- /dev/null
+++ b/test/parallel/test-runner-custom-assertions.js
@@ -0,0 +1,63 @@
+'use strict';
+require('../common');
+const assert = require('node:assert');
+const { test, assert: testAssertions } = require('node:test');
+
+testAssertions.register('isOdd', (n) => {
+  assert.strictEqual(n % 2, 1);
+});
+
+testAssertions.register('ok', () => {
+  return 'ok';
+});
+
+testAssertions.register('snapshot', () => {
+  return 'snapshot';
+});
+
+testAssertions.register('deepStrictEqual', () => {
+  return 'deepStrictEqual';
+});
+
+testAssertions.register('context', function() {
+  return this;
+});
+
+test('throws if name is not a string', () => {
+  assert.throws(() => {
+    testAssertions.register(5);
+  }, {
+    code: 'ERR_INVALID_ARG_TYPE',
+    message: 'The "name" argument must be of type string. Received type number (5)'
+  });
+});
+
+test('throws if fn is not a function', () => {
+  assert.throws(() => {
+    testAssertions.register('foo', 5);
+  }, {
+    code: 'ERR_INVALID_ARG_TYPE',
+    message: 'The "fn" argument must be of type function. Received type number (5)'
+  });
+});
+
+test('invokes a custom assertion as part of the test plan', (t) => {
+  t.plan(2);
+  t.assert.isOdd(5);
+  assert.throws(() => {
+    t.assert.isOdd(4);
+  }, {
+    code: 'ERR_ASSERTION',
+    message: /Expected values to be strictly equal/
+  });
+});
+
+test('can override existing assertions', (t) => {
+  assert.strictEqual(t.assert.ok(), 'ok');
+  assert.strictEqual(t.assert.snapshot(), 'snapshot');
+  assert.strictEqual(t.assert.deepStrictEqual(), 'deepStrictEqual');
+});
+
+test('"this" is set to the TestContext', (t) => {
+  assert.strictEqual(t.assert.context(), t);
+});