[ci-visibility] Remove usage of .asyncResource in mocha plugin (#4348)

DataDog · May 27, 2024 · 8bedb23 · 8bedb23
1 parent e77f059
commit 8bedb23
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 12 deletions.
diff --git a/packages/datadog-instrumentations/src/mocha/main.js b/packages/datadog-instrumentations/src/mocha/main.js
@@ -23,7 +23,8 @@ const {
   getOnTestEndHandler,
   getOnHookEndHandler,
   getOnFailHandler,
-  getOnPendingHandler
+  getOnPendingHandler,
+  testFileToSuiteAr
 } = require('./utils')
 require('./common')
 
@@ -39,7 +40,6 @@ let isSuitesSkippingEnabled = false
 let earlyFlakeDetectionNumRetries = 0
 let knownTests = []
 let itrCorrelationId = ''
-const testFileToSuiteAr = new Map()
 let isForcedToRun = false
 
 // We'll preserve the original coverage here
@@ -469,7 +469,7 @@ addHook({
 // Used to start and finish test session and test module
 addHook({
   name: 'mocha',
-  versions: ['>=8.0.0'],
+  versions: ['>=5.2.0'],
   file: 'lib/nodejs/parallel-buffered-runner.js'
 }, (ParallelBufferedRunner, frameworkVersion) => {
   shimmer.wrap(ParallelBufferedRunner.prototype, 'run', run => function () {

diff --git a/packages/datadog-instrumentations/src/mocha/utils.js b/packages/datadog-instrumentations/src/mocha/utils.js
@@ -21,6 +21,7 @@ const testToAr = new WeakMap()
 const originalFns = new WeakMap()
 const testToStartLine = new WeakMap()
 const testFileToSuiteAr = new Map()
+const wrappedFunctions = new WeakSet()
 
 function isNewTest (test, knownTests) {
   const testSuite = getTestSuitePath(test.file, process.cwd())
@@ -87,7 +88,7 @@ function getTestAsyncResource (test) {
   if (!test.fn) {
     return testToAr.get(test)
   }
-  if (!test.fn.asyncResource) {
+  if (!wrappedFunctions.has(test.fn)) {
     return testToAr.get(test.fn)
   }
   const originalFn = originalFns.get(test.fn)
@@ -105,9 +106,10 @@ function runnableWrapper (RunnablePackage) {
     const isTestHook = isBeforeEach || isAfterEach
 
     // we restore the original user defined function
-    if (this.fn.asyncResource) {
+    if (wrappedFunctions.has(this.fn)) {
       const originalFn = originalFns.get(this.fn)
       this.fn = originalFn
+      wrappedFunctions.delete(this.fn)
     }
 
     if (isTestHook || this.type === 'test') {
@@ -122,11 +124,7 @@ function runnableWrapper (RunnablePackage) {
         originalFns.set(newFn, this.fn)
         this.fn = newFn
 
-        // Temporarily keep functionality when .asyncResource is removed from node
-        // in https://github.com/nodejs/node/pull/46432
-        if (!this.fn.asyncResource) {
-          this.fn.asyncResource = asyncResource
-        }
+        wrappedFunctions.add(this.fn)
       }
     }
 
@@ -303,5 +301,6 @@ module.exports = {
   getOnTestEndHandler,
   getOnHookEndHandler,
   getOnFailHandler,
-  getOnPendingHandler
+  getOnPendingHandler,
+  testFileToSuiteAr
 }
diff --git a/packages/datadog-plugin-mocha/test/index.spec.js b/packages/datadog-plugin-mocha/test/index.spec.js
@@ -142,6 +142,7 @@ describe('Plugin', () => {
         mocha.addFile(testFilePath)
         mocha.run()
       })
+
       it('works with failing tests', (done) => {
         const testFilePath = path.join(__dirname, 'mocha-test-fail.js')
         const testSuite = testFilePath.replace(`${process.cwd()}/`, '')
@@ -178,6 +179,7 @@ describe('Plugin', () => {
         mocha.addFile(testFilePath)
         mocha.run()
       })
+
       it('works with skipping tests', (done) => {
         const testFilePath = path.join(__dirname, 'mocha-test-skip.js')
         const testNames = [

diff --git a/packages/dd-trace/test/setup/mocha.js b/packages/dd-trace/test/setup/mocha.js
@@ -26,7 +26,11 @@ function loadInst (plugin) {
     loadInstFile(`${plugin}/server.js`, instrumentations)
     loadInstFile(`${plugin}/client.js`, instrumentations)
   } catch (e) {
-    loadInstFile(`${plugin}.js`, instrumentations)
+    try {
+      loadInstFile(`${plugin}/main.js`, instrumentations)
+    } catch (e) {
+      loadInstFile(`${plugin}.js`, instrumentations)
+    }
   }
 
   return instrumentations
@@ -143,10 +147,12 @@ function withNamingSchema (
 function withPeerService (tracer, pluginName, spanGenerationFn, service, serviceSource, opts = {}) {
   describe('peer service computation' + (opts.desc ? ` ${opts.desc}` : ''), () => {
     let computePeerServiceSpy
+
     beforeEach(() => {
       const plugin = tracer()._pluginManager._pluginsByName[pluginName]
       computePeerServiceSpy = sinon.stub(plugin._tracerConfig, 'spanComputePeerService').value(true)
     })
+
     afterEach(() => {
       computePeerServiceSpy.restore()
     })