test_runner: add support for setImmediate

PR-URL: nodejs#49397 Reviewed-By: Chemi Atlow <[email protected]> Reviewed-By: Benjamin Gruenbaum <[email protected]> Reviewed-By: Moshe Atlow <[email protected]> Reviewed-By: Rafael Gonzaga <[email protected]>
alexfernandez · Nov 1, 2023 · f27c3c8 · f27c3c8
1 parent ed13849
commit f27c3c8
Show file tree

Hide file tree

Showing 3 changed files with 312 additions and 10 deletions.
diff --git a/doc/api/test.md b/doc/api/test.md
@@ -1579,9 +1579,10 @@ added:
 Enables timer mocking for the specified timers.
 
 * `timers` {Array} An optional array containing the timers to mock.
-  The currently supported timer values are `'setInterval'` and `'setTimeout'`.
-  If no array is provided, all timers (`'setInterval'`, `'clearInterval'`, `'setTimeout'`,
-  and `'clearTimeout'`) will be mocked by default.
+  The currently supported timer values are `'setInterval'`, `'setTimeout'`,
+  and `'setImmediate'`.  If no value is provided, all timers (`'setInterval'`,
+  `'clearInterval'`, `'setTimeout'`, `'clearTimeout'`, `'setImmediate'`,
+  and `'clearImmediate'`) will be mocked by default.
 
 **Note:** When you enable mocking for a specific timer, its associated
 clear function will also be implicitly mocked.

diff --git a/lib/internal/test_runner/mock/mock_timers.js b/lib/internal/test_runner/mock/mock_timers.js
@@ -48,13 +48,19 @@ function abortIt(signal) {
   return new AbortError(undefined, { __proto__: null, cause: signal.reason });
 }
 
-const SUPPORTED_TIMERS = ['setTimeout', 'setInterval'];
+const SUPPORTED_TIMERS = ['setTimeout', 'setInterval', 'setImmediate'];
+const TIMERS_DEFAULT_INTERVAL = {
+  __proto__: null,
+  setImmediate: -1,
+};
 
 class MockTimers {
   #realSetTimeout;
   #realClearTimeout;
   #realSetInterval;
   #realClearInterval;
+  #realSetImmediate;
+  #realClearImmediate;
 
   #realPromisifiedSetTimeout;
   #realPromisifiedSetInterval;
@@ -63,6 +69,9 @@ class MockTimers {
   #realTimersClearTimeout;
   #realTimersSetInterval;
   #realTimersClearInterval;
+  #realTimersSetImmediate;
+  #realTimersClearImmediate;
+  #realPromisifiedSetImmediate;
 
   #timersInContext = [];
   #isEnabled = false;
@@ -76,6 +85,16 @@ class MockTimers {
   #setInterval = FunctionPrototypeBind(this.#createTimer, this, true);
   #clearInterval = FunctionPrototypeBind(this.#clearTimer, this);
 
+  #setImmediate = (callback, ...args) => {
+    return this.#createTimer(
+      false,
+      callback,
+      TIMERS_DEFAULT_INTERVAL.setImmediate,
+      ...args,
+    );
+  };
+
+  #clearImmediate = FunctionPrototypeBind(this.#clearTimer, this);
   constructor() {
     emitExperimentalWarning('The MockTimers API');
   }
@@ -158,7 +177,7 @@ class MockTimers {
     yield* iterator;
   }
 
-  #setTimeoutPromisified(ms, result, options) {
+  #promisifyTimer({ timerFn, clearFn, ms, result, options }) {
     return new Promise((resolve, reject) => {
       if (options?.signal) {
         try {
@@ -173,12 +192,12 @@ class MockTimers {
       }
 
       const onabort = () => {
-        this.#clearTimeout(id);
+        clearFn(id);
         return reject(abortIt(options.signal));
       };
 
-      const id = this.#setTimeout(() => {
-        return resolve(result || id);
+      const id = timerFn(() => {
+        return resolve(result);
       }, ms);
 
       if (options?.signal) {
@@ -192,6 +211,28 @@ class MockTimers {
     });
   }
 
+  #setImmediatePromisified(result, options) {
+    return this.#promisifyTimer({
+      __proto__: null,
+      timerFn: FunctionPrototypeBind(this.#setImmediate, this),
+      clearFn: FunctionPrototypeBind(this.#clearImmediate, this),
+      ms: TIMERS_DEFAULT_INTERVAL.setImmediate,
+      result,
+      options,
+    });
+  }
+
+  #setTimeoutPromisified(ms, result, options) {
+    return this.#promisifyTimer({
+      __proto__: null,
+      timerFn: FunctionPrototypeBind(this.#setTimeout, this),
+      clearFn: FunctionPrototypeBind(this.#clearTimeout, this),
+      ms,
+      result,
+      options,
+    });
+  }
+
   #toggleEnableTimers(activate) {
     const options = {
       __proto__: null,
@@ -233,6 +274,23 @@ class MockTimers {
             this,
           );
         },
+        setImmediate: () => {
+          this.#realSetImmediate = globalThis.setImmediate;
+          this.#realClearImmediate = globalThis.clearImmediate;
+          this.#realTimersSetImmediate = nodeTimers.setImmediate;
+          this.#realTimersClearImmediate = nodeTimers.clearImmediate;
+
+          globalThis.setImmediate = this.#setImmediate;
+          globalThis.clearImmediate = this.#clearImmediate;
+
+          nodeTimers.setImmediate = this.#setImmediate;
+          nodeTimers.clearImmediate = this.#clearImmediate;
+
+          nodeTimersPromises.setImmediate = FunctionPrototypeBind(
+            this.#setImmediatePromisified,
+            this,
+          );
+        },
       },
       toReal: {
         __proto__: null,
@@ -254,6 +312,15 @@ class MockTimers {
 
           nodeTimersPromises.setInterval = this.#realPromisifiedSetInterval;
         },
+        setImmediate: () => {
+          globalThis.setImmediate = this.#realSetImmediate;
+          globalThis.clearImmediate = this.#realClearImmediate;
+
+          nodeTimers.setImmediate = this.#realTimersSetImmediate;
+          nodeTimers.clearImmediate = this.#realTimersClearImmediate;
+
+          nodeTimersPromises.setImmediate = this.#realPromisifiedSetImmediate;
+        },
       },
     };