feat(mock): support asynchronous reply options callbacks

docs/docs/api/MockPool.md

@@ -136,7 +136,8 @@ The reply behaviour of a matching request is defined through the returned
   computing all reply options dynamically rather than just the body.
   * `callback` {Function} A `(opts: MockResponseCallbackOptions) =>
     { statusCode, data, responseOptions }` function invoked with the incoming
-    request.
+    request. The callback may be asynchronous; a returned promise is awaited
+    and must resolve to the same shape.
   * Returns: {MockScope}
 * `replyWithError(error)` {Function} Defines an error for a matching request to
   throw.
@@ -263,6 +264,32 @@ for await (const data of body) {
 }
 ```
 
+```mjs displayName="Reply with an asynchronous options callback"
+import { readFile } from 'node:fs/promises'
+import { MockAgent, setGlobalDispatcher, request } from 'undici'
+
+const mockAgent = new MockAgent()
+setGlobalDispatcher(mockAgent)
+
+const mockPool = mockAgent.get('http://localhost:3000')
+
+mockPool.intercept({
+  path: '/fixture',
+  method: 'GET'
+}).reply(async ({ path }) => ({
+  statusCode: 200,
+  data: await readFile(new URL('./fixture.json', import.meta.url))
+}))
+
+const { statusCode, body } = await request('http://localhost:3000/fixture')
+
+console.log('response received', statusCode) // response received 200
+
+for await (const data of body) {
+  console.log('data', data.toString('utf8')) // contents of fixture.json
+}
+```
+
 ```mjs displayName="Multiple intercepts"
 import { MockAgent, setGlobalDispatcher, request } from 'undici'
 

lib/mock/mock-interceptor.js

@@ -12,6 +12,11 @@ const {
 } = require('./mock-symbols')
 const { InvalidArgumentError } = require('../core/errors')
 const { serializePathWithQuery } = require('../core/util')
