dd-trace 4.52.0 → 4.53.0

package/packages/datadog-instrumentations/src/jest.js

@@ -133,6 +133,7 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
       this.isEarlyFlakeDetectionEnabled = this.testEnvironmentOptions._ddIsEarlyFlakeDetectionEnabled
       this.isFlakyTestRetriesEnabled = this.testEnvironmentOptions._ddIsFlakyTestRetriesEnabled
       this.flakyTestRetriesCount = this.testEnvironmentOptions._ddFlakyTestRetriesCount
+      this.isDiEnabled = this.testEnvironmentOptions._ddIsDiEnabled
 
       if (this.isEarlyFlakeDetectionEnabled) {
         const hasKnownTests = !!knownTests.jest
@@ -237,7 +238,6 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
             name: removeEfdStringFromTestName(testName),
             suite: this.testSuite,
             testSourceFile: this.testSourceFile,
-            runner: 'jest-circus',
             displayName: this.displayName,
             testParameters,
             frameworkVersion: jestVersion,
@@ -274,13 +274,23 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
         }
       }
       if (event.name === 'test_done') {
+        const probe = {}
         const asyncResource = asyncResources.get(event.test)
         asyncResource.runInAsyncScope(() => {
           let status = 'pass'
           if (event.test.errors && event.test.errors.length) {
             status = 'fail'
-            const formattedError = formatJestError(event.test.errors[0])
-            testErrCh.publish(formattedError)
+            const numRetries = this.global[RETRY_TIMES]
+            const numTestExecutions = event.test?.invocations
+            const willBeRetried = numRetries > 0 && numTestExecutions - 1 < numRetries
+
+            const error = formatJestError(event.test.errors[0])
+            testErrCh.publish({
+              error,
+              willBeRetried,
+              probe,
+              isDiEnabled: this.isDiEnabled
+            })
           }
           testRunFinishCh.publish({
             status,
@@ -302,6 +312,9 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
             }
           }
         })
+        if (probe.setProbePromise) {
+          await probe.setProbePromise
+        }
       }
       if (event.name === 'test_skip' || event.name === 'test_todo') {
         const asyncResource = new AsyncResource('bound-anonymous-fn')
@@ -310,7 +323,6 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
             name: getJestTestName(event.test),
             suite: this.testSuite,
             testSourceFile: this.testSourceFile,
-            runner: 'jest-circus',
             displayName: this.displayName,
             frameworkVersion: jestVersion,
             testStartLine: getTestLineStart(event.test.asyncError, this.testSuite)
@@ -445,7 +457,7 @@ function cliWrapper (cli, jestVersion) {
         earlyFlakeDetectionFaultyThreshold = libraryConfig.earlyFlakeDetectionFaultyThreshold
       }
     } catch (err) {
-      log.error(err)
+      log.error('Jest library configuration error', err)
     }
 
     if (isEarlyFlakeDetectionEnabled) {
@@ -466,7 +478,7 @@ function cliWrapper (cli, jestVersion) {
           isEarlyFlakeDetectionEnabled = false
         }
       } catch (err) {
-        log.error(err)
+        log.error('Jest known tests error', err)
       }
     }
 
@@ -485,7 +497,7 @@ function cliWrapper (cli, jestVersion) {
           skippableSuites = receivedSkippableSuites
         }
       } catch (err) {
-        log.error(err)
+        log.error('Jest test-suite skippable error', err)
       }
     }
 
