dd-trace 4.45.0 → 4.47.0

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

@@ -1,7 +1,6 @@
 'use strict'
 const { createCoverageMap } = require('istanbul-lib-coverage')
 
-const { NUM_FAILED_TEST_RETRIES } = require('../../dd-trace/src/plugins/util/test')
 const { addHook, channel, AsyncResource } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 const log = require('../../dd-trace/src/log')
@@ -36,7 +35,8 @@ const {
   mergeCoverage,
   fromCoverageMapToCoverage,
   getTestSuitePath,
-  CUCUMBER_WORKER_TRACE_PAYLOAD_CODE
+  CUCUMBER_WORKER_TRACE_PAYLOAD_CODE,
+  getIsFaultyEarlyFlakeDetection
 } = require('../../dd-trace/src/plugins/util/test')
 
 const isMarkedAsUnskippable = (pickle) => {
@@ -52,7 +52,9 @@ const patched = new WeakSet()
 const lastStatusByPickleId = new Map()
 const numRetriesByPickleId = new Map()
 const numAttemptToAsyncResource = new Map()
+const newTestsByTestFullname = new Map()
 
+let eventDataCollector = null
 let pickleByFile = {}
 const pickleResultByFile = {}
 
@@ -65,7 +67,10 @@ let isUnskippable = false
 let isSuitesSkippingEnabled = false
 let isEarlyFlakeDetectionEnabled = false
 let earlyFlakeDetectionNumRetries = 0
+let earlyFlakeDetectionFaultyThreshold = 0
+let isEarlyFlakeDetectionFaulty = false
 let isFlakyTestRetriesEnabled = false
+let numTestRetries = 0
 let knownTests = []
 let skippedSuites = []
 let isSuitesSkipped = false
@@ -129,15 +134,35 @@ function getChannelPromise (channelToPublishTo) {
   })
 }
 
