dd-trace 6.13.0 → 6.15.0

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

@@ -1,21 +1,28 @@
 'use strict'
 
 const { performance } = require('node:perf_hooks')
-const dateNow = Date.now
 const { fileURLToPath } = require('node:url')
 
 const { channel } = require('dc-polyfill')
 
 const CiPlugin = require('../../dd-trace/src/plugins/ci_plugin')
 const { storage } = require('../../datadog-core')
+const { getEnvironmentVariable } = require('../../dd-trace/src/config/helper')
 const {
   getEfdRetryCountForDuration,
   hasEfdRetries,
 } = require('../../dd-trace/src/ci-visibility/efd-retry-policy')
+const {
+  SCREENSHOT_UPLOAD_RESULT_ERROR,
+  SCREENSHOT_UPLOAD_RESULT_UPLOADED,
+  setScreenshotUploadTags,
+} = require('../../dd-trace/src/ci-visibility/test-screenshot')
 const log = require('../../dd-trace/src/log')
 const {
+  requestWebdriverioScreenshotUpload,
   sendWebdriverioWorkerMessage,
   SUITE_FINISH,
+  WEBDRIVERIO_WORKER_ENV,
 } = require('../../datadog-instrumentations/src/mocha/webdriverio-protocol')
 
 const {
@@ -54,6 +61,10 @@ const {
   getFailedTestReplayPromise,
   getTestSuiteExecutionKey,
   isModifiedTest,
+  setRumTestCorrelation,
+  TEST_BROWSER_NAME,
+  TEST_BROWSER_VERSION,
+  TEST_IS_RUM_ACTIVE,
 } = require('../../dd-trace/src/plugins/util/test')
 const { COMPONENT } = require('../../dd-trace/src/constants')
 const {
@@ -83,11 +94,16 @@ const jasmineSpecExecuteStartCh = 'tracing:orchestrion:jasmine-core:Spec_execute
 const jasmineTestFunctionStartCh = 'tracing:orchestrion:@wdio/utils:testFrameworkFnWrapper:start'
 const testFinishCh = channel('ci:mocha:test:finish')
 const testRetryCh = channel('ci:mocha:test:retry')
+const WEBDRIVERIO_FRAMEWORK = 'webdriverio'
 const WEBDRIVERIO_JASMINE_ADAPTER = 'jasmine'
+const WEBDRIVERIO_SCREENSHOT_CAPTURE_TIMEOUT_MS = 30_000
 const workerFinishCh = channel('ci:mocha:worker:finish')
 const WEBDRIVERIO_JASMINE_FAILED_EXPECTATION_COUNT = Symbol('webdriverioJasmineFailedExpectationCount')
 const WEBDRIVERIO_JASMINE_FUNCTION_TYPE = Symbol('webdriverioJasmineFunctionType')
 const WEBDRIVERIO_JASMINE_TEST = Symbol('webdriverioJasmineTest')
+const isWebdriverioWorker = !!getEnvironmentVariable(WEBDRIVERIO_WORKER_ENV)
+const PNG_SIGNATURE = Buffer.from([137, 80, 78, 71, 13, 10, 26, 10])
+const BASE64_RE = /^(?:[A-Za-z\d+/]{4})*(?:[A-Za-z\d+/]{2}==|[A-Za-z\d+/]{3}=)?$/
 
 /**
  * @typedef {object} WebdriverioJasmineResult
@@ -140,6 +156,63 @@ function getJasmineError (result) {
   return error
 }
 
+/**
+ * Applies WebdriverIO RUM correlation without replacing active-browser metadata with an inactive browser.
+ *
+ * @param {object} context
+ * @param {object|undefined} activeSpan
+ * @returns {object|undefined}
+ */
+function setWebdriverioRumTestCorrelation (context, activeSpan) {
+  const browserVersion = context.browserVersion
+  const correlationContext = {
+    isRumActive: context.isRumActive,
+    testExecutionId: context.testExecutionId,
+  }
+  const testSpan = setRumTestCorrelation(correlationContext, activeSpan)
+  context.testExecutionId = correlationContext.testExecutionId
+  if (!testSpan) return
+
+  const hasActiveRum = testSpan.context().getTag(TEST_IS_RUM_ACTIVE)
+  if (!context.isRumActive && hasActiveRum) return testSpan
+  if (context.browserName) testSpan.setTag(TEST_BROWSER_NAME, context.browserName)
+  if (browserVersion) testSpan.setTag(TEST_BROWSER_VERSION, browserVersion)
+  return testSpan
+}
+
+/**
+ * Converts a screenshot capture failure into its upload result.
+ *
+ * @param {unknown} error
+ * @returns {string}
+ */
+function handleWebdriverioScreenshotError (error) {
+  log.error('Error capturing WebdriverIO failure screenshot: %s', error?.message || String(error))
+  return SCREENSHOT_UPLOAD_RESULT_ERROR
+}
+
+/**
+ * Decodes and validates PNG screenshot data returned by WebdriverIO.
+ *
+ * @param {unknown} screenshot - Base64-encoded PNG data
+ * @returns {Buffer} Decoded PNG data
+ */
+function decodeWebdriverioScreenshot (screenshot) {
+  if (typeof screenshot !== 'string' || !BASE64_RE.test(screenshot)) {
+    throw new TypeError('WebdriverIO returned invalid Base64 screenshot data')
+  }
+
+  const content = Buffer.from(screenshot, 'base64')
+  if (
+    content.toString('base64') !== screenshot ||
+    content.length < PNG_SIGNATURE.length ||
+    !content.subarray(0, PNG_SIGNATURE.length).equals(PNG_SIGNATURE)
+  ) {
+    throw new TypeError('WebdriverIO returned invalid PNG screenshot data')
+  }
+  return content
+}
+
 /**
  * Resolves the spec file responsible for a run-level Jasmine failure.
  *
@@ -174,14 +247,37 @@ class MochaPlugin extends CiPlugin {
 
     this._testTitleToParams = {}
     this.sourceRoot = process.cwd()
-    this._pendingTestSuiteSpans = []
-    this._timeOrigin = dateNow()
-    this._perfOrigin = performance.now()
-
-    this.addSub('ci:mocha:session:start', () => {
-      this._pendingTestSuiteSpans = []
-      this._timeOrigin = dateNow()
-      this._perfOrigin = performance.now()
+    this._webdriverioScreenshotUploads = new WeakMap()
+    this._pendingWebdriverioScreenshotUploads = 0
+    this._webdriverioScreenshotUploadCallbacks = []
+
+    this.addSub('ci:webdriverio:screenshot:capabilities', (ctx) => {
+      ctx.enabled = Boolean(
+        this._tracerConfig.testOptimization?.DD_TEST_FAILURE_SCREENSHOTS_ENABLED &&
+        this.tracer._exporter?.canUploadTestScreenshots?.()
+      )
+    })
+
+    this.addSub('ci:webdriverio:screenshot:upload', ({
+      capturedAtMs,
+      idempotencyKey,
+      onDone,
+      screenshot,
+      traceId,
+    }) => {
+      const exporter = this.tracer?._exporter
+      if (!exporter?.canUploadTestScreenshots?.() || !exporter.uploadTestScreenshot) {
+        onDone(new Error('WebdriverIO screenshot upload is not supported by the active Test Optimization transport'))
+        return
+      }
+      let content
+      try {
+        content = decodeWebdriverioScreenshot(screenshot)
+      } catch (error) {
+        onDone(error)
+        return
+      }
+      exporter.uploadTestScreenshot({ content, traceId, idempotencyKey, capturedAtMs }, onDone)
     })
 
     this.addSub('ci:mocha:worker:configuration', ({
@@ -195,6 +291,18 @@ class MochaPlugin extends CiPlugin {
       this.testFramework = testFramework
       this.testFrameworkAdapter = testFrameworkAdapter
       this._setRepositoryRoot(repositoryRoot)
+      if (
+        testFramework === WEBDRIVERIO_FRAMEWORK &&
+        this._tracerConfig.testOptimization.DD_TEST_FAILURE_SCREENSHOTS_ENABLED &&
+        !(isWebdriverioWorker
+          ? libraryConfig.isTestFailureScreenshotsEnabled
+          : this.tracer._exporter?.canUploadTestScreenshots?.())
+      ) {
+        log.warn(
+          'DD_TEST_FAILURE_SCREENSHOTS_ENABLED is true, but WebdriverIO screenshot upload is not supported by the %s',
+          'active Test Optimization transport.'
+        )
+      }
       if (testFrameworkAdapter === WEBDRIVERIO_JASMINE_ADAPTER) {
         this._webdriverioJasmineState = {
           completedTestStatuses: new Map(),
@@ -210,6 +318,14 @@ class MochaPlugin extends CiPlugin {
       }
     })
 
+    this.addSub('ci:webdriverio:rum:page-navigate', (ctx) => {
+      if (this.testFramework !== 'webdriverio') return
+      ctx.isTestOptimizationRunner = true
+
+      const activeSpan = storage('legacy').getStore()?.span
+      setWebdriverioRumTestCorrelation(ctx, activeSpan)
+    })
+
     this.addBind(jasmineTestFunctionStartCh, (ctx) => {
       if (this.testFrameworkAdapter !== WEBDRIVERIO_JASMINE_ADAPTER) {
         return storage('legacy').getStore()
@@ -232,7 +348,7 @@ class MochaPlugin extends CiPlugin {
 
       const functionType = currentStore[WEBDRIVERIO_JASMINE_FUNCTION_TYPE]
       if (functionType === 'Test') {
-        ctx.retryCallback = error => this.#retryWebdriverioJasmineTest(test, error)
+        ctx.retryCallback = this.#retryWebdriverioJasmineTestWithRumCorrelation.bind(this, ctx, test)
       }
       const nextStore = {
         ...test.currentStore,
@@ -271,12 +387,11 @@ class MochaPlugin extends CiPlugin {
     this.addSub(jasmineExecuteAsyncErrorCh, (ctx) => {
       const currentStore = ctx.currentStore || storage('legacy').getStore()
       const test = currentStore?.[WEBDRIVERIO_JASMINE_TEST]
-      if (
-        this.testFrameworkAdapter === WEBDRIVERIO_JASMINE_ADAPTER &&
-        test &&
-        currentStore[WEBDRIVERIO_JASMINE_FUNCTION_TYPE] === 'Hook'
-      ) {
-        test.hasFinalHookFailure = true
+      if (this.testFrameworkAdapter === WEBDRIVERIO_JASMINE_ADAPTER && test) {
+        this.#startWebdriverioScreenshotUpload(test.span)
+        if (currentStore[WEBDRIVERIO_JASMINE_FUNCTION_TYPE] === 'Hook') {
+          test.hasFinalHookFailure = true
+        }
       }
     })
 
@@ -291,6 +406,7 @@ class MochaPlugin extends CiPlugin {
         this._webdriverioJasmineState.currentResult?.failedExpectations?.length > failedExpectationCount
       ) {
         test.hasFinalHookFailure = true
+        this.#startWebdriverioScreenshotUpload(test.span)
       }
     })
 
@@ -482,19 +598,7 @@ class MochaPlugin extends CiPlugin {
         if (!testSuiteSpan.context().getTag(TEST_STATUS)) {
           testSuiteSpan.setTag(TEST_STATUS, status)
         }
-        const exporter = this.tracer._exporter
-        if (exporter.deferTestSuiteSpan) {
-          exporter.deferTestSuiteSpan(testSuiteSpan)
-          testSuiteSpan.finish()
-        } else if (exporter.exportDeferredTestSuiteSpans) {
-          const finishTime = this._now()
-          this._pendingTestSuiteSpans.push({
-            span: testSuiteSpan,
-            finishTime,
-          })
-        } else {
-          testSuiteSpan.finish()
-        }
+        testSuiteSpan.finish()
         this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'suite')
       }
     })
@@ -544,8 +648,12 @@ class MochaPlugin extends CiPlugin {
     })
 
     this.addSub('ci:mocha:worker:finish', ({ onDone } = {}) => {
-      this._finishPendingTestSuiteSpans()
-      this.tracer._exporter.flush(onDone)
+      const flush = () => this.tracer._exporter.flush(onDone)
+      if (this._pendingWebdriverioScreenshotUploads === 0) {
+        flush()
+        return
+      }
+      this._webdriverioScreenshotUploadCallbacks.push(flush)
     })
 
     this.addSub('ci:mocha:test:finish', ({
@@ -562,6 +670,7 @@ class MochaPlugin extends CiPlugin {
       earlyFlakeAbortReason,
     }) => {
       if (span) {
+        const finishTime = span._getTime()
         span.setTag(TEST_STATUS, status)
         if (finalStatus) {
           span.setTag(TEST_FINAL_STATUS, finalStatus)
@@ -590,20 +699,27 @@ class MochaPlugin extends CiPlugin {
           span.setTag(TEST_RETRY_REASON, TEST_RETRY_REASON_TYPES.atf)
         }
 
-        this.telemetry.ciVisEvent(
-          TELEMETRY_EVENT_FINISHED,
-          'test',
-          this.getTestTelemetryTags(span)
-        )
-
-        span.finish()
-        finishAllTraceSpans(span)
-        this.activeTestSpan = null
         this.cancelDiBreakpointHitWait()
         if (this.di && this.libraryConfig?.isDiEnabled && this.runningTestProbe && isLastRetry) {
           this.removeDiProbe(this.runningTestProbe)
           this.runningTestProbe = null
         }
+
+        const finishSpan = () => {
+          this.telemetry.ciVisEvent(
+            TELEMETRY_EVENT_FINISHED,
+            'test',
+            this.getTestTelemetryTags(span)
+          )
+          span.finish(finishTime)
+          finishAllTraceSpans(span)
+          if (this.activeTestSpan === span) {
+            this.activeTestSpan = null
+          }
+        }
+        if (status !== 'fail' || !this.#startWebdriverioScreenshotUpload(span, finishSpan)) {
+          finishSpan()
+        }
       }
     })
 
@@ -633,6 +749,7 @@ class MochaPlugin extends CiPlugin {
         } else {
           span.setTag(TEST_STATUS, 'fail')
           span.setTag('error', err)
+          this.#startWebdriverioScreenshotUpload(span)
         }
 
         ctx.parentStore = ctx.currentStore
@@ -657,8 +774,9 @@ class MochaPlugin extends CiPlugin {
       promises,
     }) => {
       if (span) {
+        const finishTime = span._getTime()
         const finishSpan = () => {
-          span.finish()
+          span.finish(finishTime)
           finishAllTraceSpans(span)
           if (this.activeTestSpan === span) {
             this.activeTestSpan = null
@@ -701,17 +819,42 @@ class MochaPlugin extends CiPlugin {
           }
         }
 
+        let finishWait
         if (!isFirstFailure &&
           willBeRetried &&
           this.di &&
           this.libraryConfig?.isDiEnabled &&
-          this.runningTestProbe &&
-          promises) {
-          promises.finishTestPromise = this.waitForInFlightDiBreakpointHits().then(finishSpan, finishSpan)
-          return
+          this.runningTestProbe) {
+          finishWait = this.waitForInFlightDiBreakpointHits()
         }
-
-        finishSpan()
+        let screenshotFinished = false
+        let canFinishScreenshot = false
+        let replayFinished = !finishWait
+        let spanFinished = false
+        const finishWhenReady = () => {
+          if (!spanFinished && replayFinished && (screenshotStarted === false || screenshotFinished)) {
+            spanFinished = true
+            finishSpan()
+          }
+        }
+        const screenshotStarted = this.#startWebdriverioScreenshotUpload(span, () => {
+          screenshotFinished = true
+          if (canFinishScreenshot) {
+            finishWhenReady()
+          }
+        })
+        canFinishScreenshot = true
+        if (finishWait) {
+          const finishReplay = () => {
+            replayFinished = true
+            finishWhenReady()
+          }
+          const finishTestPromise = finishWait.then(finishReplay, finishReplay)
+          if (promises) {
+            promises.finishTestPromise = finishTestPromise
+          }
+        }
+        finishWhenReady()
       }
     })
 
@@ -755,7 +898,6 @@ class MochaPlugin extends CiPlugin {
           this.testSessionSpan.setTag('error', error)
           this.testModuleSpan.setTag('error', error)
           if (isFrameworkError) {
-            this.tracer._exporter.setDeferredTestSuiteError?.(error)
             for (const testSuiteSpan of this._testSuiteSpansByTestSuite.values()) {
               testSuiteSpan.setTag(TEST_STATUS, 'fail')
               testSuiteSpan.setTag('error', error)
@@ -763,8 +905,6 @@ class MochaPlugin extends CiPlugin {
           }
         }
 
-        this._finishPendingTestSuiteSpans()
-
         if (isParallel) {
           this.testSessionSpan.setTag(MOCHA_IS_PARALLEL, 'true')
         }
@@ -802,7 +942,6 @@ class MochaPlugin extends CiPlugin {
           this.testSessionSpan.setTag(TEST_EARLY_FLAKE_ABORT_REASON, 'faulty')
         }
 
-        this.tracer._exporter.exportDeferredTestSuiteSpans?.()
         this.testModuleSpan.finish()
         this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'module')
         this.testSessionSpan.finish()
@@ -824,27 +963,6 @@ class MochaPlugin extends CiPlugin {
     })
   }
 
-  /**
-   * Returns the current time in the coordinate system used by spans.
-   *
-   * @returns {number}
-   */
-  _now () {
-    return this._timeOrigin + performance.now() - this._perfOrigin
-  }
-
-  /**
-   * Finishes suites retained for a later framework finalization event.
-   *
-   * @returns {void}
-   */
-  _finishPendingTestSuiteSpans () {
-    for (const { span, finishTime } of this._pendingTestSuiteSpans) {
-      span.finish(finishTime)
-    }
-    this._pendingTestSuiteSpans = []
-  }
-
   /**
    * Starts a Jasmine test span around WebdriverIO's test-function wrapper.
    *
@@ -1080,6 +1198,9 @@ class MochaPlugin extends CiPlugin {
     }
 
     const testStatus = getJasmineStatus(status)
+    if (testStatus === 'fail') {
+      this.#startWebdriverioScreenshotUpload(test.span)
+    }
     if (
       testStatus !== 'skip' &&
       test.isEarlyFlakeDetection &&
@@ -1177,6 +1298,25 @@ class MochaPlugin extends CiPlugin {
     this.#startWebdriverioJasmineAttempt(test, test.currentStore)
   }
 
+  /**
+   * Advances a native WebdriverIO retry without stopping its active RUM session.
+   *
+   * @param {object} context
+   * @param {object} test
+   * @param {Error|undefined} error
+   * @returns {Promise<void>}
+   */
+  async #retryWebdriverioJasmineTestWithRumCorrelation (context, test, error) {
+    const correlationContext = await context.rumRetryCallback?.() || {}
+    setWebdriverioRumTestCorrelation(correlationContext, test.span)
+
+    await this.#retryWebdriverioJasmineTest(test, error)
+
+    const testExecutionId = test.span.context().toTraceId()
+    await context.rumRetryCallback?.(testExecutionId)
+    setWebdriverioRumTestCorrelation(correlationContext, test.span)
+  }
+
   /**
    * Finishes an intermediate native WebdriverIO retry and starts its next span.
    *
@@ -1239,8 +1379,7 @@ class MochaPlugin extends CiPlugin {
     if (test._ddShouldWaitForHitProbe) {
       delete test._ddShouldWaitForHitProbe
       state.pendingTestFinishes++
-      const finishTest = () => {
-        this.#completeWebdriverioJasmineTest(test, result)
+      const onTestFinished = () => {
         state.pendingTestFinishes--
         if (state.pendingTestFinishes === 0) {
           const callbacks = state.pendingTestFinishCallbacks
@@ -1250,10 +1389,14 @@ class MochaPlugin extends CiPlugin {
           }
         }
       }
+      const finishTest = () => {
+        const finishPromise = this.#completeWebdriverioJasmineTest(test, result)
+        return finishPromise?.then(onTestFinished, onTestFinished) || onTestFinished()
+      }
       return this.waitForDiBreakpointHits().then(finishTest, finishTest)
     }
 
-    this.#completeWebdriverioJasmineTest(test, result)
+    return this.#completeWebdriverioJasmineTest(test, result)
   }
 
   /**
@@ -1312,6 +1455,191 @@ class MochaPlugin extends CiPlugin {
     }
   }
 
+  /**
+   * Captures and uploads the current WebdriverIO browser state once for a failed test attempt.
+   *
+   * @param {object} span - Failed test span
+   * @param {() => void} [onDone] - Called after the screenshot upload finishes
+   * @returns {boolean} Whether a screenshot upload exists for the span
+   */
+  #startWebdriverioScreenshotUpload (span, onDone) {
+    if (
+      !span ||
+      this.testFramework !== WEBDRIVERIO_FRAMEWORK ||
+      !this._tracerConfig.testOptimization?.DD_TEST_FAILURE_SCREENSHOTS_ENABLED
+    ) {
+      return false
+    }
+
+    const previousUpload = this._webdriverioScreenshotUploads.get(span)
+    if (previousUpload) {
+      if (onDone) {
+        if (previousUpload.finished) {
+          onDone()
+        } else {
+          previousUpload.callbacks.push(onDone)
+        }
+      }
+      return true
+    }
+
+    const exporter = this.tracer?._exporter
+    const canUpload = isWebdriverioWorker
+      ? this.libraryConfig?.isTestFailureScreenshotsEnabled
+      : exporter?.canUploadTestScreenshots?.() && exporter.uploadTestScreenshot
+    if (!canUpload) {
+      return false
+    }
+
+    const browser = globalThis._wdioGlobals?.get?.('browser') || globalThis.browser
+    const upload = {
+      callbacks: onDone ? [onDone] : [],
+      finished: false,
+    }
+    this._webdriverioScreenshotUploads.set(span, upload)
+    this._pendingWebdriverioScreenshotUploads++
+    const finishUpload = (result) => this.#finishWebdriverioScreenshotUpload(span, upload, result)
+    let captureTimeout
+    const failCapture = (error) => {
+      if (upload.finished) return
+      clearTimeout(captureTimeout)
+      captureTimeout = undefined
+      finishUpload(handleWebdriverioScreenshotError(error))
+    }
+    const uploadScreenshots = (content) => {
+      if (upload.finished) return
+      clearTimeout(captureTimeout)
+      captureTimeout = undefined
+      try {
+        this.#uploadWebdriverioScreenshots(span, content, finishUpload)
+      } catch (error) {
+        failCapture(error)
+      }
+    }
+    captureTimeout = setTimeout(() => {
+      failCapture(new Error(
+        `WebdriverIO screenshot capture timed out after ${WEBDRIVERIO_SCREENSHOT_CAPTURE_TIMEOUT_MS}ms`
+      ))
+    }, WEBDRIVERIO_SCREENSHOT_CAPTURE_TIMEOUT_MS)
+    captureTimeout.unref?.()
+    try {
+      if (typeof browser?.takeScreenshot !== 'function') {
+        throw new TypeError('browser.takeScreenshot is not available')
+      }
+      const screenshot = browser.takeScreenshot()
+      if (typeof screenshot?.then === 'function') {
+        screenshot.then(
+          uploadScreenshots,
+          failCapture
+        )
+      } else {
+        uploadScreenshots(screenshot)
+      }
+    } catch (error) {
+      failCapture(error)
+    }
+    return true
+  }
+
+  /**
+   * Completes one WebdriverIO screenshot upload and its waiters.
+   *
+   * @param {object} span - Failed test span
+   * @param {{callbacks: Array<() => void>, finished: boolean}} upload - Upload state
+   * @param {string} result - Aggregate screenshot upload result
+   * @returns {void}
+   */
+  #finishWebdriverioScreenshotUpload (span, upload, result) {
+    if (upload.finished) return
+
+    upload.finished = true
+    setScreenshotUploadTags(span, result)
+    const uploadCallbacks = upload.callbacks
+    upload.callbacks = []
+
+    this._pendingWebdriverioScreenshotUploads--
+    let workerCallbacks = []
+    if (this._pendingWebdriverioScreenshotUploads === 0) {
+      workerCallbacks = this._webdriverioScreenshotUploadCallbacks
+      this._webdriverioScreenshotUploadCallbacks = []
+    }
+    for (const callback of uploadCallbacks) {
+      callback()
+    }
+    for (const callback of workerCallbacks) {
+      callback()
+    }
+  }
+
+  /**
+   * Uploads the PNG data returned by WebdriverIO for one or more browser sessions.
+   *
+   * @param {object} span - Failed test span
+   * @param {string|string[]} screenshots - Base64-encoded PNG data
+   * @param {(result: string) => void} onDone - Aggregate upload completion callback
+   * @returns {void}
+   */
+  #uploadWebdriverioScreenshots (span, screenshots, onDone) {
+    const screenshotList = Array.isArray(screenshots) ? screenshots : [screenshots]
+    if (screenshotList.length === 0) {
+      log.error('Error capturing WebdriverIO failure screenshot: WebdriverIO returned no screenshot data')
+      onDone(SCREENSHOT_UPLOAD_RESULT_ERROR)
+      return
+    }
+    const exporter = this.tracer?._exporter
+    const traceId = span.context().toTraceId()
+    const capturedAtMs = Date.now()
+    let pendingUploads = screenshotList.length
+    let hasUploadError = false
+    const finishOne = (error) => {
+      if (error) {
+        hasUploadError = true
+        log.error('Error uploading WebdriverIO failure screenshot: %s', error?.message || String(error))
+      }
+      if (--pendingUploads === 0) {
+        onDone(hasUploadError ? SCREENSHOT_UPLOAD_RESULT_ERROR : SCREENSHOT_UPLOAD_RESULT_UPLOADED)
+      }
+    }
+    for (let index = 0; index < screenshotList.length; index++) {
+      const screenshot = screenshotList[index]
+      if (typeof screenshot !== 'string' || !screenshot) {
+        finishOne(new TypeError('WebdriverIO returned invalid screenshot data'))
+        continue
+      }
+      const options = {
+        screenshot,
+        traceId,
+        idempotencyKey: `${traceId}:webdriverio-failure-${index}.png`,
+        capturedAtMs,
+      }
+      if (isWebdriverioWorker) {
+        try {
+          requestWebdriverioScreenshotUpload(options, finishOne)
+        } catch (error) {
+          finishOne(error)
+        }
+      } else {
+        let content
+        try {
+          content = decodeWebdriverioScreenshot(screenshot)
+        } catch (error) {
+          finishOne(error)
+          continue
+        }
+        try {
+          exporter.uploadTestScreenshot({
+            content,
+            traceId,
+            idempotencyKey: options.idempotencyKey,
+            capturedAtMs,
+          }, finishOne)
+        } catch (error) {
+          finishOne(error)
+        }
+      }
+    }
+  }
+
   /**
    * Reports Jasmine suite statuses and flushes worker traces before its run settles.
    *

package/packages/datadog-plugin-mongodb-core/src/query.js

@@ -98,7 +98,13 @@ class MongodbCoreQueryPlugin extends DatabasePlugin {
 const MAX_DEPTH = 10
 const MAX_QUERY_LENGTH = 10_000
 
+/** @param {{ q?: Record<string, unknown> }[] | { documents?: { q?: Record<string, unknown> }[] }} statements */
 function extractQuery (statements) {
+  if (!Array.isArray(statements)) {
+    const { documents } = statements
+    statements = Array.isArray(documents) ? documents : []
+  }
+
   if (statements.length === 1 && statements[0].q) return statements[0].q
 
   const extractedQueries = []