feat(mock): add support reply option callbacks

docs/api/MockPool.md

@@ -113,7 +113,7 @@ for await (const data of body) {
 }
 ```
 
-#### Example - Mocked request using reply callbacks
+#### Example - Mocked request using reply data callbacks
 
 ```js
 import { MockAgent, setGlobalDispatcher, request } from 'undici'
@@ -146,6 +146,39 @@ for await (const data of body) {
 }
 ```
 
+#### Example - Mocked request using reply options callback
+
+```js
+import { MockAgent, setGlobalDispatcher, request } from 'undici'
+
+const mockAgent = new MockAgent()
+setGlobalDispatcher(mockAgent)
+
+const mockPool = mockAgent.get('http://localhost:3000')
+
+mockPool.intercept({
+  path: '/echo',
+  method: 'GET',
+  headers: {
+    'User-Agent': 'undici',
+    Host: 'example.com'
+  }
+}).reply(({ headers }) => ({ statusCode: 200, data: { message: headers.get('message') }})))
+
+const { statusCode, body, headers } = await request('http://localhost:3000', {
+  headers: {
+    message: 'hello world!'
+  }
+})
+
+console.log('response received', statusCode) // response received 200
+console.log('headers', headers) // { 'content-type': 'application/json' }
+
+for await (const data of body) {
+  console.log('data', data.toString('utf8')) // { "message":"hello world!" }
+}
+```
+
 #### Example - Basic Mocked requests with multiple intercepts
 
 ```js

lib/mock/mock-interceptor.js

@@ -9,7 +9,7 @@ const {
   kContentLength,
   kMockDispatch
 } = require('./mock-symbols')
