dd-trace/packages/datadog-instrumentations/src/selenium.js

Datadog APM tracing client for JavaScript

'use strict'

// Capture real timers at module load time, before any test can install fake timers.
const realSetTimeout = setTimeout

const shimmer = require('../../datadog-shimmer')
const { getValueFromEnvSources } = require('../../dd-trace/src/config/helper')
const {
  RUM_TEST_EXECUTION_ID_COOKIE_NAME: DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME,
} = require('../../dd-trace/src/ci-visibility/rum')
const log = require('../../dd-trace/src/log')
const { addHook, channel } = require('./helpers/instrument')

const ciSeleniumDriverGetStartCh = channel('ci:selenium:driver:get')

const RUM_STOP_SESSION_SCRIPT = `
if (window.DD_RUM && window.DD_RUM.stopSession) {
  window.DD_RUM.stopSession();
  return true;
} else {
  return false;
}
`
const IS_RUM_ACTIVE_SCRIPT = 'return !!window.DD_RUM'

const DD_CIVISIBILITY_RUM_FLUSH_WAIT_MILLIS = getValueFromEnvSources('DD_CIVISIBILITY_RUM_FLUSH_WAIT_MILLIS')

// TODO: can we increase the supported version range?
addHook({
  name: 'selenium-webdriver',
  versions: ['>=4.11.0'],
}, (seleniumPackage, seleniumVersion) => {
  // TODO: do not turn this into async. Use promises
  shimmer.wrap(seleniumPackage.WebDriver.prototype, 'get', get => async function () {
    if (!ciSeleniumDriverGetStartCh.hasSubscribers) {
      return get.apply(this, arguments)
    }
    const getResult = await get.apply(this, arguments)

    let isRumActive
    try {
      isRumActive = await this.executeScript(IS_RUM_ACTIVE_SCRIPT)
    } catch (error) {
      log.error('Selenium RUM detection error', error)
    }

    let browserName
    let browserVersion
    try {
      const capabilities = await this.getCapabilities()
      browserName = capabilities.getBrowserName()
      browserVersion = capabilities.getBrowserVersion()
    } catch (error) {
      log.error('Selenium browser metadata error', error)
    }

    const context = {
      seleniumVersion,
      browserName,
      browserVersion,
      isRumActive,
      testExecutionId: undefined,
    }
    try {
      ciSeleniumDriverGetStartCh.publish(context)
    } catch (error) {
      log.error('Selenium RUM correlation channel error', error)
    }

    if (context.testExecutionId && isRumActive) {
      try {
        await this.manage().addCookie({
          name: DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME,
          value: context.testExecutionId,
        })
      } catch (error) {
        log.error('Selenium RUM correlation cookie error', error)
      }
    }

    return getResult
  })

  shimmer.wrap(seleniumPackage.WebDriver.prototype, 'quit', quit => async function () {
    if (!ciSeleniumDriverGetStartCh.hasSubscribers) {
      return quit.apply(this, arguments)
    }
    try {
      const isRumActive = await this.executeScript(RUM_STOP_SESSION_SCRIPT)

      if (isRumActive) {
        // We'll have time for RUM to flush the events (there's no callback to know when it's done)
        await new Promise(resolve => {
          realSetTimeout(() => {
            resolve()
          }, DD_CIVISIBILITY_RUM_FLUSH_WAIT_MILLIS)
        })
        await this.manage().deleteCookie(DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME)
      }
    } catch (error) {
      log.error('Selenium RUM cleanup error', error)
    }

    return quit.apply(this, arguments)
  })

  return seleniumPackage
})