+function getShouldBeSkippedSuite (pickle, suitesToSkip) {
+  const testSuitePath = getTestSuitePath(pickle.uri, process.cwd())
+  const isUnskippable = isMarkedAsUnskippable(pickle)
+  const isSkipped = suitesToSkip.includes(testSuitePath)
+
+  return [isSkipped && !isUnskippable, testSuitePath]
+}
+
+// From cucumber@>=11
+function getFilteredPicklesNew (coordinator, suitesToSkip) {
+  return coordinator.sourcedPickles.reduce((acc, sourcedPickle) => {
+    const { pickle } = sourcedPickle
+    const [shouldBeSkipped, testSuitePath] = getShouldBeSkippedSuite(pickle, suitesToSkip)
+
+    if (shouldBeSkipped) {
+      acc.skippedSuites.add(testSuitePath)
+    } else {
+      acc.picklesToRun.push(sourcedPickle)
+    }
+    return acc
+  }, { skippedSuites: new Set(), picklesToRun: [] })
+}
+
 function getFilteredPickles (runtime, suitesToSkip) {
   return runtime.pickleIds.reduce((acc, pickleId) => {
-    const test = runtime.eventDataCollector.getPickle(pickleId)
-    const testSuitePath = getTestSuitePath(test.uri, process.cwd())
-
-    const isUnskippable = isMarkedAsUnskippable(test)
-    const isSkipped = suitesToSkip.includes(testSuitePath)
+    const pickle = runtime.eventDataCollector.getPickle(pickleId)
+    const [shouldBeSkipped, testSuitePath] = getShouldBeSkippedSuite(pickle, suitesToSkip)
 
-    if (isSkipped && !isUnskippable) {
+    if (shouldBeSkipped) {
       acc.skippedSuites.add(testSuitePath)
     } else {
       acc.picklesToRun.push(pickleId)
@@ -146,9 +171,21 @@ function getFilteredPickles (runtime, suitesToSkip) {
   }, { skippedSuites: new Set(), picklesToRun: [] })
 }
 
-function getPickleByFile (runtime) {
-  return runtime.pickleIds.reduce((acc, pickleId) => {
-    const test = runtime.eventDataCollector.getPickle(pickleId)
+// From cucumber@>=11
+function getPickleByFileNew (coordinator) {
+  return coordinator.sourcedPickles.reduce((acc, { pickle }) => {
+    if (acc[pickle.uri]) {
+      acc[pickle.uri].push(pickle)
+    } else {
+      acc[pickle.uri] = [pickle]
+    }
+    return acc
+  }, {})
+}
+
+function getPickleByFile (runtimeOrCoodinator) {
+  return runtimeOrCoodinator.pickleIds.reduce((acc, pickleId) => {
+    const test = runtimeOrCoodinator.eventDataCollector.getPickle(pickleId)
     if (acc[test.uri]) {
       acc[test.uri].push(test)
     } else {
@@ -188,7 +225,7 @@ function wrapRun (pl, isLatestVersion) {
       testStartCh.publish(testStartPayload)
     })
     try {
-      this.eventBroadcaster.on('envelope', (testCase) => {
+      this.eventBroadcaster.on('envelope', shimmer.wrapFunction(null, () => (testCase) => {
         // Only supported from >=8.0.0
         if (testCase?.testCaseFinished) {
           const { testCaseFinished: { willBeRetried } } = testCase
@@ -206,7 +243,7 @@ function wrapRun (pl, isLatestVersion) {
             })
           }
         }
-      })
+      }))
       let promise
 
       asyncResource.runInAsyncScope(() => {
@@ -294,19 +331,34 @@ function testCaseHook (TestCaseRunner) {
   return TestCaseRunner
 }
 
-function getWrappedStart (start, frameworkVersion, isParallel = false) {
+// Valid for old and new cucumber versions
+function getCucumberOptions (adapterOrCoordinator) {
+  if (adapterOrCoordinator.adapter) {
+    return adapterOrCoordinator.adapter.worker?.options || adapterOrCoordinator.adapter.options
+  }
+  return adapterOrCoordinator.options
+}
+
+function getWrappedStart (start, frameworkVersion, isParallel = false, isCoordinator = false) {
   return async function () {
     if (!libraryConfigurationCh.hasSubscribers) {
       return start.apply(this, arguments)
     }
+    const options = getCucumberOptions(this)
+
+    if (!isParallel && this.adapter?.options) {
+      isParallel = options.parallel > 0
+    }
     let errorSkippableRequest
 
     const configurationResponse = await getChannelPromise(libraryConfigurationCh)
 
     isEarlyFlakeDetectionEnabled = configurationResponse.libraryConfig?.isEarlyFlakeDetectionEnabled
     earlyFlakeDetectionNumRetries = configurationResponse.libraryConfig?.earlyFlakeDetectionNumRetries
+    earlyFlakeDetectionFaultyThreshold = configurationResponse.libraryConfig?.earlyFlakeDetectionFaultyThreshold
     isSuitesSkippingEnabled = configurationResponse.libraryConfig?.isSuitesSkippingEnabled
     isFlakyTestRetriesEnabled = configurationResponse.libraryConfig?.isFlakyTestRetriesEnabled
+    numTestRetries = configurationResponse.libraryConfig?.flakyTestRetriesCount
 
     if (isEarlyFlakeDetectionEnabled) {
       const knownTestsResponse = await getChannelPromise(knownTestsCh)
@@ -324,28 +376,49 @@ function getWrappedStart (start, frameworkVersion, isParallel = false) {
       skippableSuites = skippableResponse.skippableSuites
 
       if (!errorSkippableRequest) {
-        const filteredPickles = getFilteredPickles(this, skippableSuites)
+        const filteredPickles = isCoordinator
+          ? getFilteredPicklesNew(this, skippableSuites)
+          : getFilteredPickles(this, skippableSuites)
+
         const { picklesToRun } = filteredPickles
-        isSuitesSkipped = picklesToRun.length !== this.pickleIds.length
+        const oldPickles = isCoordinator ? this.sourcedPickles : this.pickleIds
+
+        isSuitesSkipped = picklesToRun.length !== oldPickles.length
 
         log.debug(
-          () => `${picklesToRun.length} out of ${this.pickleIds.length} suites are going to run.`
+          () => `${picklesToRun.length} out of ${oldPickles.length} suites are going to run.`
         )
 
-        this.pickleIds = picklesToRun
+        if (isCoordinator) {
+          this.sourcedPickles = picklesToRun
+        } else {
+          this.pickleIds = picklesToRun
+        }
 
         skippedSuites = Array.from(filteredPickles.skippedSuites)
         itrCorrelationId = skippableResponse.itrCorrelationId
       }
     }
 
-    pickleByFile = getPickleByFile(this)
+    pickleByFile = isCoordinator ? getPickleByFileNew(this) : getPickleByFile(this)
+
+    if (isEarlyFlakeDetectionEnabled) {
+      const isFaulty = getIsFaultyEarlyFlakeDetection(
+        Object.keys(pickleByFile),
+        knownTests.cucumber || {},
+        earlyFlakeDetectionFaultyThreshold
+      )
+      if (isFaulty) {
+        isEarlyFlakeDetectionEnabled = false
+        isEarlyFlakeDetectionFaulty = true
+      }
+    }
 
     const processArgv = process.argv.slice(2).join(' ')
     const command = process.env.npm_lifecycle_script || `cucumber-js ${processArgv}`
 
-    if (isFlakyTestRetriesEnabled && !this.options.retry) {
-      this.options.retry = NUM_FAILED_TEST_RETRIES
+    if (isFlakyTestRetriesEnabled && !options.retry && numTestRetries > 0) {
+      options.retry = numTestRetries
     }
 
     sessionAsyncResource.runInAsyncScope(() => {
@@ -387,48 +460,65 @@ function getWrappedStart (start, frameworkVersion, isParallel = false) {
         hasUnskippableSuites: isUnskippable,
         hasForcedToRunSuites: isForcedToRun,
         isEarlyFlakeDetectionEnabled,
+        isEarlyFlakeDetectionFaulty,
         isParallel
       })
     })
+    eventDataCollector = null
     return success
   }
 }
 
-function getWrappedRunTest (runTestFunction) {
-  return async function (pickleId) {
-    const test = this.eventDataCollector.getPickle(pickleId)
+// Generates suite start and finish events in the main process.
+// Handles EFD in both the main process and the worker process.
+function getWrappedRunTestCase (runTestCaseFunction, isNewerCucumberVersion = false, isWorker = false) {
+  return async function () {
+    let pickle
+    if (isNewerCucumberVersion) {
+      pickle = arguments[0].pickle
+    } else {
+      pickle = this.eventDataCollector.getPickle(arguments[0])
+    }
 
-    const testFileAbsolutePath = test.uri
+    const testFileAbsolutePath = pickle.uri
     const testSuitePath = getTestSuitePath(testFileAbsolutePath, process.cwd())
 
-    if (!pickleResultByFile[testFileAbsolutePath]) { // first test in suite
-      isUnskippable = isMarkedAsUnskippable(test)
+    // If it's a worker, suite events are handled in `getWrappedParseWorkerMessage`
+    if (!isWorker && !pickleResultByFile[testFileAbsolutePath]) { // first test in suite
+      isUnskippable = isMarkedAsUnskippable(pickle)
       isForcedToRun = isUnskippable && skippableSuites.includes(testSuitePath)
 
-      testSuiteStartCh.publish({ testSuitePath, isUnskippable, isForcedToRun, itrCorrelationId })
+      testSuiteStartCh.publish({
+        testFileAbsolutePath,
+        isUnskippable,
+        isForcedToRun,
+        itrCorrelationId
+      })
     }
 
     let isNew = false
 
     if (isEarlyFlakeDetectionEnabled) {
-      isNew = isNewTest(testSuitePath, test.name)
+      isNew = isNewTest(testSuitePath, pickle.name)
       if (isNew) {
-        numRetriesByPickleId.set(pickleId, 0)
+        numRetriesByPickleId.set(pickle.id, 0)
       }
     }
-    const runTestCaseResult = await runTestFunction.apply(this, arguments)
+    // TODO: for >=11 we could use `runTestCaseResult` instead of accumulating results in `lastStatusByPickleId`
+    let runTestCaseResult = await runTestCaseFunction.apply(this, arguments)
 
-    const testStatuses = lastStatusByPickleId.get(pickleId)
+    const testStatuses = lastStatusByPickleId.get(pickle.id)
     const lastTestStatus = testStatuses[testStatuses.length - 1]
     // If it's a new test and it hasn't been skipped, we run it again
     if (isEarlyFlakeDetectionEnabled && lastTestStatus !== 'skip' && isNew) {
       for (let retryIndex = 0; retryIndex < earlyFlakeDetectionNumRetries; retryIndex++) {
-        numRetriesByPickleId.set(pickleId, retryIndex + 1)
-        await runTestFunction.apply(this, arguments)
+        numRetriesByPickleId.set(pickle.id, retryIndex + 1)
+        runTestCaseResult = await runTestCaseFunction.apply(this, arguments)
       }
     }
     let testStatus = lastTestStatus
-    if (isEarlyFlakeDetectionEnabled) {
+    let shouldBePassedByEFD = false
+    if (isNew && isEarlyFlakeDetectionEnabled) {
       /**
        * If Early Flake Detection (EFD) is enabled the logic is as follows:
        * - If all attempts for a test are failing, the test has failed and we will let the test process fail.
@@ -438,6 +528,8 @@ function getWrappedRunTest (runTestFunction) {
        */
       testStatus = getTestStatusFromRetries(testStatuses)
       if (testStatus === 'pass') {
+        // for cucumber@>=11, setting `this.success` does not work, so we have to change the returned value
+        shouldBePassedByEFD = true
         this.success = true
       }
     }
@@ -448,8 +540,9 @@ function getWrappedRunTest (runTestFunction) {
       pickleResultByFile[testFileAbsolutePath].push(testStatus)
     }
 
-    // last test in suite
-    if (pickleResultByFile[testFileAbsolutePath].length === pickleByFile[testFileAbsolutePath].length) {
+    // If it's a worker, suite events are handled in `getWrappedParseWorkerMessage`
+    if (!isWorker && pickleResultByFile[testFileAbsolutePath].length === pickleByFile[testFileAbsolutePath].length) {
+      // last test in suite
       const testSuiteStatus = getSuiteStatusFromTestStatuses(pickleResultByFile[testFileAbsolutePath])
       if (global.__coverage__) {
         const coverageFiles = getCoveredFilenamesFromCoverage(global.__coverage__)
@@ -468,11 +561,15 @@ function getWrappedRunTest (runTestFunction) {
       testSuiteFinishCh.publish({ status: testSuiteStatus, testSuitePath })
     }
 
+    if (isNewerCucumberVersion && isEarlyFlakeDetectionEnabled && isNew) {
+      return shouldBePassedByEFD
+    }
+
     return runTestCaseResult
   }
 }
 
-function getWrappedParseWorkerMessage (parseWorkerMessageFunction) {
+function getWrappedParseWorkerMessage (parseWorkerMessageFunction, isNewVersion) {
   return function (worker, message) {
     // If the message is an array, it's a dd-trace message, so we need to stop cucumber processing,
     // or cucumber will throw an error
@@ -487,29 +584,43 @@ function getWrappedParseWorkerMessage (parseWorkerMessageFunction) {
       }
     }
 
-    const { jsonEnvelope } = message
-    if (!jsonEnvelope) {
+    let envelope
+
+    if (isNewVersion) {
+      envelope = message.envelope
+    } else {
+      envelope = message.jsonEnvelope
+    }
+
+    if (!envelope) {
       return parseWorkerMessageFunction.apply(this, arguments)
     }
-    let parsed = jsonEnvelope
+    let parsed = envelope
 
     if (typeof parsed === 'string') {
       try {
-        parsed = JSON.parse(jsonEnvelope)
+        parsed = JSON.parse(envelope)
       } catch (e) {
         // ignore errors and continue
         return parseWorkerMessageFunction.apply(this, arguments)
       }
     }
+    let pickle
+
     if (parsed.testCaseStarted) {
-      const { pickleId } = this.eventDataCollector.testCaseMap[parsed.testCaseStarted.testCaseId]
-      const pickle = this.eventDataCollector.getPickle(pickleId)
+      if (isNewVersion) {
+        pickle = this.inProgress[worker.id].pickle
+      } else {
+        const { pickleId } = this.eventDataCollector.testCaseMap[parsed.testCaseStarted.testCaseId]
+        pickle = this.eventDataCollector.getPickle(pickleId)
+      }
+      // THIS FAILS IN PARALLEL MODE
       const testFileAbsolutePath = pickle.uri
       // First test in suite
       if (!pickleResultByFile[testFileAbsolutePath]) {
         pickleResultByFile[testFileAbsolutePath] = []
         testSuiteStartCh.publish({
-          testSuitePath: getTestSuitePath(testFileAbsolutePath, process.cwd())
+          testFileAbsolutePath
         })
       }
     }
@@ -518,14 +629,47 @@ function getWrappedParseWorkerMessage (parseWorkerMessageFunction) {
 
     // after calling `parseWorkerMessageFunction`, the test status can already be read
     if (parsed.testCaseFinished) {
-      const { pickle, worstTestStepResult } =
-        this.eventDataCollector.getTestCaseAttempt(parsed.testCaseFinished.testCaseStartedId)
+      let worstTestStepResult
+      if (isNewVersion && eventDataCollector) {
+        pickle = this.inProgress[worker.id].pickle
+        worstTestStepResult =
+          eventDataCollector.getTestCaseAttempt(parsed.testCaseFinished.testCaseStartedId).worstTestStepResult
+      } else {
+        const testCase = this.eventDataCollector.getTestCaseAttempt(parsed.testCaseFinished.testCaseStartedId)
+        worstTestStepResult = testCase.worstTestStepResult
+        pickle = testCase.pickle
+      }
 
       const { status } = getStatusFromResultLatest(worstTestStepResult)
+      let isNew = false
+
+      if (isEarlyFlakeDetectionEnabled) {
+        isNew = isNewTest(pickle.uri, pickle.name)
+      }
 
       const testFileAbsolutePath = pickle.uri
       const finished = pickleResultByFile[testFileAbsolutePath]
-      finished.push(status)
+
+      if (isNew) {
+        const testFullname = `${pickle.uri}:${pickle.name}`
+        let testStatuses = newTestsByTestFullname.get(testFullname)
+        if (!testStatuses) {
+          testStatuses = [status]
+          newTestsByTestFullname.set(testFullname, testStatuses)
+        } else {
+          testStatuses.push(status)
+        }
+        // We have finished all retries
+        if (testStatuses.length === earlyFlakeDetectionNumRetries + 1) {
+          const newTestFinalStatus = getTestStatusFromRetries(testStatuses)
+          // we only push to `finished` if the retries have finished
+          finished.push(newTestFinalStatus)
+        }
+      } else {
+        // TODO: can we get error message?
+        const finished = pickleResultByFile[testFileAbsolutePath]
+        finished.push(status)
+      }
 
       if (finished.length === pickleByFile[testFileAbsolutePath].length) {
         testSuiteFinishCh.publish({
@@ -555,13 +699,16 @@ addHook({
 
 // From 7.3.0 onwards, runPickle becomes runTestCase. Not executed in parallel mode.
 // `getWrappedStart` generates session start and finish events
-// `getWrappedRunTest` generates suite start and finish events
+// `getWrappedRunTestCase` generates suite start and finish events and handles EFD.
+// TODO (fix): there is a lib/runtime/index in >=11.0.0, but we don't instrument it because it's not useful for us
+// This causes a info log saying "Found incompatible integration version".
 addHook({
   name: '@cucumber/cucumber',
-  versions: ['>=7.3.0'],
+  versions: ['>=7.3.0 <11.0.0'],
   file: 'lib/runtime/index.js'
 }, (runtimePackage, frameworkVersion) => {
-  shimmer.wrap(runtimePackage.default.prototype, 'runTestCase', runTestCase => getWrappedRunTest(runTestCase))
+  shimmer.wrap(runtimePackage.default.prototype, 'runTestCase', runTestCase => getWrappedRunTestCase(runTestCase))
+
   shimmer.wrap(runtimePackage.default.prototype, 'start', start => getWrappedStart(start, frameworkVersion))
 
   return runtimePackage
@@ -569,13 +716,13 @@ addHook({
 
 // Not executed in parallel mode.
 // `getWrappedStart` generates session start and finish events
-// `getWrappedRunTest` generates suite start and finish events
+// `getWrappedRunTestCase` generates suite start and finish events and handles EFD.
 addHook({
   name: '@cucumber/cucumber',
   versions: ['>=7.0.0 <7.3.0'],
   file: 'lib/runtime/index.js'
 }, (runtimePackage, frameworkVersion) => {
-  shimmer.wrap(runtimePackage.default.prototype, 'runPickle', runPickle => getWrappedRunTest(runPickle))
+  shimmer.wrap(runtimePackage.default.prototype, 'runPickle', runPickle => getWrappedRunTestCase(runPickle))
   shimmer.wrap(runtimePackage.default.prototype, 'start', start => getWrappedStart(start, frameworkVersion))
 
   return runtimePackage
@@ -583,11 +730,10 @@ addHook({
 
 // Only executed in parallel mode.
 // `getWrappedStart` generates session start and finish events
-// `getWrappedGiveWork` generates suite start events and sets pickleResultByFile (used by suite finish events)
-// `getWrappedParseWorkerMessage` generates suite finish events
+// `getWrappedParseWorkerMessage` generates suite start and finish events
 addHook({
   name: '@cucumber/cucumber',
-  versions: ['>=8.0.0'],
+  versions: ['>=8.0.0 <11.0.0'],
   file: 'lib/runtime/parallel/coordinator.js'
 }, (coordinatorPackage, frameworkVersion) => {
   shimmer.wrap(coordinatorPackage.default.prototype, 'start', start => getWrappedStart(start, frameworkVersion, true))
@@ -598,3 +744,95 @@ addHook({
   )
   return coordinatorPackage
 })
+
+// >=11.0.0 hooks
+// `getWrappedRunTestCase` does two things:
+// - generates suite start and finish events in the main process,
+// - handles EFD in both the main process and the worker process.
+addHook({
+  name: '@cucumber/cucumber',
+  versions: ['>=11.0.0'],
+  file: 'lib/runtime/worker.js'
+}, (workerPackage) => {
+  shimmer.wrap(
+    workerPackage.Worker.prototype,
+    'runTestCase',
+    runTestCase => getWrappedRunTestCase(runTestCase, true, !!process.env.CUCUMBER_WORKER_ID)
+  )
+  return workerPackage
+})
+
+// `getWrappedStart` generates session start and finish events
+addHook({
+  name: '@cucumber/cucumber',
+  versions: ['>=11.0.0'],
+  file: 'lib/runtime/coordinator.js'
+}, (coordinatorPackage, frameworkVersion) => {
+  shimmer.wrap(
+    coordinatorPackage.Coordinator.prototype,
+    'run',
+    run => getWrappedStart(run, frameworkVersion, false, true)
+  )
+  return coordinatorPackage
+})
+
+// Necessary because `eventDataCollector` is no longer available in the runtime instance
+addHook({
+  name: '@cucumber/cucumber',
+  versions: ['>=11.0.0'],
+  file: 'lib/formatter/helpers/event_data_collector.js'
+}, (eventDataCollectorPackage) => {
+  shimmer.wrap(eventDataCollectorPackage.default.prototype, 'parseEnvelope', parseEnvelope => function () {
+    eventDataCollector = this
+    return parseEnvelope.apply(this, arguments)
+  })
+  return eventDataCollectorPackage
+})
+
+// Only executed in parallel mode for >=11, in the main process.
+// `getWrappedParseWorkerMessage` generates suite start and finish events
+// In `startWorker` we pass early flake detection info to the worker.
+addHook({
+  name: '@cucumber/cucumber',
+  versions: ['>=11.0.0'],
+  file: 'lib/runtime/parallel/adapter.js'
+}, (adapterPackage) => {
+  shimmer.wrap(
+    adapterPackage.ChildProcessAdapter.prototype,
+    'parseWorkerMessage',
+    parseWorkerMessage => getWrappedParseWorkerMessage(parseWorkerMessage, true)
+  )
+  // EFD in parallel mode only supported in >=11.0.0
+  shimmer.wrap(adapterPackage.ChildProcessAdapter.prototype, 'startWorker', startWorker => function () {
+    if (isEarlyFlakeDetectionEnabled) {
+      this.options.worldParameters._ddKnownTests = knownTests
+      this.options.worldParameters._ddEarlyFlakeDetectionNumRetries = earlyFlakeDetectionNumRetries
+    }
+
+    return startWorker.apply(this, arguments)
+  })
+  return adapterPackage
+})
+
+// Hook executed in the worker process when in parallel mode.
+// In this hook we read the information passed in `worldParameters` and make it available for
+// `getWrappedRunTestCase`.
+addHook({
+  name: '@cucumber/cucumber',
+  versions: ['>=11.0.0'],
+  file: 'lib/runtime/parallel/worker.js'
+}, (workerPackage) => {
+  shimmer.wrap(
+    workerPackage.ChildProcessWorker.prototype,
+    'initialize',
+    initialize => async function () {
+      await initialize.apply(this, arguments)
+      isEarlyFlakeDetectionEnabled = !!this.options.worldParameters._ddKnownTests
+      if (isEarlyFlakeDetectionEnabled) {
+        knownTests = this.options.worldParameters._ddKnownTests
+        earlyFlakeDetectionNumRetries = this.options.worldParameters._ddEarlyFlakeDetectionNumRetries
+      }
+    }
+  )
+  return workerPackage
+})

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

@@ -21,16 +21,16 @@ const rrtypeMap = new WeakMap()
 const names = ['dns', 'node:dns']
 
 addHook({ name: names }, dns => {
-  dns.lookup = wrap('apm:dns:lookup', dns.lookup, 2)
-  dns.lookupService = wrap('apm:dns:lookup_service', dns.lookupService, 3)
-  dns.resolve = wrap('apm:dns:resolve', dns.resolve, 2)
-  dns.reverse = wrap('apm:dns:reverse', dns.reverse, 2)
+  shimmer.wrap(dns, 'lookup', fn => wrap('apm:dns:lookup', fn, 2))
+  shimmer.wrap(dns, 'lookupService', fn => wrap('apm:dns:lookup_service', fn, 2))
+  shimmer.wrap(dns, 'resolve', fn => wrap('apm:dns:resolve', fn, 2))
+  shimmer.wrap(dns, 'reverse', fn => wrap('apm:dns:reverse', fn, 2))
 
   patchResolveShorthands(dns)
 
   if (dns.Resolver) {
-    dns.Resolver.prototype.resolve = wrap('apm:dns:resolve', dns.Resolver.prototype.resolve, 2)
-    dns.Resolver.prototype.reverse = wrap('apm:dns:reverse', dns.Resolver.prototype.reverse, 2)
+    shimmer.wrap(dns.Resolver.prototype, 'resolve', fn => wrap('apm:dns:resolve', fn, 2))
+    shimmer.wrap(dns.Resolver.prototype, 'reverse', fn => wrap('apm:dns:reverse', fn, 2))
 
     patchResolveShorthands(dns.Resolver.prototype)
   }
@@ -43,7 +43,7 @@ function patchResolveShorthands (prototype) {
     .filter(method => !!prototype[method])
     .forEach(method => {
       rrtypeMap.set(prototype[method], rrtypes[method])
-      prototype[method] = wrap('apm:dns:resolve', prototype[method], 2, rrtypes[method])
+      shimmer.wrap(prototype, method, fn => wrap('apm:dns:resolve', fn, 2, rrtypes[method]))
     })
 }
 
@@ -72,13 +72,13 @@ function wrap (prefix, fn, expectedArgs, rrtype) {
     return asyncResource.runInAsyncScope(() => {
       startCh.publish(startArgs)
 
-      arguments[arguments.length - 1] = asyncResource.bind(function (error, result) {
+      arguments[arguments.length - 1] = shimmer.wrapFunction(cb, cb => asyncResource.bind(function (error, result) {
         if (error) {
           errorCh.publish(error)
         }
         finishCh.publish(result)
         cb.apply(this, arguments)
-      })
+      }))
 
       try {
         return fn.apply(this, arguments)
@@ -92,5 +92,5 @@ function wrap (prefix, fn, expectedArgs, rrtype) {
     })
   }
 
-  return shimmer.wrap(fn, wrapped)
+  return wrapped
 }

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

@@ -48,12 +48,12 @@ function createWrapSelect () {
     return function () {
       if (arguments.length === 1) {
         const cb = arguments[0]
-        arguments[0] = function (err, connection) {
+        arguments[0] = shimmer.wrapFunction(cb, cb => function (err, connection) {
           if (connectCh.hasSubscribers && connection && connection.host) {
             connectCh.publish({ hostname: connection.host.host, port: connection.host.port })
           }
           cb(err, connection)
-        }
+        })
       }
       return request.apply(this, arguments)
     }
@@ -86,10 +86,10 @@ function createWrapRequest (name) {
           if (typeof cb === 'function') {
             cb = parentResource.bind(cb)
 
-            arguments[lastIndex] = asyncResource.bind(function (error) {
+            arguments[lastIndex] = shimmer.wrapFunction(cb, cb => asyncResource.bind(function (error) {
               finish(params, error)
               return cb.apply(null, arguments)
-            })
+            }))
             return request.apply(this, arguments)
           } else {
             const promise = request.apply(this, arguments)

package/packages/datadog-instrumentations/src/express-mongo-sanitize.js

@@ -22,15 +22,15 @@ addHook({ name: 'express-mongo-sanitize', versions: ['>=1.0.0'] }, expressMongoS
     return sanitizedObject
   })
 
-  return shimmer.wrap(expressMongoSanitize, function () {
+  return shimmer.wrapFunction(expressMongoSanitize, expressMongoSanitize => function () {
     const middleware = expressMongoSanitize.apply(this, arguments)
 
-    return shimmer.wrap(middleware, function (req, res, next) {
+    return shimmer.wrapFunction(middleware, middleware => function (req, res, next) {
       if (!sanitizeMiddlewareFinished.hasSubscribers) {
         return middleware.apply(this, arguments)
       }
 
-      const wrappedNext = shimmer.wrap(next, function () {
+      const wrappedNext = shimmer.wrapFunction(next, next => function () {
         sanitizeMiddlewareFinished.publish({
           sanitizedProperties: propertiesToSanitize,
           req