+const {
+  types: {
+    isPromise
+  }
+} = require('node:util')
 
 /**
  * Defines the scope API for an interceptor reply
@@ -117,13 +122,9 @@ class MockInterceptor {
     // Values of reply aren't available right now as they
     // can only be available when the reply callback is invoked.
     if (typeof replyOptionsCallbackOrStatusCode === 'function') {
-      // We'll first wrap the provided callback in another function,
-      // this function will properly resolve the data from the callback
-      // when invoked.
-      const wrappedDefaultsCallback = (opts) => {
-        // Our reply options callback contains the parameter for statusCode, data and options.
-        const resolvedData = replyOptionsCallbackOrStatusCode(opts)
-
+      // Resolves the data returned by a reply options callback into
+      // dispatch data, validating its format along the way.
+      const resolveReplyCallbackData = (resolvedData) => {
         // Check if it is in the right format
         if (typeof resolvedData !== 'object' || resolvedData === null) {
           throw new InvalidArgumentError('reply options callback must return an object')
@@ -138,6 +139,23 @@ class MockInterceptor {
         }
       }
 
+      // We'll first wrap the provided callback in another function,
+      // this function will properly resolve the data from the callback
+      // when invoked.
+      const wrappedDefaultsCallback = (opts) => {
+        // Our reply options callback contains the parameter for statusCode, data and options.
+        const resolvedData = replyOptionsCallbackOrStatusCode(opts)
+
+        // An asynchronous reply options callback resolves to the reply
+        // parameters, so the dispatch data can only be resolved once the
+        // returned promise settles.
+        if (isPromise(resolvedData)) {
+          return resolvedData.then(resolveReplyCallbackData)
+        }
+
+        return resolveReplyCallbackData(resolvedData)
+      }
+
       // Add usual dispatch data, but this time set the data parameter to function that will eventually provide data.
       const newMockDispatch = addMockDispatch(this[kDispatches], this[kDispatchKey], wrappedDefaultsCallback, { ignoreTrailingSlash: this[kIgnoreTrailingSlash] })
       return new MockScope(newMockDispatch)

lib/mock/mock-utils.js

@@ -292,25 +292,54 @@ function mockDispatch (opts, handler) {
   // Get mock dispatch from built key
   const key = buildKey(opts)
   const mockDispatch = getMockDispatch(this[kDispatches], key)
+  const mockDispatches = this[kDispatches]
 
   mockDispatch.timesInvoked++
 
+  const { timesInvoked, times } = mockDispatch
+
+  // If it's used up and not persistent, mark as consumed
+  mockDispatch.consumed = !mockDispatch.persist && timesInvoked >= times
+  mockDispatch.pending = timesInvoked < times
+
   // Here's where we resolve a callback if a callback is present for the dispatch data.
   if (mockDispatch.data.callback) {
-    mockDispatch.data = { ...mockDispatch.data, ...mockDispatch.data.callback(opts) }
+    const callbackResult = mockDispatch.data.callback(opts)
+
+    // An asynchronous reply options callback resolves to the reply data, so
+    // the dispatch can only continue once the returned promise settles.
+    // A rejection cannot be thrown synchronously from the dispatch at that
+    // point, so it is surfaced as a response error instead.
+    if (isPromise(callbackResult)) {
+      callbackResult.then(
+        (resolvedData) => {
+          mockDispatch.data = { ...mockDispatch.data, ...resolvedData }
+          dispatchMockReply(mockDispatches, mockDispatch, key, opts, handler)
+        },
+        (error) => {
+          deleteMockDispatch(mockDispatches, key)
+          handler.onResponseError(null, error)
+        }
+      )
+      return true
+    }
+
+    mockDispatch.data = { ...mockDispatch.data, ...callbackResult }
   }
 
-  // Parse mockDispatch data
-  const { data: { statusCode, data, headers, trailers, error }, delay, persist } = mockDispatch
-  const { timesInvoked, times } = mockDispatch
+  return dispatchMockReply(mockDispatches, mockDispatch, key, opts, handler)
+}
 
-  // If it's used up and not persistent, mark as consumed
-  mockDispatch.consumed = !persist && timesInvoked >= times
-  mockDispatch.pending = timesInvoked < times
+/**
+ * Replies to a request once the mock dispatch data is fully resolved
+ */
+function dispatchMockReply (mockDispatches, mockDispatch, key, opts, handler) {
+  // Parse mockDispatch data
+  const { data: { statusCode, data, headers, trailers, error }, delay } = mockDispatch
 
   // If specified, trigger dispatch error
   if (error !== null) {
-    deleteMockDispatch(this[kDispatches], key)
+    deleteMockDispatch(mockDispatches, key)
     handler.onResponseError(null, error)
     return true
   }
@@ -353,10 +382,10 @@ function mockDispatch (opts, handler) {
   if (typeof delay === 'number' && delay > 0) {
     timer = setTimeout(() => {
       timer = null
-      handleReply(this[kDispatches])
+      handleReply(mockDispatches)
     }, delay)
   } else {
-    handleReply(this[kDispatches])
+    handleReply(mockDispatches)
   }
 
   function handleReply (mockDispatches, _data = data) {

test/mock-interceptor.js

@@ -5,7 +5,9 @@ const { MockInterceptor, MockScope } = require('../lib/mock/mock-interceptor')
 const MockAgent = require('../lib/mock/mock-agent')
 const { kDispatchKey } = require('../lib/mock/mock-symbols')
 const { InvalidArgumentError } = require('../lib/core/errors')
+const { MockNotMatchedError } = require('../lib/mock/mock-errors')
 const { fetch } = require('../lib/web/fetch/index')
+const { request } = require('../index')
 
 describe('MockInterceptor - path', () => {
   test('should remove hash fragment from paths', t => {
@@ -232,6 +234,132 @@ describe('MockInterceptor - reply options callback', () => {
   })
 })
 
+describe('MockInterceptor - asynchronous reply options callback', () => {
+  test('should resolve the reply from an asynchronous callback', async t => {
+    t.plan(3)
+
+    const baseUrl = 'http://localhost:9999'
+    const mockAgent = new MockAgent()
+    mockAgent.disableNetConnect()
+    after(() => mockAgent.close())
+
+    const mockPool = mockAgent.get(baseUrl)
+
+    mockPool.intercept({
+      path: '/test',
+      method: 'GET'
+    }).reply(async (options) => {
+      t.assert.strictEqual(options.path, '/test')
+      await new Promise((resolve) => setImmediate(resolve))
+      return { statusCode: 201, data: 'hello' }
+    })
+
+    const { statusCode, body } = await request(`${baseUrl}/test`, { dispatcher: mockAgent })
+    t.assert.strictEqual(statusCode, 201)
+    t.assert.strictEqual(await body.text(), 'hello')
+  })
+
+  test('should apply default headers and content length after resolution', async t => {
+    t.plan(4)
+
+    const baseUrl = 'http://localhost:9999'
+    const mockAgent = new MockAgent()
+    mockAgent.disableNetConnect()
+    after(() => mockAgent.close())
+
+    const mockPool = mockAgent.get(baseUrl)
+
+    mockPool.intercept({
+      path: '/test',
+      method: 'GET'
+    }).defaultReplyHeaders({ foo: 'bar' }).replyContentLength().reply(async () => ({
+      statusCode: 200,
+      data: 'hello'
+    }))
+
+    const { statusCode, headers, body } = await request(`${baseUrl}/test`, { dispatcher: mockAgent })
+    t.assert.strictEqual(statusCode, 200)
+    t.assert.strictEqual(headers.foo, 'bar')
+    t.assert.strictEqual(headers['content-length'], '5')
+    t.assert.strictEqual(await body.text(), 'hello')
+  })
+
+  test('should support times() with an asynchronous callback', async t => {
+    t.plan(3)
+
+    const baseUrl = 'http://localhost:9999'
+    const mockAgent = new MockAgent()
+    mockAgent.disableNetConnect()
+    after(() => mockAgent.close())
+
+    const mockPool = mockAgent.get(baseUrl)
+
+    mockPool.intercept({
+      path: '/test',
+      method: 'GET'
+    }).reply(async () => ({ statusCode: 200, data: 'hello' })).times(2)
+
+    for (let i = 0; i < 2; i++) {
+      const { statusCode, body } = await request(`${baseUrl}/test`, { dispatcher: mockAgent })
+      t.assert.strictEqual(statusCode, 200)
+      await body.text()
+    }
+
+    await t.assert.rejects(request(`${baseUrl}/test`, { dispatcher: mockAgent }), MockNotMatchedError)
+  })
+
+  test('should reject if an asynchronous callback resolves to an invalid format', async t => {
+    t.plan(2)
+
+    const baseUrl = 'http://localhost:9999'
+    const mockAgent = new MockAgent()
+    mockAgent.disableNetConnect()
+    after(() => mockAgent.close())
+
+    const mockPool = mockAgent.get(baseUrl)
+
+    mockPool.intercept({
+      path: '/test-resolve-null',
+      method: 'GET'
+    }).reply(async () => null)
+
+    mockPool.intercept({
+      path: '/test-no-status-code',
+      method: 'GET'
+    }).reply(async () => ({ data: 'hello' }))
+
+    await t.assert.rejects(
+      request(`${baseUrl}/test-resolve-null`, { dispatcher: mockAgent }),
+      new InvalidArgumentError('reply options callback must return an object')
+    )
+
+    await t.assert.rejects(
+      request(`${baseUrl}/test-no-status-code`, { dispatcher: mockAgent }),
+      new InvalidArgumentError('statusCode must be defined')
+    )
+  })
+
+  test('should reject with the reason of a rejected asynchronous callback', async t => {
+    t.plan(1)
+
+    const baseUrl = 'http://localhost:9999'
+    const mockAgent = new MockAgent()
+    mockAgent.disableNetConnect()
+    after(() => mockAgent.close())
+
+    const mockPool = mockAgent.get(baseUrl)
+
+    mockPool.intercept({
+      path: '/test',
+      method: 'GET'
+    }).reply(async () => {
+      throw new Error('kaboom')
+    })
+
+    await t.assert.rejects(request(`${baseUrl}/test`, { dispatcher: mockAgent }), new Error('kaboom'))
+  })
+})
+
 describe('MockInterceptor - replyWithError', () => {
   test('should return MockScope', t => {
     t.plan(1)

test/types/mock-interceptor.test-d.ts

@@ -54,6 +54,14 @@ expectAssignable<BodyInit | Dispatcher.DispatchOptions['body']>(mockResponseCall
     expectAssignable<MockInterceptor.MockResponseCallbackOptions['body']>(options.body)
     return { statusCode: 200, data: { foo: 'bar' } }
   })
+  expectAssignable<MockScope>(mockInterceptor.reply(async () => ({ statusCode: 200, data: { foo: 'bar' } })))
+  expectAssignable<MockScope>(mockInterceptor.reply(async () => ({
+    statusCode: 200,
+    data: { foo: 'bar' },
+    responseOptions: {
+      headers: { foo: 'bar' }
+    }
+  })))
 
   // replyWithError
   class CustomError extends Error {

types/mock-interceptor.d.ts

@@ -75,9 +75,13 @@ declare namespace MockInterceptor {
     opts: MockResponseCallbackOptions
   ) => TData | Buffer | string
 
+  export type MockReplyOptions<TData extends object = object> = {
+    statusCode: number, data?: TData | Buffer | string, responseOptions?: MockResponseOptions
+  }
+
   export type MockReplyOptionsCallback<TData extends object = object> = (
     opts: MockResponseCallbackOptions
-  ) => { statusCode: number, data?: TData | Buffer | string, responseOptions?: MockResponseOptions }
+  ) => MockReplyOptions<TData> | Promise<MockReplyOptions<TData>>
 }
 
 interface Interceptable extends Dispatcher {