dd-trace 6.13.0 → 6.15.0

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

@@ -11,13 +11,15 @@ const path = require('path')
 const satisfies = require('../../../vendor/dist/semifies')
 const { DD_MAJOR } = require('../../../version')
 const shimmer = require('../../datadog-shimmer')
-const { getEnvironmentVariable } = require('../../dd-trace/src/config/helper')
+const { getEnvironmentVariable, getValueFromEnvSources } = require('../../dd-trace/src/config/helper')
 const log = require('../../dd-trace/src/log')
 const {
   EMPTY_EFD_RETRY_POLICY,
   getEfdRetryCountForDuration,
   hasEfdRetries,
 } = require('../../dd-trace/src/ci-visibility/efd-retry-policy')
+const { FINAL_FLUSH_FALLBACK_DELAY, FINAL_FLUSH_TIMEOUT } =
+  require('../../dd-trace/src/ci-visibility/final-flush')
 const {
   getCoveredFilesFromCoverage,
   JEST_WORKER_TRACE_PAYLOAD_CODE,
@@ -62,6 +64,7 @@ const { addHook, channel } = require('./helpers/instrument')
 const testSessionStartCh = channel('ci:jest:session:start')
 const testSessionFinishCh = channel('ci:jest:session:finish')
 const codeCoverageReportCh = channel('ci:jest:coverage-report')
+const bundlerLoadCh = channel('dd-trace:bundler:load')
 
 const testSessionConfigurationCh = channel('ci:jest:session:configuration')
 
@@ -95,10 +98,10 @@ const itrSkippedSuitesCh = channel('ci:jest:itr:skipped-suites')
 // https://github.com/jestjs/jest/blob/1d682f21c7a35da4d3ab3a1436a357b980ebd0fa/packages/jest-worker/src/types.ts#L37
 const CHILD_MESSAGE_CALL = 1
 
-// Maximum time we'll wait for the tracer to flush
-// The exporter has a 10-second bounded final-flush deadline. Leave enough time
-// for its completion callback before Jest's --forceExit fallback takes over.
-const FLUSH_TIMEOUT = 12_000
+// Maximum time we'll wait for the tracer to flush.
+// Let the exporter enforce its hard deadline and invoke its completion callback
+// before Jest's --forceExit fallback takes over.
+const FLUSH_TIMEOUT = FINAL_FLUSH_TIMEOUT + FINAL_FLUSH_FALLBACK_DELAY + 5000
 const JEST_SESSION_STATE = Symbol.for('dd-trace:jest:session')
 const JEST_BAIL_REPORTER_PATH = require.resolve('./jest/bail-reporter')
 const DD_JEST_HANDLE_TEST_EVENT_WRAPPED = Symbol('dd-trace:jest:handle-test-event-wrapped')
@@ -177,7 +180,9 @@ const wrappedJestEsmLoaders = new WeakSet()
 const wrappedJestObjects = new WeakSet()
 const wrappedWorkerInitializers = new WeakSet()
 const publishedRuntimeReferenceErrors = new WeakMap()
-const jestEsmBypassModulePathsByRuntime = new WeakMap()
+const jestEsmLoggingModulePathsByRuntime = new WeakMap()
+const jestLoggingPackagesByRuntime = new WeakMap()
+const instrumentedJestLoggingModules = new WeakMap()
 const wrappedCoverageReporters = new WeakSet()
 const coverageReporterRequires = new WeakMap()
 const handledJestEvents = new WeakSet()
@@ -342,16 +347,17 @@ function getTestStats (testStatuses) {
 function formatIgnoredFailuresSummary (ignoredFailures) {
   if (!ignoredFailures?.efdFailureCount) return ''
 
-  const items = ignoredFailures.efdNames.map(text => ({ text, suffix: 'Early Flake Detection' }))
-
-  if (items.length === 0) return ''
-
-  const shown = items.slice(0, MAX_IGNORED_TEST_NAMES)
-  const more = items.length - shown.length
+  const shown = ignoredFailures.efdNames.slice(0, MAX_IGNORED_TEST_NAMES)
+  const more = ignoredFailures.efdNames.length - shown.length
   const moreSuffix = more > 0 ? `\n  ... and ${more} more` : ''
-  const formattedItems = shown
-    .map(({ text, suffix }) => `  • ${text}${suffix ? ` (${suffix})` : ''}`)
-    .join('\n') + moreSuffix
+  let formattedItems = ''
+  let isFirstItem = true
+  for (const text of shown) {
+    if (!isFirstItem) formattedItems += '\n'
+    formattedItems += `  • ${text} (Early Flake Detection)`
+    isFirstItem = false
+  }
+  formattedItems += moreSuffix
 
   return `${ignoredFailures.efdFailureCount} test failure(s) were ignored. Exit code set to 0.\n\n${formattedItems}`
 }
@@ -764,7 +770,8 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
             }
             return test.call(this, testName, testFn, ...callArgs)
           }
