fix(mock): preserve lifecycle abort semantics

Signed-off-by: marko1olo <barsukdana@gmail.com>

Author: marko1olo

lib/mock/mock-utils.js

@@ -18,6 +18,8 @@ const {
 } = require('node:util')
 const { InvalidArgumentError } = require('../core/errors')
 
+const requestAborted = Symbol('request aborted')
+
 function matchValue (match, value) {
   if (typeof match === 'string') {
     return match === value
@@ -295,13 +297,9 @@ function mockDispatch (opts, handler) {
 
   mockDispatch.timesInvoked++
 
-  // 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
+  const { data: response, delay, persist } = mockDispatch
+  const { error } = response
   const { timesInvoked, times } = mockDispatch
 
   // If it's used up and not persistent, mark as consumed
@@ -359,7 +357,7 @@ function mockDispatch (opts, handler) {
 
   if (isPromise(requestBody)) {
     requestBody.then((body) => {
-      if (body === null) {
+      if (body === requestAborted) {
         return
       }
       if (body !== opts.body) {
@@ -370,7 +368,7 @@ function mockDispatch (opts, handler) {
     return true
   }
 
-  if (requestBody === null) {
+  if (requestBody === requestAborted) {
     return true
   }
 
@@ -392,19 +390,27 @@ function mockDispatch (opts, handler) {
     }
   }
 
-  function handleReply (mockDispatches, _data = data) {
+  function handleReply (mockDispatches, _response = response) {
     // Don't send response if the request was aborted
     if (aborted) {
       return
     }
 
+    if (_response.callback) {
+      const { callback, ...responseDefaults } = _response
+      mockDispatch.data = { ...responseDefaults, ...callback(replyOpts) }
+      return handleReply(mockDispatches, mockDispatch.data)
+    }
+
+    const { statusCode, data, headers, trailers } = _response
+
     // fetch's HeadersList is a 1D string array
     const optsHeaders = Array.isArray(opts.headers)
       ? buildHeadersFromArray(opts.headers)
       : opts.headers
-    const body = typeof _data === 'function'
-      ? _data({ ...replyOpts, headers: optsHeaders })
-      : _data
+    const body = typeof data === 'function'
+      ? data({ ...replyOpts, headers: optsHeaders })
+      : data
 
     // util.types.isPromise is likely needed for jest.
     if (isPromise(body)) {
@@ -413,7 +419,7 @@ function mockDispatch (opts, handler) {
       // synchronously throw the error, which breaks some tests.
       // Rather, we wait for the callback to resolve if it is a
       // promise, and then re-run handleReply with the new body.
-      return body.then((newData) => handleReply(mockDispatches, newData))
+      return body.then((newData) => handleReply(mockDispatches, { ..._response, data: newData }))
     }
 
     // Check again if aborted after async body resolution
@@ -444,7 +450,7 @@ function dispatchRequestBody (body, handler, controller, isAborted) {
   }
 
   if (body == null) {
-    return callOnRequestSent(handler, controller) ? body : null
+    return callOnRequestSent(handler, controller, isAborted) ? body : requestAborted
   }
 
   if (body && typeof body[Symbol.asyncIterator] === 'function') {
@@ -456,42 +462,42 @@ function dispatchRequestBody (body, handler, controller, isAborted) {
 
     for (const chunk of body) {
       if (isAborted()) {
-        return null
+        return requestAborted
       }
       chunks.push(chunk)
-      if (!callOnBodySent(handler, controller, chunk)) {
-        return null
+      if (!callOnBodySent(handler, controller, chunk) || isAborted()) {
+        return requestAborted
       }
     }
 
-    return !isAborted() && callOnRequestSent(handler, controller) ? chunks : null
+    return callOnRequestSent(handler, controller, isAborted) ? chunks : requestAborted
   }
 
   if (isAborted()) {
-    return null
+    return requestAborted
   }
   if (!callOnBodySent(handler, controller, body)) {
-    return null
+    return requestAborted
   }
 
-  return !isAborted() && callOnRequestSent(handler, controller) ? body : null
+  return callOnRequestSent(handler, controller, isAborted) ? body : requestAborted
 }
 
 async function dispatchAsyncIterableBody (body, handler, controller, isAborted) {
   const chunks = []
 
   for await (const chunk of body) {
     if (isAborted()) {
-      return null
+      return requestAborted
     }
     chunks.push(chunk)
-    if (!callOnBodySent(handler, controller, chunk)) {
-      return null
+    if (!callOnBodySent(handler, controller, chunk) || isAborted()) {
+      return requestAborted
     }
   }
 
-  if (isAborted() || !callOnRequestSent(handler, controller)) {
-    return null
+  if (!callOnRequestSent(handler, controller, isAborted)) {
+    return requestAborted
   }
 
   return {
@@ -511,10 +517,10 @@ function callOnBodySent (handler, controller, chunk) {
   }
 }
 
-function callOnRequestSent (handler, controller) {
+function callOnRequestSent (handler, controller, isAborted) {
   try {
     handler.onRequestSent?.()
-    return true
+    return !isAborted()
   } catch (error) {
     controller.abort(error)
     return false