test_runner: add support for chainable mockImplementationOnce

doc/api/test.md

@@ -1318,13 +1318,19 @@ test('changes a mock behavior', (t) => {
 added:
   - v19.1.0
   - v18.13.0
+changes:
+  - version: REPLACEME
+    pr-url: https://github.com/nodejs/node/pull/49088
+    description: this now return the context as well as appending when not providing the onCall argument.
 -->
 
 * `implementation` {Function|AsyncFunction} The function to be used as the
   mock's implementation for the invocation number specified by `onCall`.
 * `onCall` {integer} The invocation number that will use `implementation`. If
   the specified invocation has already occurred then an exception is thrown.
-  **Default:** The number of the next invocation.
+  **Default:** The number of the next invocation that is not mocked by `mockImplementationOnce`.
+* Returns: {MockFunctionContext} A reference to the [`MockFunctionContext`][],
+  so that calls can be chained.
 
 This function is used to change the behavior of an existing mock for a single
 invocation. Once invocation `onCall` has occurred, the mock will revert to
@@ -1352,8 +1358,15 @@ test('changes a mock behavior once', (t) => {
   const fn = t.mock.fn(addOne);
 
   assert.strictEqual(fn(), 1);
-  fn.mock.mockImplementationOnce(addTwo);
+  assert.strictEqual(fn(), 2);
+  fn.mock
+    .mockImplementationOnce(() => 17)
+    .mockImplementationOnce(() => 77, 5)
+    .mockImplementationOnce(() => 42);
+  assert.strictEqual(fn(), 17);
+  assert.strictEqual(fn(), 42);
   assert.strictEqual(fn(), 3);
+  assert.strictEqual(fn(), 77);
   assert.strictEqual(fn(), 4);
 });
 ```

lib/internal/test_runner/mock/mock.js

@@ -61,9 +61,24 @@ 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);
+    // eslint-disable-next-line eqeqeq
+    if (onCall != undefined) {
+      validateInteger(onCall, 'onCall', nextCall);
+      this.#mocks.set(onCall, implementation);
+
+      return this;
+    }
+
+    const mocksSize = this.#mocks.size;
+    for (let i = 0; i < mocksSize; i++) {
+      if (!this.#mocks.has(i + nextCall)) {
+        this.#mocks.set(i + nextCall, implementation);
+        return this;
+      }
+    }
+    this.#mocks.set(mocksSize + nextCall, implementation);
+
+    return this;
   }
 
   restore() {

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(() => 77, 5)
+    .mockImplementationOnce(() => 17)
+    .mockImplementationOnce(() => 42);
+  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() {},

tools/doc/type-parser.mjs

@@ -212,6 +212,7 @@ const customTypesMap = {
   'Timer': 'timers.html#timers',
 
   'TestsStream': 'test.html#class-testsstream',
+  'MockFunctionContext': 'test.html#class-mockfunctioncontext',
 
   'tls.SecureContext': 'tls.html#tlscreatesecurecontextoptions',
   'tls.Server': 'tls.html#class-tlsserver',

typings/primordials.d.ts

@@ -120,6 +120,7 @@ declare namespace primordials {
   export const ArrayPrototypeLastIndexOf: UncurryThis<typeof Array.prototype.lastIndexOf>
   export const ArrayPrototypePop: UncurryThis<typeof Array.prototype.pop>
   export const ArrayPrototypePush: UncurryThis<typeof Array.prototype.push>
+  export const ArrayPrototypeAt: UncurryThis<typeof Array.prototype.at>
   export const ArrayPrototypePushApply: UncurryThisStaticApply<typeof Array.prototype.push>
   export const ArrayPrototypeReverse: UncurryThis<typeof Array.prototype.reverse>
   export const ArrayPrototypeShift: UncurryThis<typeof Array.prototype.shift>