From 2439b3357a88e4f1e17b8ed3d55a7be0b112a799 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Juan=20Antonio=20Fern=C3=A1ndez=20de=20Alba?=
 <juan.fernandezdealba@datadoghq.com>
Date: Mon, 20 May 2024 16:43:51 +0200
Subject: [PATCH] [ci-visibility] Fix selenium when run outside of a supported
 test framework (#4330)

---
 .../test/selenium-no-framework.js             | 30 +++++++++++++++++++
 integration-tests/selenium/selenium.spec.js   | 29 ++++++++++++++++++
 .../datadog-instrumentations/src/selenium.js  | 19 ++++++++----
 3 files changed, 72 insertions(+), 6 deletions(-)
 create mode 100644 integration-tests/ci-visibility/test/selenium-no-framework.js

diff --git a/integration-tests/ci-visibility/test/selenium-no-framework.js b/integration-tests/ci-visibility/test/selenium-no-framework.js
new file mode 100644
index 00000000000..cca24586bfd
--- /dev/null
+++ b/integration-tests/ci-visibility/test/selenium-no-framework.js
@@ -0,0 +1,30 @@
+const { By, Builder } = require('selenium-webdriver')
+const chrome = require('selenium-webdriver/chrome')
+
+async function run () {
+  const options = new chrome.Options()
+  options.addArguments('--headless')
+  const build = new Builder().forBrowser('chrome').setChromeOptions(options)
+  const driver = await build.build()
+
+  await driver.get(process.env.WEB_APP_URL)
+
+  await driver.getTitle()
+
+  await driver.manage().setTimeouts({ implicit: 500 })
+
+  const helloWorld = await driver.findElement(By.className('hello-world'))
+
+  await helloWorld.getText()
+
+  return driver.quit()
+}
+
+run()
+  .then(() => {
+    process.exit(0)
+  }).catch((err) => {
+    // eslint-disable-next-line no-console
+    console.error(err)
+    process.exit(1)
+  })
diff --git a/integration-tests/selenium/selenium.spec.js b/integration-tests/selenium/selenium.spec.js
index 7cec0de791b..e7f2404a2db 100644
--- a/integration-tests/selenium/selenium.spec.js
+++ b/integration-tests/selenium/selenium.spec.js
@@ -113,5 +113,34 @@ versionRange.forEach(version => {
         })
       })
     })
+
+    it('does not crash when used outside a known test framework', (done) => {
+      let testOutput = ''
+      childProcess = exec(
+        'node ./ci-visibility/test/selenium-no-framework.js',
+        {
+          cwd,
+          env: {
+            ...getCiVisAgentlessConfig(receiver.port),
+            WEB_APP_URL: `http://localhost:${webAppPort}`,
+            TESTS_TO_RUN: '**/ci-visibility/test/selenium-test*'
+          },
+          stdio: 'pipe'
+        }
+      )
+
+      childProcess.on('exit', (code) => {
+        assert.equal(code, 0)
+        assert.notInclude(testOutput, 'InvalidArgumentError')
+        done()
+      })
+
+      childProcess.stdout.on('data', (chunk) => {
+        testOutput += chunk.toString()
+      })
+      childProcess.stderr.on('data', (chunk) => {
+        testOutput += chunk.toString()
+      })
+    })
   })
 })
diff --git a/packages/datadog-instrumentations/src/selenium.js b/packages/datadog-instrumentations/src/selenium.js
index d018c97dfb9..141aa967e40 100644
--- a/packages/datadog-instrumentations/src/selenium.js
+++ b/packages/datadog-instrumentations/src/selenium.js
@@ -23,6 +23,9 @@ addHook({
 }, (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)
+    }
     let traceId
     const setTraceId = (inputTraceId) => {
       traceId = inputTraceId
@@ -40,15 +43,20 @@ addHook({
       isRumActive
     })
 
-    await this.manage().addCookie({
-      name: DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME,
-      value: traceId
-    })
+    if (traceId && isRumActive) {
+      await this.manage().addCookie({
+        name: DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME,
+        value: traceId
+      })
+    }
 
     return getResult
   })
 
   shimmer.wrap(seleniumPackage.WebDriver.prototype, 'quit', quit => async function () {
+    if (!ciSeleniumDriverGetStartCh.hasSubscribers) {
+      return quit.apply(this, arguments)
+    }
     const isRumActive = await this.executeScript(RUM_STOP_SESSION_SCRIPT)
 
     if (isRumActive) {
@@ -58,10 +66,9 @@ addHook({
           resolve()
         }, DD_CIVISIBILITY_RUM_FLUSH_WAIT_MILLIS)
       })
+      await this.manage().deleteCookie(DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME)
     }
 
-    await this.manage().deleteCookie(DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME)
-
     return quit.apply(this, arguments)
   })