-          const eachBind = bind(eachTest).apply(this, eachArgs)
+          const boundTest = bind(eachTest)
+          const eachBind = boundTest.apply(this, eachArgs)
           return eachBind.apply(this, testArgs)
         }
       }
@@ -877,7 +884,8 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
             }
             return concurrentTest.call(this, testName, testFn, ...callArgs)
           }
-          const eachBind = bind(concurrentEachTest, false, needsEachError).apply(this, eachArgs)
+          const boundTest = bind(concurrentEachTest, false, needsEachError)
+          const eachBind = boundTest.apply(this, eachArgs)
           return eachBind.apply(this, testArgs)
         }
       }
@@ -926,10 +934,12 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
 
       let jestEach
       try {
-        jestEach = createRequire(path.join(this.rootDir, 'package.json'))('jest-each')
+        const requireFromPath = createRequire(path.join(this.rootDir, 'package.json'))
+        jestEach = requireFromPath('jest-each')
       } catch {
         try {
-          jestEach = createRequire(path.join(process.cwd(), 'package.json'))('jest-each')
+          const requireFromPath = createRequire(path.join(process.cwd(), 'package.json'))
+          jestEach = requireFromPath('jest-each')
         } catch {
           jestEach = undefined
         }
@@ -3067,9 +3077,7 @@ function getCliWrapper (isNewJestVersion) {
         result = await runCLI.apply(this, arguments)
       } catch (error) {
         try {
-          await waitForTestSessionFinish(getTestSessionFinishPayload('fail', error, {
-            isTestSessionFinalizationError: true,
-          }))
+          await waitForTestSessionFinish(getTestSessionFinishPayload('fail', error))
         } catch (finalizationError) {
           log.error('Jest test session finalization error: %s', finalizationError)
         }
@@ -3728,13 +3736,15 @@ if (DD_MAJOR < 6) {
   }, jestConfigSyncWrapper)
 }
 
