test_runner: add support for chainable mockImplementationOnce

Fixes: nodejs#47718
rluvaton · Aug 10, 2023 · 4bb7656 · 4bb7656
1 parent f1b3ade
commit 4bb7656
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 3 deletions.
diff --git a/lib/internal/test_runner/mock/mock.js b/lib/internal/test_runner/mock/mock.js
@@ -61,9 +61,14 @@ class MockFunctionContext {
   mockImplementationOnce(implementation, onCall) {
     validateFunction(implementation, 'implementation');
     const nextCall = this.#calls.length;
-    const call = onCall ?? nextCall;
-    validateInteger(call, 'onCall', nextCall);
-    this.#mocks.set(call, implementation);
+    if(onCall) {
+        validateInteger(onCall, 'onCall', nextCall);
+        this.#mocks.set(onCall, implementation);
+    } else {
+        this.#mocks.set(this.#mocks.size + nextCall, implementation);
+    }
+
+    return this;
   }
 
   restore() {

diff --git a/test/parallel/test-runner-mocking.js b/test/parallel/test-runner-mocking.js
@@ -891,6 +891,55 @@ test('mock implementation can be changed dynamically', (t) => {
   assert.strictEqual(fn.mock.callCount(), 12);
 });
 
+test('mockImplementationOnce chaining', (t) => {
+  let cnt = 0;
+  function addOne() {
+    cnt++;
+    return cnt;
+  }
+
+  const fn = t.mock.fn(addOne);
+  assert.strictEqual(fn(), 1);
+  assert.strictEqual(fn(), 2);
+  fn.mock
+    .mockImplementationOnce(() => 17)
+    .mockImplementationOnce(() => 42)
+    .mockImplementationOnce(() => 77, 5)
+  assert.strictEqual(fn(), 17);
+  assert.strictEqual(fn(), 42);
+  assert.strictEqual(fn(), 3);
+  assert.strictEqual(fn(), 77);
+  assert.strictEqual(fn(), 4);
+  assert.strictEqual(fn(), 5);
+
+  fn.mock.mockImplementationOnce(() => 52)
+  assert.strictEqual(fn(), 52);
+  assert.strictEqual(fn(), 6);
+});
+
+test('mockImplementationOnce chaining with default value', (t) => {
+  let cnt = 0;
+  function addOne() {
+    cnt++;
+    return cnt;
+  }
+
+  const fn = t.mock.fn(addOne);
+  assert.strictEqual(fn(), 1);
+  fn.mock
+    .mockImplementationOnce(() => 17)
+    .mockImplementationOnce(() => 42)
+    .mockImplementation(() => 99)
+  assert.strictEqual(fn(), 17);
+  assert.strictEqual(fn(), 42);
+  assert.strictEqual(fn(), 99);
+  assert.strictEqual(fn(), 99);
+
+  fn.mock.mockImplementationOnce(() => 52)
+  assert.strictEqual(fn(), 52);
+  assert.strictEqual(fn(), 99);
+});
+
 test('local mocks are auto restored after the test finishes', async (t) => {
   const obj = {
     foo() {},