Add interface for communication with ddtestpy

ddtrace/contrib/internal/pytest/ddtestpy_interface.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from pathlib import Path
+import typing as t
+
+from ddtestpy.ddtrace_interface import CoverageData
+from ddtestpy.ddtrace_interface import PushSpanProtocol
+from ddtestpy.ddtrace_interface import TraceContext
+from ddtestpy.ddtrace_interface import TracerInterface
+from ddtestpy.ddtrace_interface import register_tracer_interface
+
+import ddtrace
+from ddtrace.internal.coverage.code import ModuleCodeCollector
+from ddtrace.internal.coverage.installer import install
+from ddtrace.trace import Span
+from ddtrace.trace import TraceFilter
+
+
+class TraceForwarder(TraceFilter):
+    def __init__(self, push_span: PushSpanProtocol) -> None:
+        self.push_span = push_span
+
+    def process_trace(self, trace: t.List[Span]) -> t.Optional[t.List[Span]]:
+        for span in trace:
+            self.push_span(
+                trace_id=span.trace_id,
+                parent_id=span.parent_id,
+                span_id=span.span_id,
+                service=span.service,
+                resource=span.resource,
+                name=span.name,
+                error=span.error,
+                start_ns=span.start_ns,
+                duration_ns=span.duration_ns,
+                meta=span.get_tags(),
+                metrics=span.get_metrics(),
+                span_type=span.span_type,
+            )
+        return None
+
+
+class DDTraceInterface(TracerInterface):
+    def __init__(self) -> None:
+        self.workspace_path: t.Optional[Path] = None
+        self._should_enable_test_optimization: bool = False
+        self._should_enable_trace_collection: bool = False
+
+    def should_enable_test_optimization(self) -> bool:
+        return self._should_enable_test_optimization
+
+    def should_enable_trace_collection(self) -> bool:
+        return self._should_enable_trace_collection
+
+    def enable_trace_collection(self, push_span: PushSpanProtocol) -> None:
+        ddtrace.tracer.configure(trace_processors=[TraceForwarder(push_span)])
+
+    def disable_trace_collection(self) -> None:
+        ddtrace.tracer.configure(trace_processors=[])
+
+    @contextmanager
+    def trace_context(self, resource: str) -> t.Generator[TraceContext, None, None]:
+        import ddtrace
+
+        # TODO: check if this breaks async tests.
+        # This seems to be necessary because buggy ddtrace integrations can leave spans
+        # unfinished, and spans for subsequent tests will have the wrong parent.
+        ddtrace.tracer.context_provider.activate(None)
+
+        with ddtrace.tracer.trace(resource) as root_span:
+            yield TraceContext(root_span.trace_id, root_span.span_id)
+
+    def enable_coverage_collection(self, workspace_path: Path) -> None:
+        self.workspace_path = workspace_path
+        install(include_paths=[workspace_path], collect_import_time_coverage=True)
+        ModuleCodeCollector.start_coverage()
+
+    def disable_coverage_collection(self) -> None:
+        ModuleCodeCollector.stop_coverage()
+
+    @contextmanager
+    def coverage_context(self) -> t.Generator[CoverageData, None, None]:
+        with ModuleCodeCollector.CollectInContext() as coverage_collector:
+            coverage_data = CoverageData()
+            yield coverage_data
+            covered_lines_by_file = coverage_collector.get_covered_lines()
+
+        coverage_data.bitmaps = {}
+        for absolute_path, covered_lines in covered_lines_by_file.items():
+            try:
+                relative_path = Path(absolute_path).relative_to(self.workspace_path)
+            except ValueError:
+                continue  # covered file does not belong to current repo
+
+            path_str = f"/{relative_path}"
+            coverage_data.bitmaps[path_str] = covered_lines.to_bytes()
+
+
+ddtrace_interface = DDTraceInterface()
+
+register_tracer_interface(ddtrace_interface)

ddtrace/contrib/internal/pytest/entry_point.py