-const LOGGING_LIBRARIES_BYPASSING_JEST_REQUIRE_ENGINE = new Set([
+const JEST_LOGGING_LIBRARIES = new Set([
   'bunyan',
   'pino',
   'winston',
 ])
+const disabledJestInstrumentations = new Set(
+  getValueFromEnvSources('DD_TRACE_DISABLED_INSTRUMENTATIONS')?.split(',')
+)
 const LIBRARIES_BYPASSING_JEST_REQUIRE_ENGINE = new Set([
-  ...LOGGING_LIBRARIES_BYPASSING_JEST_REQUIRE_ENGINE,
   'selenium-webdriver',
   'selenium-webdriver/chrome',
   'selenium-webdriver/edge',
@@ -3879,13 +3889,17 @@ function requireOutsideJestRequireEngine (runtime, moduleName) {
  * @param {string} moduleName
  * @returns {void}
  */
-function recordJestEsmBypassModulePath (runtime, from, moduleName) {
-  if (typeof from !== 'string' || !LOGGING_LIBRARIES_BYPASSING_JEST_REQUIRE_ENGINE.has(moduleName)) return
+function recordJestEsmLoggingModulePath (runtime, from, moduleName) {
+  if (
+    typeof from !== 'string' ||
+    !JEST_LOGGING_LIBRARIES.has(moduleName) ||
+    disabledJestInstrumentations.has(moduleName)
+  ) return
 
-  let pathsByParent = jestEsmBypassModulePathsByRuntime.get(runtime)
+  let pathsByParent = jestEsmLoggingModulePathsByRuntime.get(runtime)
   if (!pathsByParent) {
     pathsByParent = new Map()
-    jestEsmBypassModulePathsByRuntime.set(runtime, pathsByParent)
+    jestEsmLoggingModulePathsByRuntime.set(runtime, pathsByParent)
   }
 
   let modulePaths = pathsByParent.get(from)
@@ -3906,16 +3920,15 @@ function recordJestEsmBypassModulePath (runtime, from, moduleName) {
  * @param {object} runtime
  * @param {string} from
  * @param {string} modulePath
- * @returns {boolean}
+ * @returns {string | undefined}
  */
-function hasJestEsmBypassModulePath (runtime, from, modulePath) {
-  const modulePaths = jestEsmBypassModulePathsByRuntime.get(runtime)?.get(from)
-  if (!modulePaths) return false
+function getJestEsmLoggingModuleName (runtime, from, modulePath) {
+  const modulePaths = jestEsmLoggingModulePathsByRuntime.get(runtime)?.get(from)
+  if (!modulePaths) return
 
-  for (const resolvedPath of modulePaths.values()) {
-    if (resolvedPath === modulePath) return true
+  for (const [moduleName, resolvedPath] of modulePaths) {
+    if (resolvedPath === modulePath) return moduleName
   }
-  return false
 }
 
 /**
@@ -3929,7 +3942,7 @@ function wrapJestEsmLoader (runtime) {
   wrappedJestEsmLoaders.add(esmLoader)
   if (typeof esmLoader.resolveModule === 'function') {
     shimmer.wrap(esmLoader, 'resolveModule', resolveModule => function (moduleName, from) {
-      recordJestEsmBypassModulePath(runtime, from, moduleName)
+      recordJestEsmLoggingModulePath(runtime, from, moduleName)
       return resolveModule.apply(this, arguments)
     })
   }
@@ -3938,7 +3951,7 @@ function wrapJestEsmLoader (runtime) {
       esmLoader,
       'resolveSpecifierForSyncGraph',
       resolveSpecifier => function (from, moduleName) {
-        recordJestEsmBypassModulePath(runtime, from, moduleName)
+        recordJestEsmLoggingModulePath(runtime, from, moduleName)
         return resolveSpecifier.apply(this, arguments)
       }
     )
@@ -3955,11 +3968,10 @@ function getJestBypassModulePath (runtime, from, moduleName) {
   if (typeof from !== 'string' || typeof moduleName !== 'string') return
 
   if (!LIBRARIES_BYPASSING_JEST_REQUIRE_ENGINE.has(moduleName)) {
-    // Jest passes the resolved path when a CommonJS package is imported from an ESM test.
-    if (path.isAbsolute(moduleName) && hasJestEsmBypassModulePath(runtime, from, moduleName)) {
-      return moduleName
-    }
-    return
+    // Keep the native fallback for logging packages that cannot be instrumented in Jest's realm, such as a
+    // package-name-preserving symlink whose target is a user wrapper rather than the package itself.
+    if (path.isAbsolute(moduleName) && getJestEsmLoggingModuleName(runtime, from, moduleName)) return moduleName
+    if (!JEST_LOGGING_LIBRARIES.has(moduleName) || disabledJestInstrumentations.has(moduleName)) return
   }
 
   try {
@@ -3978,6 +3990,113 @@ function getJestBypassModulePath (runtime, from, moduleName) {
   }
 }
 
+/**
+ * @param {object} runtime
+ * @param {string} from
+ * @param {string} moduleName
+ * @returns {string}
+ */
+function resolveJestModulePath (runtime, from, moduleName) {
+  if (path.isAbsolute(moduleName)) return moduleName
+
+  if (typeof runtime._resolveCjsModule === 'function') {
+    return runtime._resolveCjsModule(from, moduleName)
+  } else if (typeof runtime.cjsLoader?.resolution?.resolveCjs === 'function') {
+    return runtime.cjsLoader.resolution.resolveCjs(from, moduleName)
+  }
+  return runtime._resolveModule(from, moduleName)
+}
+
+/**
+ * @param {object} runtime
+ * @param {string} from
+ * @param {string} moduleName
+ * @returns {{ name: string, path: string, version: string } | undefined}
+ */
+function getJestLoggingPackage (runtime, from, moduleName) {
+  if (typeof from !== 'string' || typeof moduleName !== 'string') return
+
+  const directModuleName = JEST_LOGGING_LIBRARIES.has(moduleName) && !disabledJestInstrumentations.has(moduleName)
+    ? moduleName
+    : getJestEsmLoggingModuleName(runtime, from, moduleName)
+  const packages = jestLoggingPackagesByRuntime.get(runtime)
+  let containingPackage
+  if (!directModuleName) {
+    if (!packages) return
+
+    const normalizedFrom = from.replaceAll(path.sep, '/')
+    for (const [packageRoot, loggingPackage] of packages) {
+      if (normalizedFrom.startsWith(`${packageRoot}/`)) {
+        containingPackage = { packageRoot, ...loggingPackage }
+        break
+      }
+    }
+    if (!containingPackage) return
+  }
+
+  try {
+    const modulePath = resolveJestModulePath(runtime, from, moduleName)
+    const normalizedModulePath = modulePath.replaceAll(path.sep, '/')
+
+    if (directModuleName) {
+      if (modulePath !== createRequire(from).resolve(directModuleName)) return
+
+      const nodeModulesPath = `/node_modules/${directModuleName}/`
+      const packagePathIndex = normalizedModulePath.lastIndexOf(nodeModulesPath)
+      if (packagePathIndex === -1) return
+
+      const packageRoot = normalizedModulePath.slice(0, packagePathIndex + nodeModulesPath.length - 1)
+      const { version } = JSON.parse(readFileSync(`${packageRoot}/package.json`, 'utf8'))
+      if (typeof version !== 'string') return
+
+      let runtimePackages = packages
+      if (!runtimePackages) {
+        runtimePackages = new Map()
+        jestLoggingPackagesByRuntime.set(runtime, runtimePackages)
+      }
+      runtimePackages.set(packageRoot, { name: directModuleName, version })
+      return { name: directModuleName, path: directModuleName, version }
+    }
+
+    const { packageRoot, name, version } = containingPackage
+    if (!normalizedModulePath.startsWith(`${packageRoot}/`)) return
+
+    const modulePathWithinPackage = normalizedModulePath.slice(packageRoot.length + 1)
+    return { name, path: `${name}/${modulePathWithinPackage}`, version }
+  } catch {
+    // Let Jest load unresolved, resolver-only, or virtual modules without instrumentation.
+  }
+}
+
+/**
+ * @param {unknown} moduleExports
+ * @param {{ name: string, path: string, version: string } | undefined} loggingPackage
+ * @returns {unknown}
+ */
+function instrumentJestLoggingModule (moduleExports, loggingPackage) {
+  if (!loggingPackage || (typeof moduleExports !== 'object' && typeof moduleExports !== 'function')) {
+    return moduleExports
+  }
+  if (moduleExports === null) return moduleExports
+
+  const instrumentedModule = instrumentedJestLoggingModules.get(moduleExports)
+  if (instrumentedModule) return instrumentedModule
+
+  const payload = {
+    module: moduleExports,
+    package: loggingPackage.name,
+    path: loggingPackage.path,
+    version: loggingPackage.version,
+  }
+  // Apply the regular instrumentation after Jest has evaluated the module so its built-ins stay in Jest's realm.
+  bundlerLoadCh.publish(payload)
+  instrumentedJestLoggingModules.set(moduleExports, payload.module)
+  if (payload.module && (typeof payload.module === 'object' || typeof payload.module === 'function')) {
+    instrumentedJestLoggingModules.set(payload.module, payload.module)
+  }
+  return payload.module
+}
+
 function formatDefaultStackTrace (error, structuredStackTrace) {
   const errorString = Error.prototype.toString.call(error)
   if (structuredStackTrace.length === 0) return errorString
@@ -4004,7 +4123,7 @@ addHook({
   // Jest 28 through 30.3 keeps ESM dependency resolution on Runtime itself.
   if (typeof Runtime.prototype.resolveModule === 'function') {
     shimmer.wrap(Runtime.prototype, 'resolveModule', resolveModule => function (moduleName, from) {
-      recordJestEsmBypassModulePath(this, from, moduleName)
+      recordJestEsmLoggingModulePath(this, from, moduleName)
       return resolveModule.apply(this, arguments)
     })
   }
@@ -4019,11 +4138,13 @@ addHook({
   shimmer.wrap(Runtime.prototype, 'requireModule', requireModule => function (from, moduleName) {
     wrapJestGlobalsForRuntime(this)
     try {
+      const loggingPackage = getJestLoggingPackage(this, from, moduleName)
       // Jest calls requireModule only after deciding that the module should not be mocked.
-      const bypassModulePath = getJestBypassModulePath(this, from, moduleName)
-      const returnedValue = bypassModulePath
+      const bypassModulePath = loggingPackage ? undefined : getJestBypassModulePath(this, from, moduleName)
+      let returnedValue = bypassModulePath
         ? requireOutsideJestRequireEngine(this, bypassModulePath)
         : requireModule.apply(this, arguments)
+      returnedValue = instrumentJestLoggingModule(returnedValue, loggingPackage)
       if (moduleName === '@jest/globals') {
         wrapConcurrentJestGlobalsForRuntime(this, returnedValue)
       }

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

@@ -3,6 +3,11 @@
 const shimmer = require('../../datadog-shimmer')
 const { addHook, channel } = require('./helpers/instrument')
 const { wrapThen } = require('./helpers/promise')
+const {
+  wrapPoolRelease,
+  wrapPromisePoolAcquire,
+  wrapPromisePoolQueryMethod,
+} = require('./helpers/pool-acquire')
 
 const startRawQueryCh = channel('datadog:knex:raw:start')
 const rawQuerySubscribes = channel('datadog:knex:raw:subscribes')
@@ -32,6 +37,14 @@ addHook({
   versions: ['>=2'],
   file: 'lib/knex-builder/Knex.js',
 }, Knex => {
+  shimmer.wrap(Knex.Client.prototype, 'acquireConnection', acquireConnection => wrapPromisePoolAcquire(
+    acquireConnection,
+    resolveKnexDriver,
+    resolveKnexConfig,
+    knexHasIdleConnection
+  ))
+  shimmer.wrap(Knex.Client.prototype, 'releaseConnection', wrapPoolRelease)
+
   shimmer.wrap(Knex.Client.prototype, 'raw', raw => function (...args) {
     if (!startRawQueryCh.hasSubscribers) {
       return raw.apply(this, args)
@@ -77,6 +90,34 @@ addHook({
   return Knex
 })
 
+addHook({
+  name: 'knex',
+  versions: ['>=2'],
+  file: 'lib/execution/runner.js',
+}, Runner => {
+  shimmer.wrap(Runner.prototype, 'ensureConnection', ensureConnection => wrapPromisePoolQueryMethod(
+    ensureConnection,
+    runner => runner.connection || runner.builder?._connection ? undefined : runner.client,
+    runner => resolveKnexDriver(runner.client)
+  ))
+
+  return Runner
+})
+
+addHook({
+  name: 'knex',
+  versions: ['>=2'],
+  file: 'lib/execution/transaction.js',
+}, Transaction => {
+  shimmer.wrap(Transaction.prototype, 'acquireConnection', acquireConnection => wrapPromisePoolQueryMethod(
+    acquireConnection,
+    (transaction, args) => args[0]?.connection ? undefined : transaction.client,
+    transaction => resolveKnexDriver(transaction.client)
+  ))
+
+  return Transaction
+})
+
 function wrapCallbackWithFinish (callback, finish, context) {
   if (typeof callback !== 'function') return callback
 
@@ -84,3 +125,38 @@ function wrapCallbackWithFinish (callback, finish, context) {
     finish(context, () => callback.apply(this, args))
   })
 }
+
+/**
+ * @param {{ pool?: { numFree?: () => number, numPendingAcquires?: () => number } }} client
+ * @returns {boolean}
+ */
+function knexHasIdleConnection (client) {
+  const pool = client.pool
+  if (typeof pool?.numFree !== 'function') return false
+  return pool.numFree() > pool.numPendingAcquires()
+}
+
+/**
+ * @param {{ driverName?: string }|undefined} client
+ * @returns {'mysql'|'mysql2'|'pg'|undefined}
+ */
+function resolveKnexDriver (client) {
+  switch (client?.driverName) {
+    case 'mysql':
+      return 'mysql'
+    case 'mysql2':
+      return 'mysql2'
+    case 'pg':
+    case 'cockroachdb':
+    case 'pg-redshift':
+      return 'pg'
+  }
+}
+
+/**
+ * @param {{ connectionSettings?: Record<string, unknown> }} client
+ * @returns {Record<string, unknown>}
+ */
+function resolveKnexConfig (client) {
+  return client.connectionSettings ?? {}
+}

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

@@ -11,6 +11,7 @@ const {
   runOutsidePoolQueryAcquire,
   runPoolAcquireError,
   runWithPoolWait,
+  setPromisePoolConnectionConfig,
   takePoolWaitTime,
   wrapPoolQueryMethod,
 } = require('./helpers/pool-acquire')
@@ -97,6 +98,8 @@ function createWrapQuery (options) {
       if (!startCh.hasSubscribers) return query.apply(this, arguments)
 
       const ctx = { sql, conf: options }
+      const poolWaitTime = takePoolWaitTime(this)
+      if (poolWaitTime !== undefined) ctx.poolWaitTime = poolWaitTime
 
       return startCh.runStores(ctx, query, this, ...arguments)
         .then(result => {
@@ -146,8 +149,11 @@ function createWrapQueryCallback (options) {
 function wrapConnection (promiseMethod, Connection) {
   return function (options) {
     Connection.apply(this, arguments)
+    setPromisePoolConnectionConfig(this, options)
 
+    const query = this[promiseMethod]
     shimmer.wrap(this, promiseMethod, createWrapQuery(options))
+    if (this.query === query) this.query = this[promiseMethod]
     shimmer.wrap(this, '_queryCallback', createWrapQueryCallback(options))
   }
 }

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

@@ -11,6 +11,60 @@ const MINIMUM_MOCHA_VERSION = DD_MAJOR >= 6 ? '>=8.0.0' : '>=5.2.0'
 const parameterizedTestCh = channel('ci:mocha:test:parameterize')
 const patched = new WeakSet()
 
+/**
+ * Registers every Mocha package view that can expose the constructor.
+ *
+ * @param {string[]} versions
+ * @param {(Mocha: Function, frameworkVersion: string) => Function} wrapMochaRun
+ * @returns {void}
+ */
+function addMochaRunHooks (versions, wrapMochaRun) {
+  const patchedMochaConstructors = new WeakSet()
+
+  /**
+   * @param {Function} Mocha
+   * @param {string} frameworkVersion
+   * @returns {Function}
+   */
+  function wrapMochaOnce (Mocha, frameworkVersion) {
+    if (patchedMochaConstructors.has(Mocha)) return Mocha
+    patchedMochaConstructors.add(Mocha)
+    return wrapMochaRun(Mocha, frameworkVersion)
+  }
+
+  /**
+   * @param {object|Function} MochaPackage
+   * @param {string} frameworkVersion
+   * @returns {object|Function}
+   */
+  function wrapMochaPackage (MochaPackage, frameworkVersion) {
+    wrapMochaOnce(MochaPackage.default ?? MochaPackage, frameworkVersion)
+    return MochaPackage
+  }
+
+  addHook({
+    name: 'mocha',
+    versions,
+    filePattern: String.raw`lib/mocha\.(?:c?js)$`,
+  }, wrapMochaOnce)
+
+  // Mocha 12's ESM package root can load lib/mocha.cjs before the CJS hook observes it. ESM imports report the
+  // package root while synchronous require reports index.js, so hook both views. Keeping the namespace intact makes
+  // the ESM hook skip its separate default-export callback, and wrapMochaOnce deduplicates loaders exposing both.
+  addHook({
+    name: 'mocha',
+    versions: ['>=12.0.0'],
+    patchDefault: false,
+  }, wrapMochaPackage)
+
+  addHook({
+    name: 'mocha',
+    versions: ['>=12.0.0'],
+    file: 'index.js',
+    patchDefault: false,
+  }, wrapMochaPackage)
+}
+
 // mocha-each support
 addHook({
   name: 'mocha-each',
@@ -38,7 +92,8 @@ addHook({
   name: 'mocha',
   versions: [MINIMUM_MOCHA_VERSION],
   file: 'lib/suite.js',
-}, (Suite) => {
+}, (SuitePackage) => {
+  const Suite = SuitePackage.Suite ?? SuitePackage
   shimmer.wrap(Suite.prototype, 'addTest', addTest => function (test) {
     const callSites = getCallSites()
     const testCallSite = callSites.find(site => site.getFileName() === test.file)
@@ -47,5 +102,7 @@ addHook({
     }
     return addTest.apply(this, arguments)
   })
-  return Suite
+  return SuitePackage
 })
+
+module.exports = { addMochaRunHooks }