feat(ci_visibility): support selenium in test visibility (#11610)

Introduces support for Selenium, including integration with the RUM
product for session recording/replays (requested in #10203 ).

A new `hatch` environment is added which tests support for the feature
in both `v1` and `v2` versions of the `pytest` plugin, with a matching
addition to test suite and specs.

A few other changes are added:
- properly adding the `test.type` tag when using the test API (used by
manual test runners and the new pytest plugin)
- adding the `type` tag when the test span is originally created (so
that it may be accessed by the Selenium integration)
- telemetry for events created/finished for tests is being split out
into its own function because the multipurpose single function was
becoming unreasonably branch-y with the addition of yet another set of
tags

## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

---------

Co-authored-by: Vítor De Araújo <[email protected]>

.github/CODEOWNERS

@@ -51,7 +51,7 @@ tests/opentracer                    @DataDog/apm-core-python
 tests/runtime                       @DataDog/apm-core-python
 tests/tracer                        @DataDog/apm-core-python
 
-# CI App and related
+# Test Visibility and related
 ddtrace/contrib/asynctest                             @DataDog/ci-app-libraries
 ddtrace/contrib/coverage                              @DataDog/ci-app-libraries
 ddtrace/contrib/pytest                                @DataDog/ci-app-libraries
@@ -81,6 +81,11 @@ scripts/ci_visibility/*                               @DataDog/ci-app-libraries
 ddtrace/contrib/freezegun                             @DataDog/ci-app-libraries
 ddtrace/contrib/internal/freezegun                    @DataDog/ci-app-libraries
 tests/contrib/freezegun                               @DataDog/ci-app-libraries
+# Test Visibility: Selenium integration
+ddtrace/contrib/selenium                              @DataDog/ci-app-libraries
+ddtrace/internal/selenium                             @DataDog/ci-app-libraries
+tests/contrib/selenium                                @DataDog/ci-app-libraries
+tests/snapshots/test_selenium_*                       @DataDog/ci-app-libraries
 
 # Debugger
 ddtrace/debugging/                  @DataDog/debugger-python

ddtrace/_monkey.py

@@ -101,6 +101,7 @@
     "subprocess": True,
     "unittest": True,
     "coverage": False,
+    "selenium": True,
 }
 
 

ddtrace/contrib/internal/selenium/patch.py

@@ -0,0 +1,185 @@
+import os
+import time
+import typing as t
+
+from wrapt.importer import when_imported
+
+from ddtrace import config
+from ddtrace.internal.logger import get_logger
+from ddtrace.internal.wrapping.context import WrappingContext
+import ddtrace.tracer
+
+
+if t.TYPE_CHECKING:
+    import selenium.webdriver.remote.webdriver
+
+log = get_logger(__name__)
+
+T = t.TypeVar("T")
+
+_RUM_STOP_SESSION_SCRIPT = """
+if (window.DD_RUM && window.DD_RUM.stopSession) {
+    window.DD_RUM.stopSession();
+    return true;
+} else {
+    return false;
+}
+"""
+
+_DEFAULT_FLUSH_SLEEP_MS = 500
+
+
+def _get_flush_sleep_ms() -> int:
+    env_flush_sleep_ms = os.getenv("DD_CIVISIBILITY_RUM_FLUSH_WAIT_MILLIS")
+    if env_flush_sleep_ms is None:
+        return _DEFAULT_FLUSH_SLEEP_MS
+
+    try:
+        return int(env_flush_sleep_ms)
+    except Exception:  # noqa E722
+        log.warning(
+            "Could not convert DD_CIVISIBILITY_RUM_FLUSH_WAIT_MILLIS value %s to int, using default: %s",
+            env_flush_sleep_ms,
+            _DEFAULT_FLUSH_SLEEP_MS,
+        )
+        return _DEFAULT_FLUSH_SLEEP_MS
+
+
+config._add(
+    "selenium",
+    dict(flush_sleep_ms=_get_flush_sleep_ms()),
+)
+
+
+class SeleniumWrappingContextBase(WrappingContext):
+    def _handle_enter(self) -> None:
+        pass
+
+    def _handle_return(self) -> None:
+        pass
+
+    def _get_webdriver_instance(self) -> "selenium.webdriver.remote.webdriver.WebDriver":
+        try:
+            return self.get_local("self")
+        except KeyError:
+            log.debug("Could not get Selenium WebDriver instance")
+            return None
+
+    def __enter__(self) -> "SeleniumWrappingContextBase":
+        super().__enter__()
+
+        try:
+            self._handle_enter()
+        except Exception:  # noqa: E722
+            log.debug("Error handling selenium instrumentation enter", exc_info=True)
+
+        return self
+
+    def __return__(self, value: T) -> T:
+        """Always return the original value no matter what our instrumentation does"""
+        try:
+            self._handle_return()
+        except Exception:  # noqa: E722
+            log.debug("Error handling instrumentation return", exc_info=True)
+
+        return value
+
+
+class SeleniumGetWrappingContext(SeleniumWrappingContextBase):
+    def _handle_return(self) -> None:
+        root_span = ddtrace.tracer.current_root_span()
+        test_trace_id = root_span.trace_id
+
+        if root_span is None or root_span.get_tag("type") != "test":
+            return
+
+        webdriver_instance = self._get_webdriver_instance()
+
+        if webdriver_instance is None:
+            return
+
+        # The trace IDs for Test Visibility data using the CIVisibility protocol are 64-bit
+        # TODO[ci_visibility]: properly identify whether to use 64 or 128 bit trace_ids
+        trace_id_64bit = test_trace_id % 2**64
+
+        webdriver_instance.add_cookie({"name": "datadog-ci-visibility-test-execution-id", "value": str(trace_id_64bit)})
+
+        root_span.set_tag("test.is_browser", "true")
+        root_span.set_tag("test.browser.driver", "selenium")
+        root_span.set_tag("test.browser.driver_version", get_version())
+
+        # Submit empty values for browser names or version if multiple are found
+        browser_name = webdriver_instance.capabilities.get("browserName")
+        browser_version = webdriver_instance.capabilities.get("browserVersion")
+
+        existing_browser_name = root_span.get_tag("test.browser.name")
+        if existing_browser_name is None:
+            root_span.set_tag("test.browser.name", browser_name)
+        elif existing_browser_name not in ["", browser_name]:
+            root_span.set_tag("test.browser.name", "")
+
+        existing_browser_version = root_span.get_tag("test.browser.version")
+        if existing_browser_version is None:
+            root_span.set_tag("test.browser.version", browser_version)
+        elif existing_browser_version not in ["", browser_version]:
+            root_span.set_tag("test.browser.version", "")
+
+
+class SeleniumQuitWrappingContext(SeleniumWrappingContextBase):
+    def _handle_enter(self) -> None:
+        root_span = ddtrace.tracer.current_root_span()
+
+        if root_span is None or root_span.get_tag("type") != "test":
+            return
+
+        webdriver_instance = self._get_webdriver_instance()
+
+        if webdriver_instance is None:
+            return
+
+        is_rum_active = webdriver_instance.execute_script(_RUM_STOP_SESSION_SCRIPT)
+        time.sleep(config.selenium.flush_sleep_ms / 1000)
+
+        if is_rum_active:
+            root_span.set_tag("test.is_rum_active", "true")
+
+        webdriver_instance.delete_cookie("datadog-ci-visibility-test-execution-id")
+
+
+def get_version() -> str:
+    import selenium
+
+    try:
+        return selenium.__version__
+    except AttributeError:
+        log.debug("Could not get Selenium version")
+        return ""
+
+
+def patch() -> None:
+    import selenium
+
+    if getattr(selenium, "_datadog_patch", False):
+        return
+
+    @when_imported("selenium.webdriver.remote.webdriver")
+    def _(m):
+        SeleniumGetWrappingContext(m.WebDriver.get).wrap()
+        SeleniumQuitWrappingContext(m.WebDriver.quit).wrap()
+        SeleniumQuitWrappingContext(m.WebDriver.close).wrap()
+
+    selenium._datadog_patch = True
+
+
+def unpatch() -> None:
+    import selenium
+    from selenium.webdriver.remote.webdriver import WebDriver
+
+    if not getattr(selenium, "_datadog_patch", False):
+        return
+
+    SeleniumGetWrappingContext.extract(WebDriver.get).unwrap()
+    SeleniumQuitWrappingContext.extract(WebDriver.quit).unwrap()
+    SeleniumQuitWrappingContext.extract(WebDriver.close).unwrap()
+
+    selenium._datadog_patch = False

ddtrace/contrib/selenium/__init__.py

@@ -0,0 +1,37 @@
+"""
+The Selenium integration enriches Test Visibility data with extra tags and, if available,
+Real User Monitoring session replays.
+
+Enabling
+~~~~~~~~
+
+The Selenium integration is enabled by default in test contexts (eg: pytest, or unittest). Use
+:func:`patch()<ddtrace.patch>` to enable the integration::
+
+    from ddtrace import patch
+    patch(selenium=True)
+
+
+When using pytest, the `--ddtrace-patch-all` flag is required in order for this integration to
+be enabled.
+
+Configuration
+~~~~~~~~~~~~~
+
+The Selenium integration can be configured using the following options:
+
+DD_CIVISIBILITY_RUM_FLUSH_WAIT_MILLIS: The time in milliseconds to wait after flushing the RUM session.
+"""
+from ...internal.utils.importlib import require_modules
+
+
+required_modules = ["selenium"]
+
+with require_modules(required_modules) as missing_modules:
+    if not missing_modules:
+        # Expose public methods
+        from ..internal.selenium.patch import get_version
+        from ..internal.selenium.patch import patch
+        from ..internal.selenium.patch import unpatch
+
+        __all__ = ["get_version", "patch", "unpatch"]

ddtrace/internal/ci_visibility/api/_base.py

@@ -183,6 +183,9 @@ def _start_span(self) -> None:
             span_type=SpanTypes.TEST,
             activate=True,
         )
+        # Setting initial tags is necessary for integrations that might look at the span before it is finished
+        self._span.set_tag(EVENT_TYPE, self._event_type)
+        self._span.set_tag(SPAN_KIND, "test")
         log.debug("Started span %s for item %s", self._span, self)
 
     @_require_span
@@ -219,8 +222,6 @@ def _set_default_tags(self) -> None:
 
         self.set_tags(
             {
-                EVENT_TYPE: self._event_type,
-                SPAN_KIND: "test",
                 COMPONENT: self._session_settings.test_framework,
                 test.FRAMEWORK: self._session_settings.test_framework,
                 test.FRAMEWORK_VERSION: self._session_settings.test_framework_version,

ddtrace/internal/ci_visibility/api/_test.py

@@ -5,6 +5,7 @@
 from typing import Optional
 from typing import Union
 
+from ddtrace.ext import SpanTypes
 from ddtrace.ext import test
 from ddtrace.ext.test_visibility import ITR_SKIPPING_LEVEL
 from ddtrace.ext.test_visibility._item_ids import TestId
@@ -21,8 +22,8 @@
 from ddtrace.internal.ci_visibility.constants import TEST_IS_NEW
 from ddtrace.internal.ci_visibility.constants import TEST_IS_RETRY
 from ddtrace.internal.ci_visibility.telemetry.constants import EVENT_TYPES
-from ddtrace.internal.ci_visibility.telemetry.events import record_event_created
-from ddtrace.internal.ci_visibility.telemetry.events import record_event_finished
+from ddtrace.internal.ci_visibility.telemetry.events import record_event_created_test
+from ddtrace.internal.ci_visibility.telemetry.events import record_event_finished_test
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.test_visibility._efd_mixins import EFDTestStatus
 from ddtrace.internal.test_visibility._internal_item_ids import InternalTestId
@@ -119,20 +120,20 @@ def _set_span_tags(self) -> None:
             self._span.set_exc_info(self._exc_info.exc_type, self._exc_info.exc_value, self._exc_info.exc_traceback)
 
     def _telemetry_record_event_created(self):
-        record_event_created(
-            event_type=self._event_type_metric_name,
+        record_event_created_test(
             test_framework=self._session_settings.test_framework_metric_name,
-            is_benchmark=self._is_benchmark if self._is_benchmark is not None else None,
+            is_benchmark=self._is_benchmark,
         )
 
     def _telemetry_record_event_finished(self):
-        record_event_finished(
-            event_type=self._event_type_metric_name,
+        record_event_finished_test(
             test_framework=self._session_settings.test_framework_metric_name,
-            is_benchmark=self._is_benchmark if self._is_benchmark is not None else None,
-            is_new=self._is_new if self._is_new is not None else None,
+            is_benchmark=self._is_benchmark,
+            is_new=self.is_new(),
             is_retry=self._efd_is_retry or self._atr_is_retry,
             early_flake_detection_abort_reason=self._efd_abort_reason,
+            is_rum=self._is_rum(),
+            browser_driver=self._get_browser_driver(),
         )
 
     def finish_test(
@@ -143,6 +144,9 @@ def finish_test(
         override_finish_time: Optional[float] = None,
     ) -> None:
         log.debug("Test Visibility: finishing %s, with status: %s, reason: %s", self, status, reason)
+
+        self.set_tag(test.TYPE, SpanTypes.TEST)
+
         if status is not None:
             self.set_status(status)
         if reason is not None:
@@ -379,3 +383,16 @@ def atr_get_final_status(self) -> TestStatus:
             return TestStatus.PASS
 
         return TestStatus.FAIL
+
+    #
+    # Selenium / RUM functionality
+    #
+    def _is_rum(self):
+        if self._span is None:
+            return False
+        return self._span.get_tag("is_rum_active") == "true"
+
+    def _get_browser_driver(self):
+        if self._span is None:
+            return None
+        return self._span.get_tag("test.browser.driver")