-const { InvalidArgumentError } = require('../core/errors')
+const { InvalidArgumentError, InvalidReturnValueError } = require('../core/errors')
 
 /**
  * Defines the scope API for a interceptor reply
@@ -74,10 +74,16 @@ class MockInterceptor {
     this[kContentLength] = false
   }
 
-  /**
-   * Mock an undici request with a defined reply.
-   */
-  reply (statusCode, data, responseOptions = {}) {
+  createMockScopeDispatchData(statusCode, data, responseOptions = {}) {
+    const responseData = getResponseData(data)
+    const contentLength = this[kContentLength] ? { 'content-length': responseData.length } : {}
+    const headers = { ...this[kDefaultHeaders], ...contentLength, ...responseOptions.headers }
+    const trailers = { ...this[kDefaultTrailers], ...responseOptions.trailers }
+
+    return { statusCode, data, headers, trailers };
+  }
+
+  validateReplyParameters(statusCode, data, responseOptions) {
     if (typeof statusCode === 'undefined') {
       throw new InvalidArgumentError('statusCode must be defined')
     }
@@ -87,13 +93,53 @@ class MockInterceptor {
     if (typeof responseOptions !== 'object') {
       throw new InvalidArgumentError('responseOptions must be an object')
     }
+  }
 
-    const responseData = getResponseData(data)
-    const contentLength = this[kContentLength] ? { 'content-length': responseData.length } : {}
-    const headers = { ...this[kDefaultHeaders], ...contentLength, ...responseOptions.headers }
-    const trailers = { ...this[kDefaultTrailers], ...responseOptions.trailers }
-    const newMockDispatch = addMockDispatch(this[kDispatches], this[kDispatchKey], { statusCode, data, headers, trailers })
+  /**
+   * Mock an undici request with a defined reply.
+   */
+  reply (replyData) {
+    // Values of reply aren't available right now as they
+    // can only be available when the reply callback is invoked.
+    if (typeof replyData === '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 = replyData(opts);
+
+        // Check if it is in the right format
+        if (typeof resolvedData !== 'object') {
+          throw new InvalidArgumentError('reply options callback must return an object')
+        }
+
+        const { statusCode, data, responseOptions = {}} = resolvedData;
+        this.validateReplyParameters(statusCode, data, responseOptions);
+        // Since the values can be obtained immediately we return them
+        // from this higher order function that will be resolved later.
+        return { 
+          ...this.createMockScopeDispatchData(statusCode, data, responseOptions)
+        }
+      }
+
+      // 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)
+      return new MockScope(newMockDispatch);
+    }
+
+    // We can have either one or three parameters, if we get here,
+    // we should have 2-3 parameters. So we spread the arguments of
+    // this function to obtain the parameters, since replyData will always
+    // just be the statusCode. 
+    const [statusCode, data, responseOptions = {}] = [...arguments];   
+    this.validateReplyParameters(statusCode, data, responseOptions);
+
+    // Send in-already provided data like usual
+    const dispatchData = this.createMockScopeDispatchData(statusCode, data, responseOptions);
+    const newMockDispatch = addMockDispatch(this[kDispatches], this[kDispatchKey], dispatchData)
     return new MockScope(newMockDispatch)
+
   }
 
   /**

lib/mock/mock-utils.js

@@ -90,7 +90,8 @@ function getMockDispatch (mockDispatches, key) {
 
 function addMockDispatch (mockDispatches, key, data) {
   const baseData = { times: null, persist: false, consumed: false }
-  const newMockDispatch = { ...baseData, ...key, data: { error: null, ...data } }
+  const replyData = typeof data === 'function' ? { callback: data } : { ...data };
+  const newMockDispatch = { ...baseData, ...key, data: { error: null, ...replyData } }
   mockDispatches.push(newMockDispatch)
   return newMockDispatch
 }
@@ -135,7 +136,12 @@ async function getResponse (body) {
 function mockDispatch (opts, handler) {
   // Get mock dispatch from built key
   const key = buildKey(opts)
-  const mockDispatch = getMockDispatch(this[kDispatches], key)
+  let mockDispatch = getMockDispatch(this[kDispatches], key)
+
+  // 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) }
+  }
 
   // Parse mockDispatch data
   const { data: { statusCode, data, headers, trailers, error }, delay, persist } = mockDispatch

test/mock-interceptor.js

@@ -2,6 +2,7 @@
 
 const { test } = require('tap')
 const { MockInterceptor, MockScope } = require('../lib/mock/mock-interceptor')
+const MockAgent = require('../lib/mock/mock-agent');
 const { InvalidArgumentError } = require('../lib/core/errors')
 
 test('MockInterceptor - reply', t => {
@@ -53,6 +54,125 @@ test('MockInterceptor - reply callback', t => {
     t.throws(() => mockInterceptor.reply(), new InvalidArgumentError('statusCode must be defined'))
     t.throws(() => mockInterceptor.reply(200, () => {}, 'hello'), new InvalidArgumentError('responseOptions must be an object'))
   })
+})
+
+test('MockInterceptor - reply options callback', t => {
+  t.plan(2)
+
+  t.test('should return MockScope', t => {
+    t.plan(2)
+
+    const mockInterceptor = new MockInterceptor({
+      path: '',
+      method: ''
+    }, [])
+    const result = mockInterceptor.reply((options) => ({
+      statusCode: 200,
+      data: 'hello'
+    }))
+    t.type(result, MockScope)
+
+    // Test parameters
+
+    const baseUrl = 'http://localhost:9999'
+    const mockAgent = new MockAgent()
+    t.teardown(mockAgent.close.bind(mockAgent))
+
+    const mockPool = mockAgent.get(baseUrl)
+
+    mockPool.intercept({
+      path: '/test',
+      method: 'GET',
+    }).reply((options) => {
+      t.strictSame(options, { path: '/test', method: 'GET', headers: { foo: 'bar' }})
+      return { statusCode: 200, data: 'hello' }
+    });
+
+    mockPool.dispatch({
+      path: '/test',
+      method: 'GET',
+      headers: { foo: 'bar' }
+    }, {
+      onHeaders: () => {},
+      onData: () => {},
+      onComplete: () => {},
+    })
+  })
+
+  t.test('should error if passed options invalid', async (t) => {
+    t.plan(4)
+
+    const baseUrl = 'http://localhost:9999'
+    const mockAgent = new MockAgent()
+    t.teardown(mockAgent.close.bind(mockAgent))
+
+    const mockPool = mockAgent.get(baseUrl)
+
+    mockPool.intercept({
+      path: '/test',
+      method: 'GET',
+    }).reply(() => {});
+
+    mockPool.intercept({
+      path: '/test2',
+      method: 'GET',
+    }).reply(() => ({
+      statusCode: 200,
+    }));
+
+    mockPool.intercept({
+      path: '/test3',
+      method: 'GET',
+    }).reply(() => ({
+      statusCode: 200,
+      data: 'hello',
+      responseOptions: 42,
+    }));
+
+    mockPool.intercept({
+      path: '/test4',
+      method: 'GET',
+    }).reply(() => ({
+      data: 'hello',
+      responseOptions: 42,
+    }));
+
+    t.throws(() => mockPool.dispatch({
+      path: '/test',
+      method: 'GET',
+    }, {
+      onHeaders: () => {},
+      onData: () => {},
+      onComplete: () => {},
+    }), new InvalidArgumentError('reply options callback must return an object'))
+
+    t.throws(() => mockPool.dispatch({
+      path: '/test2',
+      method: 'GET',
+    }, {
+      onHeaders: () => {},
+      onData: () => {},
+      onComplete: () => {},
+    }), new InvalidArgumentError('data must be defined'))
+
+    t.throws(() => mockPool.dispatch({
+      path: '/test3',
+      method: 'GET',
+    }, {
+      onHeaders: () => {},
+      onData: () => {},
+      onComplete: () => {},
+    }), new InvalidArgumentError('responseOptions must be an object'))
+
+    t.throws(() => mockPool.dispatch({
+      path: '/test4',
+      method: 'GET',
+    }, {
+      onHeaders: () => {},
+      onData: () => {},
+      onComplete: () => {},
+    }), new InvalidArgumentError('statusCode must be defined'))
+  })
 }) 
 
 test('MockInterceptor - replyWithError', t => {

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

@@ -19,6 +19,17 @@ import { MockInterceptor, MockScope } from '../../types/mock-interceptor'
   expectAssignable<MockScope>(mockInterceptor.reply(200, () => ({}), { trailers: { foo: 'bar' }}))
   expectAssignable<MockScope<{ foo: string }>>(mockInterceptor.reply<{ foo: string }>(200, { foo: 'bar' }))
   expectAssignable<MockScope<{ foo: string }>>(mockInterceptor.reply<{ foo: string }>(200, () => ({ foo: 'bar' })))
+  expectAssignable<MockScope>(mockInterceptor.reply(() => ({ statusCode: 200, data: { foo: 'bar' }})))
+  expectAssignable<MockScope>(mockInterceptor.reply(() => ({ statusCode: 200, data: { foo: 'bar' }, responseOptions: {
+    headers: { foo: 'bar' }
+  }})))
+  expectAssignable<MockScope>(mockInterceptor.reply((options) => { 
+    expectAssignable<MockInterceptor.MockResponseCallbackOptions>(options);
+    return { statusCode: 200, data: { foo: 'bar'}
+  }}))
+  expectAssignable<MockScope>(mockInterceptor.reply(() => ({ statusCode: 200, data: { foo: 'bar' }, responseOptions: {
+    trailers: { foo: 'bar' }
+  }})))
 
   // replyWithError
   class CustomError extends Error {

types/mock-interceptor.d.ts

@@ -21,6 +21,7 @@ declare class MockScope<TData extends object = object> {
 declare class MockInterceptor {
   constructor(options: MockInterceptor.Options, mockDispatches: MockInterceptor.MockDispatch[]);
   /** Mock an undici request with the defined reply. */
+  reply<TData extends object = object>(replyOptionsCallback: MockInterceptor.MockReplyOptionsCallback<TData>): MockScope<TData>;
   reply<TData extends object = object>(
     statusCode: number,
     data: TData | Buffer | string | MockInterceptor.MockResponseDataHandler<TData>,
@@ -76,6 +77,10 @@ declare namespace MockInterceptor {
   export type MockResponseDataHandler<TData extends object = object> = (
     opts: MockResponseCallbackOptions
   ) => TData | Buffer | string;
+
+  export type MockReplyOptionsCallback<TData extends object = object> = (
+    opts: MockResponseCallbackOptions
+  ) => { statusCode: number, data: TData | Buffer | string, responseOptions?: MockResponseOptions }
 }
 
 interface Interceptable {