From 1cc050eaa89ea0bd42668c3be070111bd2f5ba0f Mon Sep 17 00:00:00 2001
From: Moshe Atlow <moshe@atlow.co.il>
Date: Thu, 11 Aug 2022 22:27:34 +0300
Subject: [PATCH] assert: add `getCalls` and `reset` to callTracker

PR-URL: https://github.com/nodejs/node/pull/44191
Reviewed-By: Erick Wendel <erick.workspace@gmail.com>
Reviewed-By: Benjamin Gruenbaum <benjamingr@gmail.com>
---
 doc/api/assert.md                             |  83 +++++++++++++
 lib/internal/assert/calltracker.js            | 117 +++++++++++++-----
 .../test-assert-calltracker-getCalls.js       |  73 +++++++++++
 3 files changed, 239 insertions(+), 34 deletions(-)
 create mode 100644 test/parallel/test-assert-calltracker-getCalls.js

diff --git a/doc/api/assert.md b/doc/api/assert.md
index 245ec0d138191f..312aee4ba353d6 100644
--- a/doc/api/assert.md
+++ b/doc/api/assert.md
@@ -322,6 +322,47 @@ function func() {}
 const callsfunc = tracker.calls(func);
 ```
 
+### `tracker.getCalls(fn)`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+* `fn` {Function}.
+
+* Returns: {Array} with all the calls to a tracked function.
+
+* Object {Object}
+  * `thisArg` {Object}
+  * `arguments` {Array} the arguments passed to the tracked function
+
+```mjs
+import assert from 'node:assert';
+
+const tracker = new assert.CallTracker();
+
+function func() {}
+const callsfunc = tracker.calls(func);
+callsfunc(1, 2, 3);
+
+assert.deepStrictEqual(tracker.getCalls(callsfunc),
+                       [{ thisArg: this, arguments: [1, 2, 3 ] }]);
+```
+
+```cjs
+const assert = require('node:assert');
+
+// Creates call tracker.
+const tracker = new assert.CallTracker();
+
+function func() {}
+const callsfunc = tracker.calls(func);
+callsfunc(1, 2, 3);
+
+assert.deepStrictEqual(tracker.getCalls(callsfunc),
+                       [{ thisArg: this, arguments: [1, 2, 3 ] }]);
+```
+
 ### `tracker.report()`
 
 <!-- YAML
@@ -395,6 +436,48 @@ tracker.report();
 // ]
 ```
 
+### `tracker.reset([fn])`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+* `fn` {Function} a tracked function to reset.
+
+reset calls of the call tracker.
+if a tracked function is passed as an argument, the calls will be reset for it.
+if no arguments are passed, all tracked functions will be reset
+
+```mjs
+import assert from 'node:assert';
+
+const tracker = new assert.CallTracker();
+
+function func() {}
+const callsfunc = tracker.calls(func);
+
+callsfunc();
+// Tracker was called once
+tracker.getCalls(callsfunc).length === 1;
+
+tracker.reset(callsfunc);
+tracker.getCalls(callsfunc).length === 0;
+```
+
+```cjs
+const assert = require('node:assert');
+
+function func() {}
+const callsfunc = tracker.calls(func);
+
+callsfunc();
+// Tracker was called once
+tracker.getCalls(callsfunc).length === 1;
+
+tracker.reset(callsfunc);
+tracker.getCalls(callsfunc).length === 0;
+```
+
 ### `tracker.verify()`
 
 <!-- YAML
diff --git a/lib/internal/assert/calltracker.js b/lib/internal/assert/calltracker.js
index cf53b1af10fe28..2046634932d13d 100644
--- a/lib/internal/assert/calltracker.js
+++ b/lib/internal/assert/calltracker.js
@@ -2,16 +2,20 @@
 
 const {
   ArrayPrototypePush,
+  ArrayPrototypeSlice,
   Error,
   FunctionPrototype,
+  ObjectFreeze,
   Proxy,
   ReflectApply,
   SafeSet,
+  SafeWeakMap,
 } = primordials;
 
 const {
   codes: {
     ERR_UNAVAILABLE_DURING_EXIT,
+    ERR_INVALID_ARG_VALUE,
   },
 } = require('internal/errors');
 const AssertionError = require('internal/assert/assertion_error');