@@ -780,6 +792,7 @@ addHook({
       _ddRepositoryRoot,
       _ddIsFlakyTestRetriesEnabled,
       _ddFlakyTestRetriesCount,
+      _ddIsDiEnabled,
       ...restOfTestEnvironmentOptions
     } = testEnvironmentOptions
 
@@ -855,12 +868,18 @@ addHook({
 
 const LIBRARIES_BYPASSING_JEST_REQUIRE_ENGINE = [
   'selenium-webdriver',
+  'selenium-webdriver/chrome',
+  'selenium-webdriver/edge',
+  'selenium-webdriver/safari',
+  'selenium-webdriver/firefox',
+  'selenium-webdriver/ie',
+  'selenium-webdriver/chromium',
   'winston'
 ]
 
 function shouldBypassJestRequireEngine (moduleName) {
   return (
-    LIBRARIES_BYPASSING_JEST_REQUIRE_ENGINE.some(library => moduleName.includes(library))
+    LIBRARIES_BYPASSING_JEST_REQUIRE_ENGINE.includes(moduleName)
   )
 }
 

package/packages/datadog-instrumentations/src/mocha/utils.js

@@ -284,8 +284,9 @@ function getOnTestRetryHandler () {
     const asyncResource = getTestAsyncResource(test)
     if (asyncResource) {
       const isFirstAttempt = test._currentRetry === 0
+      const willBeRetried = test._currentRetry < test._retries
       asyncResource.runInAsyncScope(() => {
-        testRetryCh.publish({ isFirstAttempt, err })
+        testRetryCh.publish({ isFirstAttempt, err, willBeRetried })
       })
     }
     const key = getTestToArKey(test)

package/packages/datadog-instrumentations/src/mysql2.js

@@ -8,7 +8,7 @@ const {
 const shimmer = require('../../datadog-shimmer')
 const semver = require('semver')
 
-addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, (Connection, version) => {
+function wrapConnection (Connection, version) {
   const startCh = channel('apm:mysql2:query:start')
   const finishCh = channel('apm:mysql2:query:finish')
   const errorCh = channel('apm:mysql2:query:error')
@@ -151,9 +151,8 @@ addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['>=1'] }, (Conne
       }
     }, cmd))
   }
-})
-
-addHook({ name: 'mysql2', file: 'lib/pool.js', versions: ['>=1'] }, (Pool, version) => {
+}
+function wrapPool (Pool, version) {
   const startOuterQueryCh = channel('datadog:mysql2:outerquery:start')
   const shouldEmitEndAfterQueryAbort = semver.intersects(version, '>=1.3.3')
 
@@ -221,10 +220,9 @@ addHook({ name: 'mysql2', file: 'lib/pool.js', versions: ['>=1'] }, (Pool, versi
   })
 
   return Pool
-})
+}
 
-// PoolNamespace.prototype.query does not exist in mysql2<2.3.0
-addHook({ name: 'mysql2', file: 'lib/pool_cluster.js', versions: ['>=2.3.0'] }, PoolCluster => {
+function wrapPoolCluster (PoolCluster) {
   const startOuterQueryCh = channel('datadog:mysql2:outerquery:start')
   const wrappedPoolNamespaces = new WeakSet()
 
@@ -297,4 +295,11 @@ addHook({ name: 'mysql2', file: 'lib/pool_cluster.js', versions: ['>=2.3.0'] },
   })
 
   return PoolCluster
-})
+}
+
+addHook({ name: 'mysql2', file: 'lib/base/connection.js', versions: ['>=3.11.5'] }, wrapConnection)
+addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['1 - 3.11.4'] }, wrapConnection)
+addHook({ name: 'mysql2', file: 'lib/pool.js', versions: ['1 - 3.11.4'] }, wrapPool)
+
+// PoolNamespace.prototype.query does not exist in mysql2<2.3.0
+addHook({ name: 'mysql2', file: 'lib/pool_cluster.js', versions: ['2.3.0 - 3.11.4'] }, wrapPoolCluster)

package/packages/datadog-instrumentations/src/next.js

@@ -2,7 +2,6 @@
 
 const { channel, addHook } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
-const { DD_MAJOR } = require('../../../version')
 
 const startChannel = channel('apm:next:request:start')
 const finishChannel = channel('apm:next:request:finish')