@@ -0,0 +1,91 @@
+import os
+import typing as t
+
+import pytest
+
+from ddtrace.contrib.internal.pytest.ddtestpy_interface import ddtrace_interface
+
+
+if os.getenv("DD_USE_DDTESTPY", "false").lower() in ("true", "1"):
+
+    def _is_option_true(option: str, early_config: pytest.Config, args: t.List[str]) -> bool:
+        """
+        Check whether a command line option is true.
+        """
+        return early_config.getoption(option) or early_config.getini(option) or f"--{option}" in args
+
+    def _is_enabled_early(early_config: pytest.Config, args: t.List[str]) -> bool:
+        """
+        Check whether ddtrace is enabled via command line during early initialization.
+        """
+        if _is_option_true("no-ddtrace", early_config, args):
+            return False
+
+        return _is_option_true("ddtrace", early_config, args)
+
+    @pytest.hookimpl(tryfirst=True, hookwrapper=True)
+    def pytest_load_initial_conftests(
+        early_config: pytest.Config, parser: pytest.Parser, args: t.List[str]
+    ) -> t.Generator[None, None, None]:
+        """
+        Ensure we enable ddtrace early enough so that coverage collection is enabled by the time we collect tests.
+        """
+        if _is_enabled_early(early_config, args):
+            ddtrace_interface._should_enable_test_optimization = True
+
+        yield
+
+    def pytest_addoption(parser: pytest.Parser) -> None:
+        """Add ddtrace options."""
+        group = parser.getgroup("ddtrace")
+
+        group.addoption(
+            "--ddtrace",
+            action="store_true",
+            dest="ddtrace",
+            default=False,
+            help="Enable Datadog Test Optimization with tracer features",
+        )
+        group.addoption(
+            "--no-ddtrace",
+            action="store_true",
+            dest="no-ddtrace",
+            default=False,
+            help="Disable Datadog Test Optimization with tracer features (overrides --ddtrace)",
+        )
+        group.addoption(
+            "--ddtrace-patch-all",
+            action="store_true",
+            dest="ddtrace-patch-all",
+            default=False,
+            help="Enable all tracer integrations during tests",
+        )
+        group.addoption(
+            "--ddtrace-iast-fail-tests",
+            action="store_true",
+            dest="ddtrace-iast-fail-tests",
+            default=False,
+            help="When IAST is enabled, fail tests that have detected vulnerabilities",
+        )
+
+        parser.addini("ddtrace", "Enable Datadog Test Optimization with tracer features", type="bool")
+        parser.addini(
+            "no-ddtrace", "Disable Datadog Test Optimization with tracer features (overrides 'ddtrace')", type="bool"
+        )
+        parser.addini("ddtrace-patch-all", "Enable all tracer integrations during tests", type="bool")
+
+    def pytest_configure(config: pytest.Config) -> None:
+        yes_ddtrace = config.getoption("ddtrace") or config.getini("ddtrace")
+        no_ddtrace = config.getoption("no-ddtrace") or config.getini("no-ddtrace")
+
+        patch_all = config.getoption("ddtrace-patch-all") or config.getini("ddtrace-patch-all")
+
+        if yes_ddtrace and not no_ddtrace:
+            ddtrace_interface._should_enable_test_optimization = True
+
+        if patch_all:
+            ddtrace_interface._should_enable_trace_collection = True
+
+else:
+    # If DD_USE_DDTESTPY is not enabled, behave like the old pytest plugin.
+    from ddtrace.contrib.internal.pytest.plugin import *  # noqa: F403

pyproject.toml

@@ -63,9 +63,7 @@ _dd_crashtracker_receiver = "ddtrace.commands._dd_crashtracker_receiver:main"
 ddcontextvars_context = "ddtrace.internal.opentelemetry.context:DDRuntimeContext"
 
 [project.entry-points.pytest11]
-ddtrace = "ddtrace.contrib.internal.pytest.plugin"
-"ddtrace.pytest_bdd" = "ddtrace.contrib.internal.pytest_bdd.plugin"
-"ddtrace.pytest_benchmark" = "ddtrace.contrib.internal.pytest_benchmark.plugin"
+ddtrace = "ddtrace.contrib.internal.pytest.entry_point"
 
 [project.entry-points.'ddtrace.products']
 "apm-tracing-rc" = "ddtrace.internal.remoteconfig.products.apm_tracing"