test_runner: support function mocking

cjihrig · cjihrig · commit ac9bac0b8e3a · 2022-11-05T13:35:14.000-04:00
This commit allows tests in the test runner to mock functions
and methods.
diff --git a/lib/internal/test_runner/mock.js b/lib/internal/test_runner/mock.js
@@ -0,0 +1,304 @@
+'use strict';
+const {
+  ArrayPrototypePush,
+  ArrayPrototypeSlice,
+  Error,
+  FunctionPrototypeCall,
+  ObjectDefineProperty,
+  ObjectGetOwnPropertyDescriptor,
+  Proxy,
+  ReflectApply,
+  ReflectConstruct,
+  ReflectGet,
+  SafeMap,
+} = primordials;
+const {
+  codes: {
+    ERR_INVALID_ARG_TYPE,
+    ERR_INVALID_ARG_VALUE,
+  }
+} = require('internal/errors');
+const { kEmptyObject } = require('internal/util');
+const {
+  validateBoolean,
+  validateFunction,
+  validateInteger,
+  validateObject,
+} = require('internal/validators');
+
+function kDefaultFunction() {}
+
+class MockFunctionContext {
+  #calls;
+  #mocks;
+  #implementation;
+  #restore;
+  #times;
+
+  constructor(implementation, restore, times) {
+    this.#calls = [];
+    this.#mocks = new SafeMap();
+    this.#implementation = implementation;
+    this.#restore = restore;
+    this.#times = times;
+  }
+
+  get calls() {
+    return ArrayPrototypeSlice(this.#calls, 0);
+  }
+
+  callCount() {
+    return this.#calls.length;
+  }
+
+  mockImplementation(implementation) {
+    validateFunction(implementation, 'implementation');
+    this.#implementation = implementation;
+  }
+
+  mockImplementationOnce(implementation, onCall) {
+    validateFunction(implementation, 'implementation');
+    const nextCall = this.#calls.length;
+    const call = onCall ?? nextCall;
+    validateInteger(call, 'onCall', 0);
+
+    if (call < nextCall) {
+      // The call number has already passed.
+      return;
+    }
+
+    this.#mocks.set(call, implementation);
+  }
+
+  restore() {
+    const { descriptor, object, original, methodName } = this.#restore;
+
+    if (typeof methodName === 'string') {
+      // This is an object method spy.
+      ObjectDefineProperty(object, methodName, descriptor);
+    } else {
+      // This is a bare function spy. There isn't much to do here but make
+      // the mock call the original function.
+      this.#implementation = original;
+    }
+  }
+
+  trackCall(call) {
+    ArrayPrototypePush(this.#calls, call);
+  }
+
+  nextImpl() {
+    const nextCall = this.#calls.length;
+    const mock = this.#mocks.get(nextCall);
+    const impl = mock ?? this.#implementation;
+
+    if (nextCall + 1 === this.#times) {
+      this.restore();
+    }
+
+    this.#mocks.delete(nextCall);
+    return impl;
+  }
+}
+
+const { nextImpl, restore, trackCall } = MockFunctionContext.prototype;
+delete MockFunctionContext.prototype.trackCall;
+delete MockFunctionContext.prototype.nextImpl;
+
+class MockTracker {
+  #mocks;
+
+  constructor() {
+    this.#mocks = [];
+  }
+
+  fn(
+    original = function() {},
+    implementation = original,
+    options = kEmptyObject,
+  ) {
+    if (original !== null && typeof original === 'object') {
+      options = original;
+      original = function() {};
+      implementation = original;
+    } else if (implementation !== null && typeof implementation === 'object') {
+      options = implementation;
+      implementation = original;
+    }
+
+    validateFunction(original, 'original');
+    validateFunction(implementation, 'implementation');
+    validateObject(options, 'options');
+    const { times = Infinity } = options;
+    validateTimes(times, 'options.times');
+    const ctx = new MockFunctionContext(implementation, { original }, times);
+    return this.#setupMock(ctx, original);
+  }
+
+  method(
+    object,
+    methodName,
+    implementation = kDefaultFunction,
+    options = kEmptyObject,
+  ) {
+    validateObject(object, 'object');
+    validateStringOrSymbol(methodName, 'methodName');
+
+    if (implementation !== null && typeof implementation === 'object') {
+      options = implementation;
+      implementation = kDefaultFunction;
+    }
+
+    validateFunction(implementation, 'implementation');
+    validateObject(options, 'options');
+
+    const {
+      getter = false,
+      setter = false,
+      times = Infinity,
+    } = options;
+
+    validateBoolean(getter, 'options.getter');
+    validateBoolean(setter, 'options.setter');
+    validateTimes(times, 'options.times');
+
+    if (setter && getter) {
+      throw new ERR_INVALID_ARG_VALUE(
+        'options.setter', setter, "cannot be used with 'options.getter'"
+      );
+    }
+
+    const descriptor = ObjectGetOwnPropertyDescriptor(object, methodName);
+    let original;
+
+    if (getter) {
+      original = descriptor.get;
+    } else if (setter) {
+      original = descriptor.set;
+    } else {
+      original = descriptor.value;
+    }
+
+    if (typeof original !== 'function') {
+      throw new ERR_INVALID_ARG_VALUE(
+        'methodName', original, 'must be a method'
+      );
+    }
+
+    const restore = { descriptor, object, methodName };
+    const impl = implementation === kDefaultFunction ?
+      original : implementation;
+    const ctx = new MockFunctionContext(impl, restore, times);
+    const mock = this.#setupMock(ctx, original);
+    const mockDescriptor = {
+      configurable: descriptor.configurable,
+      enumerable: descriptor.enumerable,
+    };
+
+    if (getter) {
+      mockDescriptor.get = mock;
+      mockDescriptor.set = descriptor.set;
+    } else if (setter) {
+      mockDescriptor.get = descriptor.get;
+      mockDescriptor.set = mock;
+    } else {
+      mockDescriptor.writable = descriptor.writable;
+      mockDescriptor.value = mock;
+    }
+
+    ObjectDefineProperty(object, methodName, mockDescriptor);
+
+    return mock;
+  }
+
+  reset() {
+    this.restoreAll();
+    this.#mocks = [];
+  }
+
+  restoreAll() {
+    for (let i = 0; i < this.#mocks.length; i++) {
+      FunctionPrototypeCall(restore, this.#mocks[i]);
+    }
+  }
+
+  #setupMock(ctx, fnToMatch) {
+    const mock = new Proxy(fnToMatch, {
+      __proto__: null,
+      apply(_fn, thisArg, argList) {
+        const fn = FunctionPrototypeCall(nextImpl, ctx);
+        let result;
+        let error;
+
+        try {
+          result = ReflectApply(fn, thisArg, argList);
+        } catch (err) {
+          error = err;
+          throw err;
+        } finally {
+          FunctionPrototypeCall(trackCall, ctx, {
+            arguments: argList,
+            error,
+            result,
+            // eslint-disable-next-line no-restricted-syntax
+            stack: new Error(),
+            target: undefined,
+            this: thisArg,
+          });
+        }
+
+        return result;
+      },
+      construct(target, argList, newTarget) {
+        const realTarget = FunctionPrototypeCall(nextImpl, ctx);
+        let result;
+        let error;
+
+        try {
+          result = ReflectConstruct(realTarget, argList, newTarget);
+        } catch (err) {
+          error = err;
+          throw err;
+        } finally {
+          FunctionPrototypeCall(trackCall, ctx, {
+            arguments: argList,
+            error,
+            result,
+            // eslint-disable-next-line no-restricted-syntax
+            stack: new Error(),
+            target,
+            this: result,
+          });
+        }
+
+        return result;
+      },
+      get(target, property, receiver) {
+        if (property === 'mock') {
+          return ctx;
+        }
+
+        return ReflectGet(target, property, receiver);
+      },
+    });
+
+    this.#mocks.push(ctx);
+    return mock;
+  }
+}
+
+function validateStringOrSymbol(value, name) {
+  if (typeof value !== 'string' && typeof value !== 'symbol') {
+    throw new ERR_INVALID_ARG_TYPE(name, ['string', 'symbol'], value);
+  }
+}
+
+function validateTimes(value, name) {
+  if (value === Infinity) {
+    return;
+  }
+
+  validateInteger(value, name, 1);
+}
+
+module.exports = { MockTracker };
diff --git a/lib/internal/test_runner/test.js b/lib/internal/test_runner/test.js
@@ -32,6 +32,7 @@ const {
   AbortError,
 } = require('internal/errors');
 const { getOptionValue } = require('internal/options');
+const { MockTracker } = require('internal/test_runner/mock');
 const { TapStream } = require('internal/test_runner/tap_stream');
 const {
   convertStringToRegExp,
@@ -111,6 +112,11 @@ class TestContext {
     this.#test.diagnostic(message);
   }
 
+  get mock() {
+    this.#test.mock ??= new MockTracker();
+    return this.#test.mock;
+  }
+
   runOnly(value) {
     this.#test.runOnlySubtests = !!value;
   }
@@ -238,6 +244,7 @@ class Test extends AsyncResource {
     this.#outerSignal?.addEventListener('abort', this.#abortHandler);
 
     this.fn = fn;
+    this.mock = null;
     this.name = name;
     this.parent = parent;
     this.cancelled = false;
@@ -588,6 +595,7 @@ class Test extends AsyncResource {
     }
 
     this.#outerSignal?.removeEventListener('abort', this.#abortHandler);
+    this.mock?.reset();
 
     if (this.parent !== null) {
       this.parent.activeSubtests--;
diff --git a/lib/test.js b/lib/test.js
@@ -1,5 +1,5 @@
 'use strict';
-const { ObjectAssign } = primordials;
+const { ObjectAssign, ObjectDefineProperty } = primordials;
 const { test, describe, it, before, after, beforeEach, afterEach } = require('internal/test_runner/harness');
 const { run } = require('internal/test_runner/runner');
 
@@ -14,3 +14,20 @@ ObjectAssign(module.exports, {
   run,
   test,
 });
+
+let lazyMock;
+
+ObjectDefineProperty(module.exports, 'mock', {
+  __proto__: null,
+  configurable: true,
+  enumerable: true,
+  get() {
+    if (lazyMock === undefined) {
+      const { MockTracker } = require('internal/test_runner/mock');
+
+      lazyMock = new MockTracker();
+    }
+
+    return lazyMock;
+  },
+});
diff --git a/test/parallel/test-runner-mocking.js b/test/parallel/test-runner-mocking.js
