|
| 1 | +import os |
| 2 | +import time |
| 3 | +import typing as t |
| 4 | + |
| 5 | +from wrapt.importer import when_imported |
| 6 | + |
| 7 | +from ddtrace import config |
| 8 | +from ddtrace.internal.logger import get_logger |
| 9 | +from ddtrace.internal.wrapping.context import WrappingContext |
| 10 | +import ddtrace.tracer |
| 11 | + |
| 12 | + |
| 13 | +if t.TYPE_CHECKING: |
| 14 | + import selenium.webdriver.remote.webdriver |
| 15 | + |
| 16 | +log = get_logger(__name__) |
| 17 | + |
| 18 | +T = t.TypeVar("T") |
| 19 | + |
| 20 | +_RUM_STOP_SESSION_SCRIPT = """ |
| 21 | +if (window.DD_RUM && window.DD_RUM.stopSession) { |
| 22 | + window.DD_RUM.stopSession(); |
| 23 | + return true; |
| 24 | +} else { |
| 25 | + return false; |
| 26 | +} |
| 27 | +""" |
| 28 | + |
| 29 | +_DEFAULT_FLUSH_SLEEP_MS = 500 |
| 30 | + |
| 31 | + |
| 32 | +def _get_flush_sleep_ms() -> int: |
| 33 | + env_flush_sleep_ms = os.getenv("DD_CIVISIBILITY_RUM_FLUSH_WAIT_MILLIS") |
| 34 | + if env_flush_sleep_ms is None: |
| 35 | + return _DEFAULT_FLUSH_SLEEP_MS |
| 36 | + |
| 37 | + try: |
| 38 | + return int(env_flush_sleep_ms) |
| 39 | + except Exception: # noqa E722 |
| 40 | + log.warning( |
| 41 | + "Could not convert DD_CIVISIBILITY_RUM_FLUSH_WAIT_MILLIS value %s to int, using default: %s", |
| 42 | + env_flush_sleep_ms, |
| 43 | + _DEFAULT_FLUSH_SLEEP_MS, |
| 44 | + ) |
| 45 | + return _DEFAULT_FLUSH_SLEEP_MS |
| 46 | + |
| 47 | + |
| 48 | +config._add( |
| 49 | + "selenium", |
| 50 | + dict(flush_sleep_ms=_get_flush_sleep_ms()), |
| 51 | +) |
| 52 | + |
| 53 | + |
| 54 | +class SeleniumWrappingContextBase(WrappingContext): |
| 55 | + def _handle_enter(self) -> None: |
| 56 | + pass |
| 57 | + |
| 58 | + def _handle_return(self) -> None: |
| 59 | + pass |
| 60 | + |
| 61 | + def _get_webdriver_instance(self) -> "selenium.webdriver.remote.webdriver.WebDriver": |
| 62 | + try: |
| 63 | + return self.get_local("self") |
| 64 | + except KeyError: |
| 65 | + log.debug("Could not get Selenium WebDriver instance") |
| 66 | + return None |
| 67 | + |
| 68 | + def __enter__(self) -> "SeleniumWrappingContextBase": |
| 69 | + super().__enter__() |
| 70 | + |
| 71 | + try: |
| 72 | + self._handle_enter() |
| 73 | + except Exception: # noqa: E722 |
| 74 | + log.debug("Error handling selenium instrumentation enter", exc_info=True) |
| 75 | + |
| 76 | + return self |
| 77 | + |
| 78 | + def __return__(self, value: T) -> T: |
| 79 | + """Always return the original value no matter what our instrumentation does""" |
| 80 | + try: |
| 81 | + self._handle_return() |
| 82 | + except Exception: # noqa: E722 |
| 83 | + log.debug("Error handling instrumentation return", exc_info=True) |
| 84 | + |
| 85 | + return value |
| 86 | + |
| 87 | + |
| 88 | +class SeleniumGetWrappingContext(SeleniumWrappingContextBase): |
| 89 | + def _handle_return(self) -> None: |
| 90 | + root_span = ddtrace.tracer.current_root_span() |
| 91 | + test_trace_id = root_span.trace_id |
| 92 | + |
| 93 | + if root_span is None or root_span.get_tag("type") != "test": |
| 94 | + return |
| 95 | + |
| 96 | + webdriver_instance = self._get_webdriver_instance() |
| 97 | + |
| 98 | + if webdriver_instance is None: |
| 99 | + return |
| 100 | + |
| 101 | + # The trace IDs for Test Visibility data using the CIVisibility protocol are 64-bit |
| 102 | + # TODO[ci_visibility]: properly identify whether to use 64 or 128 bit trace_ids |
| 103 | + trace_id_64bit = test_trace_id % 2**64 |
| 104 | + |
| 105 | + webdriver_instance.add_cookie({"name": "datadog-ci-visibility-test-execution-id", "value": str(trace_id_64bit)}) |
| 106 | + |
| 107 | + root_span.set_tag("test.is_browser", "true") |
| 108 | + root_span.set_tag("test.browser.driver", "selenium") |
| 109 | + root_span.set_tag("test.browser.driver_version", get_version()) |
| 110 | + |
| 111 | + # Submit empty values for browser names or version if multiple are found |
| 112 | + browser_name = webdriver_instance.capabilities.get("browserName") |
| 113 | + browser_version = webdriver_instance.capabilities.get("browserVersion") |
| 114 | + |
| 115 | + existing_browser_name = root_span.get_tag("test.browser.name") |
| 116 | + if existing_browser_name is None: |
| 117 | + root_span.set_tag("test.browser.name", browser_name) |
| 118 | + elif existing_browser_name not in ["", browser_name]: |
| 119 | + root_span.set_tag("test.browser.name", "") |
| 120 | + |
| 121 | + existing_browser_version = root_span.get_tag("test.browser.version") |
| 122 | + if existing_browser_version is None: |
| 123 | + root_span.set_tag("test.browser.version", browser_version) |
| 124 | + elif existing_browser_version not in ["", browser_version]: |
| 125 | + root_span.set_tag("test.browser.version", "") |
| 126 | + |
| 127 | + |
| 128 | +class SeleniumQuitWrappingContext(SeleniumWrappingContextBase): |
| 129 | + def _handle_enter(self) -> None: |
| 130 | + root_span = ddtrace.tracer.current_root_span() |
| 131 | + |
| 132 | + if root_span is None or root_span.get_tag("type") != "test": |
| 133 | + return |
| 134 | + |
| 135 | + webdriver_instance = self._get_webdriver_instance() |
| 136 | + |
| 137 | + if webdriver_instance is None: |
| 138 | + return |
| 139 | + |
| 140 | + is_rum_active = webdriver_instance.execute_script(_RUM_STOP_SESSION_SCRIPT) |
| 141 | + time.sleep(config.selenium.flush_sleep_ms / 1000) |
| 142 | + |
| 143 | + if is_rum_active: |
| 144 | + root_span.set_tag("test.is_rum_active", "true") |
| 145 | + |
| 146 | + webdriver_instance.delete_cookie("datadog-ci-visibility-test-execution-id") |
| 147 | + |
| 148 | + |
| 149 | +def get_version() -> str: |
| 150 | + import selenium |
| 151 | + |
| 152 | + try: |
| 153 | + return selenium.__version__ |
| 154 | + except AttributeError: |
| 155 | + log.debug("Could not get Selenium version") |
| 156 | + return "" |
| 157 | + |
| 158 | + |
| 159 | +def patch() -> None: |
| 160 | + import selenium |
| 161 | + |
| 162 | + if getattr(selenium, "_datadog_patch", False): |
| 163 | + return |
| 164 | + |
| 165 | + @when_imported("selenium.webdriver.remote.webdriver") |
| 166 | + def _(m): |
| 167 | + SeleniumGetWrappingContext(m.WebDriver.get).wrap() |
| 168 | + SeleniumQuitWrappingContext(m.WebDriver.quit).wrap() |
| 169 | + SeleniumQuitWrappingContext(m.WebDriver.close).wrap() |
| 170 | + |
| 171 | + selenium._datadog_patch = True |
| 172 | + |
| 173 | + |
| 174 | +def unpatch() -> None: |
| 175 | + import selenium |
| 176 | + from selenium.webdriver.remote.webdriver import WebDriver |
| 177 | + |
| 178 | + if not getattr(selenium, "_datadog_patch", False): |
| 179 | + return |
| 180 | + |
| 181 | + SeleniumGetWrappingContext.extract(WebDriver.get).unwrap() |
| 182 | + SeleniumQuitWrappingContext.extract(WebDriver.quit).unwrap() |
| 183 | + SeleniumQuitWrappingContext.extract(WebDriver.close).unwrap() |
| 184 | + |
| 185 | + selenium._datadog_patch = False |
0 commit comments