fix(test-optimization): wait for WebDriverIO worker payloads (#9565)

Author: juan-fernandez

packages/datadog-instrumentations/src/mocha/worker.js

@@ -42,20 +42,22 @@ let configurationRequestId = 0
  *
  * @param {object} message
  * @param {() => void} [onError]
+ * @param {() => void} [onDone]
  * @returns {void}
  */
-function sendWebdriverioMessage (message, onError) {
+function sendWebdriverioMessage (message, onError, onDone) {
   if (!process.send || !process.connected) {
     onError?.()
+    onDone?.()
     return
   }
 
   process.send(createWebdriverioWorkerMessage(message), (error) => {
-    if (!error) {
-      return
+    if (error) {
+      log.error('WebdriverIO Test Optimization IPC error', error)
+      onError?.()
     }
-    log.error('WebdriverIO Test Optimization IPC error', error)
-    onError?.()
+    onDone?.()
   })
 }
 
@@ -266,10 +268,12 @@ function getWebdriverioSuiteResults (runner) {
  * Sends suite results to the WebdriverIO launcher.
  *
  * @param {object} runner
+ * @param {() => void} [onDone]
  * @returns {void}
  */
-function reportWebdriverioSuiteResults (runner) {
+function reportWebdriverioSuiteResults (runner, onDone) {
   if (!isWebdriverioWorker) {
+    onDone?.()
     return
   }
 
@@ -279,7 +283,32 @@ function reportWebdriverioSuiteResults (runner) {
     content: {
       results: getWebdriverioSuiteResults(runner),
     },
-  })
+  }, undefined, onDone)
+}
+
+/**
+ * Flushes worker payloads before reporting suite results and completing Mocha.
+ *
+ * @param {object} runner
+ * @param {() => void} onDone
+ * @returns {void}
+ */
+function finishWebdriverioWorker (runner, onDone) {
+  try {
+    workerFinishCh.publish({
+      onDone: () => {
+        try {
+          reportWebdriverioSuiteResults(runner, onDone)
+        } catch (error) {
+          log.error('WebdriverIO Test Optimization worker completion error', error)
+          onDone()
+        }
+      },
+    })
+  } catch (error) {
+    log.error('WebdriverIO Test Optimization worker completion error', error)
+    onDone()
+  }
 }
 
 function isFailedTestReplayEnabled () {
@@ -376,11 +405,22 @@ addHook({
     if (isFailedTestReplayEnabled()) {
       patchFailedTestReplayHookUp(Runner)
     }
-    // Flush after the worker finishes its Mocha run, including grouped spec files.
-    this.once('end', () => {
-      workerFinishCh.publish()
-      reportWebdriverioSuiteResults(this)
-    })
+    const onRunDone = args[0]
+    if (isWebdriverioWorker && typeof onRunDone === 'function') {
+      args[0] = (...onRunDoneArgs) => {
+        finishWebdriverioWorker(this, () => onRunDone(...onRunDoneArgs))
+      }
+    } else {
+      // Flush after the worker finishes its Mocha run, including grouped spec files.
+      this.once('end', () => {
+        try {
+          workerFinishCh.publish()
+          reportWebdriverioSuiteResults(this)
+        } catch (error) {
+          log.error('WebdriverIO Test Optimization worker completion error', error)
+        }
+      })
+    }
     this.on('test', getOnTestHandler(false))
 
     this.on('test end', getOnTestEndHandler(config))

packages/datadog-instrumentations/test/fixtures/webdriverio-delayed-worker.js

@@ -0,0 +1,94 @@
+'use strict'
+
+const assert = require('node:assert/strict')
+const { EventEmitter } = require('node:events')
+
+const { channel } = require('../../src/helpers/instrument')
+const instrumentations = require('../../src/helpers/instrumentations')
+const {
+  SUITE_FINISH,
+  WEBDRIVERIO_WORKER_ENV,
+  WEBDRIVERIO_WORKER_EVENT,
+  WEBDRIVERIO_WORKER_ORIGIN,
+} = require('../../src/mocha/webdriverio-protocol')
+
+process.env[WEBDRIVERIO_WORKER_ENV] = 'true'
+const existingMochaHookCount = instrumentations.mocha?.length || 0
+require('../../src/mocha/worker')
+
+const runnerHook = instrumentations.mocha
+  .slice(existingMochaHookCount)
+  .find(({ file }) => file === 'lib/runner.js')
+
+assert.ok(runnerHook)
+
+const workerFinishCh = channel('ci:mocha:worker:finish')
+let flushDone
+
+function onWorkerFinish ({ onDone }) {
+  flushDone = onDone
+}
+
+workerFinishCh.subscribe(onWorkerFinish)
+
+class FakeRunner extends EventEmitter {
+  constructor () {
+    super()
+    this.failures = 0
+    this.suite = {
+      eachTest () {},
+    }
+  }
+
+  runTests () {}
+
+  run (onDone) {
+    this.emit('end')
+    onDone()
+  }
+}
+
+runnerHook.hook(FakeRunner)
+
+let runDone = false
+let suiteMessageDone
+process.connected = true
+process.send = (message, onDone) => {
+  assert.strictEqual(message.origin, WEBDRIVERIO_WORKER_ORIGIN)
+  assert.strictEqual(message.name, WEBDRIVERIO_WORKER_EVENT)
+  assert.strictEqual(message.args.name, SUITE_FINISH)
+  suiteMessageDone = onDone
+}
+
+new FakeRunner().run(() => {
+  runDone = true
+})
+
+assert.strictEqual(runDone, false)
+assert.strictEqual(typeof flushDone, 'function')
+assert.strictEqual(suiteMessageDone, undefined)
+
+flushDone()
+assert.strictEqual(runDone, false)
+assert.strictEqual(typeof suiteMessageDone, 'function')
+
+suiteMessageDone()
+assert.strictEqual(runDone, true)
+
+runDone = false
+flushDone = undefined
+suiteMessageDone = undefined
+process.connected = false
+
+new FakeRunner().run(() => {
+  runDone = true
+})
+
+assert.strictEqual(runDone, false)
+assert.strictEqual(typeof flushDone, 'function')
+
+flushDone()
+assert.strictEqual(runDone, true)
+assert.strictEqual(suiteMessageDone, undefined)
+
+workerFinishCh.unsubscribe(onWorkerFinish)

packages/datadog-instrumentations/test/webdriverio.spec.js

@@ -24,6 +24,7 @@ const {
 } = require('../src/mocha/webdriverio-protocol')
 
 const fixturePath = path.join(__dirname, 'fixtures', 'webdriverio-local-runner.mjs')
+const delayedWorkerFixturePath = path.join(__dirname, 'fixtures', 'webdriverio-delayed-worker.js')
 const disconnectedWorkerFixturePath = path.join(__dirname, 'fixtures', 'webdriverio-disconnected-worker.js')
 const regularMochaWorkerFixturePath = path.join(__dirname, 'fixtures', 'mocha-regular-worker.js')
 const fixtureModulePath = path.join(
@@ -150,6 +151,10 @@ describe('webdriverio instrumentation', () => {
     await execFileAsync(process.execPath, [disconnectedWorkerFixturePath])
   })
 
+  it('waits for worker payloads before completing Mocha', async () => {
+    await execFileAsync(process.execPath, [delayedWorkerFixturePath])
+  })
+
   it('does not track WebdriverIO hook failures in regular Mocha workers', async () => {
     await execFileAsync(process.execPath, [regularMochaWorkerFixturePath])
   })

packages/datadog-plugin-mocha/src/index.js

@@ -216,8 +216,8 @@ class MochaPlugin extends CiPlugin {
       return ctx.currentStore
     })
 
-    this.addSub('ci:mocha:worker:finish', () => {
-      this.tracer._exporter.flush()
+    this.addSub('ci:mocha:worker:finish', ({ onDone } = {}) => {
+      this.tracer._exporter.flush(onDone)
     })
 
     this.addSub('ci:mocha:test:finish', ({