From bbc5623d8acf642d4aba929faa283810e405f8df Mon Sep 17 00:00:00 2001
From: Colin Ihrig <cjihrig@gmail.com>
Date: Wed, 15 Jan 2025 21:32:37 -0500
Subject: [PATCH] test_runner: add TestContext.prototype.waitFor()
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This commit adds a waitFor() method to the TestContext class in
the test runner. As the name implies, this method allows tests to
more easily wait for things to happen.

PR-URL: https://github.com/nodejs/node/pull/56595
Reviewed-By: Pietro Marchini <pietro.marchini94@gmail.com>
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
Reviewed-By: Chemi Atlow <chemi@atlow.co.il>
Reviewed-By: Michaël Zasso <targos@protonmail.com>
---
 doc/api/test.md                       |  21 +++++
 lib/internal/test_runner/test.js      |  62 ++++++++++++-
 test/parallel/test-runner-wait-for.js | 124 ++++++++++++++++++++++++++
 3 files changed, 206 insertions(+), 1 deletion(-)
 create mode 100644 test/parallel/test-runner-wait-for.js

diff --git a/doc/api/test.md b/doc/api/test.md
index 3b0c040d852334..68b4a0f8a6a969 100644
--- a/doc/api/test.md
+++ b/doc/api/test.md
@@ -3604,6 +3604,27 @@ test('top level test', async (t) => {
 });
 ```
 
+### `context.waitFor(condition[, options])`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+* `condition` {Function|AsyncFunction} An assertion function that is invoked
+  periodically until it completes successfully or the defined polling timeout
+  elapses. Successful completion is defined as not throwing or rejecting. This
+  function does not accept any arguments, and is allowed to return any value.
+* `options` {Object} An optional configuration object for the polling operation.
+  The following properties are supported:
+  * `interval` {number} The number of milliseconds to wait after an unsuccessful
+    invocation of `condition` before trying again. **Default:** `50`.
+  * `timeout` {number} The poll timeout in milliseconds. If `condition` has not
+    succeeded by the time this elapses, an error occurs. **Default:** `1000`.
+* Returns: {Promise} Fulfilled with the value returned by `condition`.
+
+This method polls a `condition` function until that function either returns
+successfully or the operation times out.
+
 ## Class: `SuiteContext`
 
 <!-- YAML
diff --git a/lib/internal/test_runner/test.js b/lib/internal/test_runner/test.js
index 0dc61255874a04..b13d13e105e9ba 100644
--- a/lib/internal/test_runner/test.js
+++ b/lib/internal/test_runner/test.js
@@ -8,6 +8,7 @@ const {
   ArrayPrototypeSplice,
   ArrayPrototypeUnshift,
   ArrayPrototypeUnshiftApply,
+  Error,
   FunctionPrototype,
   MathMax,
   Number,
@@ -58,11 +59,16 @@ const {
 const { isPromise } = require('internal/util/types');
 const {
   validateAbortSignal,
+  validateFunction,
   validateNumber,
+  validateObject,
   validateOneOf,
   validateUint32,
 } = require('internal/validators');
-const { setTimeout } = require('timers');
+const {
+  clearTimeout,
+  setTimeout,
+} = require('timers');
 const { TIMEOUT_MAX } = require('internal/timers');
 const { fileURLToPath } = require('internal/url');
 const { availableParallelism } = require('os');
@@ -340,6 +346,60 @@ class TestContext {
       loc: getCallerLocation(),
     });
   }
+
+  waitFor(condition, options = kEmptyObject) {
+    validateFunction(condition, 'condition');
+    validateObject(options, 'options');
+
+    const {
+      interval = 50,
+      timeout = 1000,
+    } = options;
+
+    validateNumber(interval, 'options.interval', 0, TIMEOUT_MAX);
+    validateNumber(timeout, 'options.timeout', 0, TIMEOUT_MAX);
+
+    const { promise, resolve, reject } = PromiseWithResolvers();
+    const noError = Symbol();
+    let cause = noError;
+    let pollerId;
+    let timeoutId;
+    const done = (err, result) => {
+      clearTimeout(pollerId);
+      clearTimeout(timeoutId);
+
+      if (err === noError) {
+        resolve(result);
+      } else {
+        reject(err);
+      }
+    };
+
+    timeoutId = setTimeout(() => {
+      // eslint-disable-next-line no-restricted-syntax
+      const err = new Error('waitFor() timed out');
+
+      if (cause !== noError) {
+        err.cause = cause;
+      }
+
+      done(err);
+    }, timeout);
+
+    const poller = async () => {
+      try {
+        const result = await condition();
+
+        done(noError, result);
+      } catch (err) {
+        cause = err;
+        pollerId = setTimeout(poller, interval);
+      }
+    };
+
+    poller();
+    return promise;
+  }
 }
 
 class SuiteContext {
diff --git a/test/parallel/test-runner-wait-for.js b/test/parallel/test-runner-wait-for.js
new file mode 100644
index 00000000000000..8f1e28f12868e4
--- /dev/null
+++ b/test/parallel/test-runner-wait-for.js
@@ -0,0 +1,124 @@
+'use strict';
+require('../common');
+const { suite, test } = require('node:test');
+
+suite('input validation', () => {
+  test('throws if condition is not a function', (t) => {
+    t.assert.throws(() => {
+      t.waitFor(5);
+    }, {
+      code: 'ERR_INVALID_ARG_TYPE',
+      message: /The "condition" argument must be of type function/,
+    });
+  });
+
+  test('throws if options is not an object', (t) => {
+    t.assert.throws(() => {
+      t.waitFor(() => {}, null);
+    }, {
+      code: 'ERR_INVALID_ARG_TYPE',
+      message: /The "options" argument must be of type object/,
+    });
+  });
+
+  test('throws if options.interval is not a number', (t) => {
+    t.assert.throws(() => {
+      t.waitFor(() => {}, { interval: 'foo' });
+    }, {
+      code: 'ERR_INVALID_ARG_TYPE',
+      message: /The "options\.interval" property must be of type number/,
+    });
+  });
+
+  test('throws if options.timeout is not a number', (t) => {
+    t.assert.throws(() => {
+      t.waitFor(() => {}, { timeout: 'foo' });
+    }, {
+      code: 'ERR_INVALID_ARG_TYPE',
+      message: /The "options\.timeout" property must be of type number/,
+    });
+  });
+});
+
+test('returns the result of the condition function', async (t) => {
+  const result = await t.waitFor(() => {
+    return 42;
+  });
+
+  t.assert.strictEqual(result, 42);
+});
+
+test('returns the result of an async condition function', async (t) => {
+  const result = await t.waitFor(async () => {
+    return 84;
+  });
+
+  t.assert.strictEqual(result, 84);
+});
+
+test('errors if the condition times out', async (t) => {
+  await t.assert.rejects(async () => {
+    await t.waitFor(() => {
+      return new Promise(() => {});
+    }, {
+      interval: 60_000,
+      timeout: 1,
+    });
+  }, {
+    message: /waitFor\(\) timed out/,
+  });
+});
+
+test('polls until the condition returns successfully', async (t) => {
+  let count = 0;
+  const result = await t.waitFor(() => {
+    ++count;
+    if (count < 4) {
+      throw new Error('resource is not ready yet');
+    }
+
+    return 'success';
+  }, {
+    interval: 1,
+    timeout: 60_000,
+  });
+
+  t.assert.strictEqual(result, 'success');
+  t.assert.strictEqual(count, 4);
+});
+
+test('sets last failure as error cause on timeouts', async (t) => {
+  const error = new Error('boom');
+  await t.assert.rejects(async () => {
+    await t.waitFor(() => {
+      return new Promise((_, reject) => {
+        reject(error);
+      });
+    });
+  }, (err) => {
+    t.assert.match(err.message, /timed out/);
+    t.assert.strictEqual(err.cause, error);
+    return true;
+  });
+});
+
+test('limits polling if condition takes longer than interval', async (t) => {
+  let count = 0;
+
+  function condition() {
+    count++;
+    return new Promise((resolve) => {
+      setTimeout(() => {
+        resolve('success');
+      }, 200);
+    });
+  }
+
+  const result = await t.waitFor(condition, {
+    interval: 1,
+    timeout: 60_000,
+  });
+
+  t.assert.strictEqual(result, 'success');
+  t.assert.strictEqual(count, 1);
+});