@@ -221,7 +220,7 @@ addHook({
 
 addHook({
   name: 'next',
-  versions: DD_MAJOR >= 4 ? ['>=10.2 <11.1'] : ['>=9.5 <11.1'],
+  versions: ['>=10.2 <11.1'],
   file: 'dist/next-server/server/serve-static.js'
 }, serveStatic => shimmer.wrap(serveStatic, 'serveStatic', wrapServeStatic))
 
@@ -248,7 +247,11 @@ addHook({ name: 'next', versions: ['>=13.2'], file: 'dist/server/next-server.js'
   return nextServer
 })
 
-addHook({ name: 'next', versions: ['>=11.1 <13.2'], file: 'dist/server/next-server.js' }, nextServer => {
+addHook({
+  name: 'next',
+  versions: ['>=11.1 <13.2'],
+  file: 'dist/server/next-server.js'
+}, nextServer => {
   const Server = nextServer.default
   shimmer.wrap(Server.prototype, 'handleApiRequest', wrapHandleApiRequest)
   return nextServer
@@ -256,7 +259,7 @@ addHook({ name: 'next', versions: ['>=11.1 <13.2'], file: 'dist/server/next-serv
 
 addHook({
   name: 'next',
-  versions: DD_MAJOR >= 4 ? ['>=10.2 <11.1'] : ['>=9.5 <11.1'],
+  versions: ['>=10.2 <11.1'],
   file: 'dist/next-server/server/next-server.js'
 }, nextServer => {
   const Server = nextServer.default

package/packages/datadog-instrumentations/src/passport-http.js

@@ -1,22 +1,10 @@
 'use strict'
 
-const shimmer = require('../../datadog-shimmer')
 const { addHook } = require('./helpers/instrument')
-const { wrapVerify } = require('./passport-utils')
+const { strategyHook } = require('./passport-utils')
 
 addHook({
   name: 'passport-http',
   file: 'lib/passport-http/strategies/basic.js',
   versions: ['>=0.3.0']
-}, BasicStrategy => {
-  return shimmer.wrapFunction(BasicStrategy, BasicStrategy => function () {
-    const type = 'http'
-
-    if (typeof arguments[0] === 'function') {
-      arguments[0] = wrapVerify(arguments[0], false, type)
-    } else {
-      arguments[1] = wrapVerify(arguments[1], (arguments[0] && arguments[0].passReqToCallback), type)
-    }
-    return BasicStrategy.apply(this, arguments)
-  })
-})
+}, strategyHook)

package/packages/datadog-instrumentations/src/passport-local.js

@@ -1,22 +1,10 @@
 'use strict'
 
-const shimmer = require('../../datadog-shimmer')
 const { addHook } = require('./helpers/instrument')
-const { wrapVerify } = require('./passport-utils')
+const { strategyHook } = require('./passport-utils')
 
 addHook({
   name: 'passport-local',
   file: 'lib/strategy.js',
   versions: ['>=1.0.0']
-}, Strategy => {
-  return shimmer.wrapFunction(Strategy, Strategy => function () {
-    const type = 'local'
-
-    if (typeof arguments[0] === 'function') {
-      arguments[0] = wrapVerify(arguments[0], false, type)
-    } else {
-      arguments[1] = wrapVerify(arguments[1], (arguments[0] && arguments[0].passReqToCallback), type)
-    }
-    return Strategy.apply(this, arguments)
-  })
-})
+}, strategyHook)

package/packages/datadog-instrumentations/src/passport-utils.js

@@ -5,33 +5,57 @@ const { channel } = require('./helpers/instrument')
 
 const passportVerifyChannel = channel('datadog:passport:verify:finish')
 
-function wrapVerifiedAndPublish (username, password, verified, type) {
-  if (!passportVerifyChannel.hasSubscribers) {
-    return verified
-  }
+function wrapVerifiedAndPublish (framework, username, verified) {
+  return shimmer.wrapFunction(verified, function wrapVerified (verified) {
+    return function wrappedVerified (err, user) {
+      // if there is an error, it's neither an auth success nor a failure
+      if (!err) {
+        const abortController = new AbortController()
+
+        passportVerifyChannel.publish({ framework, login: username, user, success: !!user, abortController })
+
+        if (abortController.signal.aborted) return
+      }
 
-  // eslint-disable-next-line n/handle-callback-err
-  return shimmer.wrapFunction(verified, verified => function (err, user, info) {
-    const credentials = { type, username }
-    passportVerifyChannel.publish({ credentials, user })
-    return verified.apply(this, arguments)
+      return verified.apply(this, arguments)
+    }
   })
 }
 
-function wrapVerify (verify, passReq, type) {
-  if (passReq) {
-    return function (req, username, password, verified) {
-      arguments[3] = wrapVerifiedAndPublish(username, password, verified, type)
-      return verify.apply(this, arguments)
+function wrapVerify (verify) {
+  return function wrappedVerify (req, username, password, verified) {
+    if (passportVerifyChannel.hasSubscribers) {
+      const framework = `passport-${this.name}`
+
+      // replace the callback with our own wrapper to get the result
+      if (this._passReqToCallback) {
+        arguments[3] = wrapVerifiedAndPublish(framework, arguments[1], arguments[3])
+      } else {
+        arguments[2] = wrapVerifiedAndPublish(framework, arguments[0], arguments[2])
+      }
     }
-  } else {
-    return function (username, password, verified) {
-      arguments[2] = wrapVerifiedAndPublish(username, password, verified, type)
-      return verify.apply(this, arguments)
+
+    return verify.apply(this, arguments)
+  }
+}
+
+function wrapStrategy (Strategy) {
+  return function wrappedStrategy () {
+    // verify function can be either the first or second argument
+    if (typeof arguments[0] === 'function') {
+      arguments[0] = wrapVerify(arguments[0])
+    } else {
+      arguments[1] = wrapVerify(arguments[1])
     }
+
+    return Strategy.apply(this, arguments)
   }
 }
 
+function strategyHook (Strategy) {
+  return shimmer.wrapFunction(Strategy, wrapStrategy)
+}
+
 module.exports = {
-  wrapVerify
+  strategyHook
 }

package/packages/datadog-instrumentations/src/pg.js

@@ -62,17 +62,17 @@ function wrapQuery (query) {
         abortController
       })
 
-      const finish = asyncResource.bind(function (error) {
+      const finish = asyncResource.bind(function (error, res) {
         if (error) {
           errorCh.publish(error)
         }
-        finishCh.publish()
+        finishCh.publish({ result: res?.rows })
       })
 
       if (abortController.signal.aborted) {
         const error = abortController.signal.reason || new Error('Aborted')
 
-        // eslint-disable-next-line max-len
+        // eslint-disable-next-line @stylistic/js/max-len
         // Based on: https://github.com/brianc/node-postgres/blob/54eb0fa216aaccd727765641e7d1cf5da2bc483d/packages/pg/lib/client.js#L510
         const reusingQuery = typeof pgQuery.submit === 'function'
         const callback = arguments[arguments.length - 1]
@@ -119,15 +119,15 @@ function wrapQuery (query) {
       if (newQuery.callback) {
         const originalCallback = callbackResource.bind(newQuery.callback)
         newQuery.callback = function (err, res) {
-          finish(err)
+          finish(err, res)
           return originalCallback.apply(this, arguments)
         }
       } else if (newQuery.once) {
         newQuery
           .once('error', finish)
-          .once('end', () => finish())
+          .once('end', (res) => finish(null, res))
       } else {
-        newQuery.then(() => finish(), finish)
+        newQuery.then((res) => finish(null, res), finish)
       }
 
       try {

package/packages/datadog-instrumentations/src/playwright.js

@@ -47,7 +47,7 @@ function isNewTest (test) {
   const testSuite = getTestSuitePath(test._requireFile, rootDir)
   const testsForSuite = knownTests?.playwright?.[testSuite] || []
 
-  return !testsForSuite.includes(test.title)
+  return !testsForSuite.includes(getTestFullname(test))
 }
 
 function getSuiteType (test, type) {
@@ -224,10 +224,21 @@ function testWillRetry (test, testStatus) {
   return testStatus === 'fail' && test.results.length <= test.retries
 }
 
+function getTestFullname (test) {
+  let parent = test.parent
+  const names = [test.title]
+  while (parent?._type === 'describe' || parent?._isDescribe) {
+    if (parent.title) {
+      names.unshift(parent.title)
+    }
+    parent = parent.parent
+  }
+  return names.join(' ')
+}
+
 function testBeginHandler (test, browserName) {
   const {
     _requireFile: testSuiteAbsolutePath,
-    title: testName,
     _type,
     location: {
       line: testSourceLine
@@ -238,6 +249,8 @@ function testBeginHandler (test, browserName) {
     return
   }
 
+  const testName = getTestFullname(test)
+
   const isNewTestSuite = !startedSuites.includes(testSuiteAbsolutePath)
 
   if (isNewTestSuite) {
@@ -412,7 +425,7 @@ function runnerHook (runnerExport, playwrightVersion) {
       }
     } catch (e) {
       isEarlyFlakeDetectionEnabled = false
-      log.error(e)
+      log.error('Playwright session start error', e)
     }
 
     if (isEarlyFlakeDetectionEnabled && semver.gte(playwrightVersion, MINIMUM_SUPPORTED_VERSION_EFD)) {
@@ -425,7 +438,7 @@ function runnerHook (runnerExport, playwrightVersion) {
         }
       } catch (err) {
         isEarlyFlakeDetectionEnabled = false
-        log.error(err)
+        log.error('Playwright known tests error', err)
       }
     }
 

package/packages/datadog-instrumentations/src/router.js

@@ -169,11 +169,107 @@ function createWrapRouterMethod (name) {
 
 const wrapRouterMethod = createWrapRouterMethod('router')
 
-addHook({ name: 'router', versions: ['>=1'] }, Router => {
+addHook({ name: 'router', versions: ['>=1 <2'] }, Router => {
   shimmer.wrap(Router.prototype, 'use', wrapRouterMethod)
   shimmer.wrap(Router.prototype, 'route', wrapRouterMethod)
 
   return Router
 })
 
+const queryParserReadCh = channel('datadog:query:read:finish')
+
+addHook({ name: 'router', versions: ['>=2'] }, Router => {
+  const WrappedRouter = shimmer.wrapFunction(Router, function (originalRouter) {
+    return function wrappedMethod () {
+      const router = originalRouter.apply(this, arguments)
+
+      shimmer.wrap(router, 'handle', function wrapHandle (originalHandle) {
+        return function wrappedHandle (req, res, next) {
+          const abortController = new AbortController()
+
+          if (queryParserReadCh.hasSubscribers && req) {
+            queryParserReadCh.publish({ req, res, query: req.query, abortController })
+
+            if (abortController.signal.aborted) return
+          }
+
+          return originalHandle.apply(this, arguments)
+        }
+      })
+
+      return router
+    }
+  })
+
+  shimmer.wrap(WrappedRouter.prototype, 'use', wrapRouterMethod)
+  shimmer.wrap(WrappedRouter.prototype, 'route', wrapRouterMethod)
+
+  return WrappedRouter
+})
+
+const routerParamStartCh = channel('datadog:router:param:start')
+const visitedParams = new WeakSet()
+
+function wrapHandleRequest (original) {
+  return function wrappedHandleRequest (req, res, next) {
+    if (routerParamStartCh.hasSubscribers && Object.keys(req.params).length && !visitedParams.has(req.params)) {
+      visitedParams.add(req.params)
+
+      const abortController = new AbortController()
+
+      routerParamStartCh.publish({
+        req,
+        res,
+        params: req?.params,
+        abortController
+      })
+
+      if (abortController.signal.aborted) return
+    }
+
+    return original.apply(this, arguments)
+  }
+}
+
+addHook({
+  name: 'router', file: 'lib/layer.js', versions: ['>=2']
+}, Layer => {
+  shimmer.wrap(Layer.prototype, 'handleRequest', wrapHandleRequest)
+  return Layer
+})
+
+function wrapParam (original) {
+  return function wrappedProcessParams () {
+    arguments[1] = shimmer.wrapFunction(arguments[1], (originalFn) => {
+      return function wrappedFn (req, res) {
+        if (routerParamStartCh.hasSubscribers && Object.keys(req.params).length && !visitedParams.has(req.params)) {
+          visitedParams.add(req.params)
+
+          const abortController = new AbortController()
+
+          routerParamStartCh.publish({
+            req,
+            res,
+            params: req?.params,
+            abortController
+          })
+
+          if (abortController.signal.aborted) return
+        }
+
+        return originalFn.apply(this, arguments)
+      }
+    })
+
+    return original.apply(this, arguments)
+  }
+}
+
+addHook({
+  name: 'router', versions: ['>=2']
+}, router => {
+  shimmer.wrap(router.prototype, 'param', wrapParam)
+  return router
+})
+
 module.exports = { createWrapRouterMethod }

package/packages/datadog-instrumentations/src/sequelize.js

@@ -13,7 +13,7 @@ addHook({ name: 'sequelize', versions: ['>=4'] }, Sequelize => {
   const finishCh = channel('datadog:sequelize:query:finish')
 
   shimmer.wrap(Sequelize.prototype, 'query', query => {
-    return function (sql) {
+    return function (sql, options) {
       if (!startCh.hasSubscribers) {
         return query.apply(this, arguments)
       }
@@ -27,9 +27,14 @@ addHook({ name: 'sequelize', versions: ['>=4'] }, Sequelize => {
         dialect = this.dialect.name
       }
 
-      function onFinish () {
+      function onFinish (result) {
+        const type = options?.type || 'RAW'
+        if (type === 'RAW' && result?.length > 1) {
+          result = result[0]
+        }
+
         asyncResource.bind(function () {
-          finishCh.publish()
+          finishCh.publish({ result })
         }, this).apply(this)
       }
 
@@ -40,7 +45,7 @@ addHook({ name: 'sequelize', versions: ['>=4'] }, Sequelize => {
         })
 
         const promise = query.apply(this, arguments)
-        promise.then(onFinish, onFinish)
+        promise.then(onFinish, () => { onFinish() })
 
         return promise
       }, this).apply(this, arguments)

package/packages/datadog-instrumentations/src/url.js

@@ -59,6 +59,10 @@ addHook({ name: names }, function (url) {
           isURL: true
         })
       }
+
+      static [Symbol.hasInstance] (instance) {
+        return instance instanceof URL
+      }
     }
   })
 

package/packages/datadog-instrumentations/src/vitest.js

@@ -117,6 +117,7 @@ function getSortWrapper (sort) {
     let isEarlyFlakeDetectionEnabled = false
     let earlyFlakeDetectionNumRetries = 0
     let isEarlyFlakeDetectionFaulty = false
+    let isDiEnabled = false
     let knownTests = {}
 
     try {
@@ -126,10 +127,12 @@ function getSortWrapper (sort) {
         flakyTestRetriesCount = libraryConfig.flakyTestRetriesCount
         isEarlyFlakeDetectionEnabled = libraryConfig.isEarlyFlakeDetectionEnabled
         earlyFlakeDetectionNumRetries = libraryConfig.earlyFlakeDetectionNumRetries
+        isDiEnabled = libraryConfig.isDiEnabled
       }
     } catch (e) {
       isFlakyTestRetriesEnabled = false
       isEarlyFlakeDetectionEnabled = false
+      isDiEnabled = false
     }
 
     if (isFlakyTestRetriesEnabled && !this.ctx.config.retry && flakyTestRetriesCount > 0) {
@@ -169,6 +172,15 @@ function getSortWrapper (sort) {
       }
     }
 
+    if (isDiEnabled) {
+      try {
+        const workspaceProject = this.ctx.getCoreWorkspaceProject()
+        workspaceProject._provided._ddIsDiEnabled = isDiEnabled
+      } catch (e) {
+        log.warn('Could not send Dynamic Instrumentation configuration to workers.')
+      }
+    }
+
     let testCodeCoverageLinesTotal
 
     if (this.ctx.coverageProvider?.generateCoverage) {
@@ -298,13 +310,16 @@ addHook({
     const testName = getTestName(task)
     let isNew = false
     let isEarlyFlakeDetectionEnabled = false
+    let isDiEnabled = false
 
     try {
       const {
-        _ddIsEarlyFlakeDetectionEnabled
+        _ddIsEarlyFlakeDetectionEnabled,
+        _ddIsDiEnabled
       } = globalThis.__vitest_worker__.providedContext
 
       isEarlyFlakeDetectionEnabled = _ddIsEarlyFlakeDetectionEnabled
+      isDiEnabled = _ddIsDiEnabled
 
       if (isEarlyFlakeDetectionEnabled) {
         isNew = newTasks.has(task)
@@ -316,12 +331,22 @@ addHook({
 
     // We finish the previous test here because we know it has failed already
     if (numAttempt > 0) {
+      const probe = {}
       const asyncResource = taskToAsync.get(task)
       const testError = task.result?.errors?.[0]
       if (asyncResource) {
         asyncResource.runInAsyncScope(() => {
-          testErrorCh.publish({ error: testError })
+          testErrorCh.publish({
+            error: testError,
+            willBeRetried: true,
+            probe,
+            isDiEnabled
+          })
         })
+        // We wait for the probe to be set
+        if (probe.setProbePromise) {
+          await probe.setProbePromise
+        }
       }
     }