@@ -21,66 +25,111 @@ const {
 
 const noop = FunctionPrototype;
 
+class CallTrackerContext {
+  #expected;
+  #calls;
+  #name;
+  #stackTrace;
+  constructor({ expected, stackTrace, name }) {
+    this.#calls = [];
+    this.#expected = expected;
+    this.#stackTrace = stackTrace;
+    this.#name = name;
+  }
+
+  track(thisArg, args) {
+    const argsClone = ObjectFreeze(ArrayPrototypeSlice(args));
+    ArrayPrototypePush(this.#calls, ObjectFreeze({ thisArg, arguments: argsClone }));
+  }
+
+  get delta() {
+    return this.#calls.length - this.#expected;
+  }
+
+  reset() {
+    this.#calls = [];
+  }
+  getCalls() {
+    return ObjectFreeze(ArrayPrototypeSlice(this.#calls));
+  }
+
+  report() {
+    if (this.delta !== 0) {
+      const message = `Expected the ${this.#name} function to be ` +
+                      `executed ${this.#expected} time(s) but was ` +
+                      `executed ${this.#calls.length} time(s).`;
+      return {
+        message,
+        actual: this.#calls.length,
+        expected: this.#expected,
+        operator: this.#name,
+        stack: this.#stackTrace
+      };
+    }
+  }
+}
+
 class CallTracker {
 
   #callChecks = new SafeSet();
+  #trackedFunctions = new SafeWeakMap();
+
+  #getTrackedFunction(tracked) {
+    if (!this.#trackedFunctions.has(tracked)) {
+      throw new ERR_INVALID_ARG_VALUE('tracked', tracked, 'is not a tracked function');
+    }
+    return this.#trackedFunctions.get(tracked);
+  }
+
+  reset(tracked) {
+    if (tracked === undefined) {
+      this.#callChecks.forEach((check) => check.reset());
+      return;
+    }
 
-  calls(fn, exact = 1) {
+    this.#getTrackedFunction(tracked).reset();
+  }
+
+  getCalls(tracked) {
+    return this.#getTrackedFunction(tracked).getCalls();
+  }
+
+  calls(fn, expected = 1) {
     if (process._exiting)
       throw new ERR_UNAVAILABLE_DURING_EXIT();
     if (typeof fn === 'number') {
-      exact = fn;
+      expected = fn;
       fn = noop;
     } else if (fn === undefined) {
       fn = noop;
     }
 
-    validateUint32(exact, 'exact', true);
+    validateUint32(expected, 'expected', true);
 
-    const context = {
-      exact,
-      actual: 0,
+    const context = new CallTrackerContext({
+      expected,
       // eslint-disable-next-line no-restricted-syntax
       stackTrace: new Error(),
       name: fn.name || 'calls'
-    };
-    const callChecks = this.#callChecks;
-    callChecks.add(context);
-
-    return new Proxy(fn, {
+    });
+    const tracked = new Proxy(fn, {
       __proto__: null,
       apply(fn, thisArg, argList) {
-        context.actual++;
-        if (context.actual === context.exact) {
-          // Once function has reached its call count remove it from
-          // callChecks set to prevent memory leaks.
-          callChecks.delete(context);
-        }
-        // If function has been called more than expected times, add back into
-        // callchecks.
-        if (context.actual === context.exact + 1) {
-          callChecks.add(context);
-        }
+        context.track(thisArg, argList);
         return ReflectApply(fn, thisArg, argList);
       },
     });
+    this.#callChecks.add(context);
+    this.#trackedFunctions.set(tracked, context);
+    return tracked;
   }
 
   report() {
     const errors = [];
     for (const context of this.#callChecks) {
-      // If functions have not been called exact times
-      if (context.actual !== context.exact) {
-        const message = `Expected the ${context.name} function to be ` +
-                        `executed ${context.exact} time(s) but was ` +
-                        `executed ${context.actual} time(s).`;
-        ArrayPrototypePush(errors, {
-          message,
-          actual: context.actual,
-          expected: context.exact,
-          operator: context.name,
-          stack: context.stackTrace
-        });
+      const message = context.report();
+      if (message !== undefined) {
+        ArrayPrototypePush(errors, message);
       }
     }
     return errors;
diff --git a/test/parallel/test-assert-calltracker-getCalls.js b/test/parallel/test-assert-calltracker-getCalls.js
new file mode 100644
index 00000000000000..243f72c90f1def
--- /dev/null
+++ b/test/parallel/test-assert-calltracker-getCalls.js
@@ -0,0 +1,73 @@
+'use strict';
+require('../common');
+const assert = require('assert');
+const { describe, it } = require('node:test');
+
+
+describe('assert.CallTracker.getCalls()', { concurrency: true }, () => {
+  const tracker = new assert.CallTracker();
+
+  it('should return empty list when no calls', () => {
+    const fn = tracker.calls();
+    assert.deepStrictEqual(tracker.getCalls(fn), []);
+  });
+
+  it('should return calls', () => {
+    const fn = tracker.calls(() => {});
+    const arg1 = {};
+    const arg2 = {};
+    fn(arg1, arg2);
+    fn.call(arg2, arg2);
+    assert.deepStrictEqual(tracker.getCalls(fn), [
+      { arguments: [arg1, arg2], thisArg: undefined },
+      { arguments: [arg2], thisArg: arg2 }]);
+  });
+
+  it('should throw when getting calls of a non-tracked function', () => {
+    [() => {}, 1, true, null, undefined, {}, []].forEach((fn) => {
+      assert.throws(() => tracker.getCalls(fn), { code: 'ERR_INVALID_ARG_VALUE' });
+    });
+  });
+
+  it('should return a frozen object', () => {
+    const fn = tracker.calls();
+    fn();
+    const calls = tracker.getCalls(fn);
+    assert.throws(() => calls.push(1), /object is not extensible/);
+    assert.throws(() => Object.assign(calls[0], { foo: 'bar' }), /object is not extensible/);
+    assert.throws(() => calls[0].arguments.push(1), /object is not extensible/);
+  });
+});
+
+describe('assert.CallTracker.reset()', () => {
+  const tracker = new assert.CallTracker();
+
+  it('should reset calls', () => {
+    const fn = tracker.calls();
+    fn();
+    fn();
+    fn();
+    assert.strictEqual(tracker.getCalls(fn).length, 3);
+    tracker.reset(fn);
+    assert.deepStrictEqual(tracker.getCalls(fn), []);
+  });
+
+  it('should reset all calls', () => {
+    const fn1 = tracker.calls();
+    const fn2 = tracker.calls();
+    fn1();
+    fn2();
+    assert.strictEqual(tracker.getCalls(fn1).length, 1);
+    assert.strictEqual(tracker.getCalls(fn2).length, 1);
+    tracker.reset();
+    assert.deepStrictEqual(tracker.getCalls(fn1), []);
+    assert.deepStrictEqual(tracker.getCalls(fn2), []);
+  });
+
+
+  it('should throw when resetting a non-tracked function', () => {
+    [() => {}, 1, true, null, {}, []].forEach((fn) => {
+      assert.throws(() => tracker.reset(fn), { code: 'ERR_INVALID_ARG_VALUE' });
+    });
+  });
+});