Support for running with xdist disabled or not installed (#10)

Fixes #9

Author: chrisguidry

.flake8

@@ -0,0 +1,3 @@
+[flake8]
+max-line-length = 88
+extend-ignore = E203

.gitignore

@@ -6,3 +6,4 @@ dist
 __pycache__
 .python-version
 *.egg-info
+.vscode

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = pytest-opentelemetry
-version = 0.3.2
+version = 0.4.0
 author = Chris Guidry
 author_email = [email protected]
 description = A pytest plugin for instrumenting test runs via OpenTelemetry

src/pytest_opentelemetry/instrumentation.py

@@ -17,32 +17,23 @@
 
 from .resource import CodebaseResourceDetector
 
-try:
-    from xdist.workermanage import WorkerController  # pylint: disable=unused-import
-except ImportError:  # pragma: no cover
-    WorkerController = None
-
 tracer = trace.get_tracer('pytest-opentelemetry')
 
 
 class OpenTelemetryPlugin:
     """A pytest plugin which produces OpenTelemetry spans around test sessions and
     individual test runs."""
 
-    @staticmethod
-    def get_trace_parent(config: Config) -> Optional[Context]:
+    @classmethod
+    def get_trace_parent(cls, config: Config) -> Optional[Context]:
         if trace_parent := config.getvalue('--trace-parent'):
             from_arguments = {'traceparent': trace_parent}
             return propagate.extract(from_arguments)
 
-        if workerinput := getattr(config, 'workerinput', None):
-            return propagate.extract(workerinput)
-
         return None
 
     def pytest_configure(self, config: Config) -> None:
         self.trace_parent = self.get_trace_parent(config)
-        self.worker_id = getattr(config, 'workerinput', {}).get('workerid')
 
         # This can't be tested both ways in one process
         if config.getoption('--export-traces'):  # pragma: no cover
@@ -52,13 +43,13 @@ def pytest_configure(self, config: Config) -> None:
         configurator.resource_detectors.append(CodebaseResourceDetector())
         configurator.configure()
 
-    def pytest_sessionstart(self, session: Session) -> None:
-        session_name = f'test worker {self.worker_id}' if self.worker_id else 'test run'
-        self.session_span = tracer.start_span(session_name, context=self.trace_parent)
+    session_name: str = 'test run'
 
-    def pytest_configure_node(self, node: WorkerController) -> None:  # pragma: no cover
-        with trace.use_span(self.session_span, end_on_exit=False):
-            propagate.inject(node.workerinput)
+    def pytest_sessionstart(self, session: Session) -> None:
+        self.session_span = tracer.start_span(
+            self.session_name,
+            context=self.trace_parent,
+        )
 
     def pytest_sessionfinish(self, session: Session) -> None:
         self.session_span.end()
@@ -110,3 +101,31 @@ def pytest_runtest_logreport(report: TestReport) -> None:
 
         status_code = StatusCode.ERROR if report.outcome == 'failed' else StatusCode.OK
         trace.get_current_span().set_status(Status(status_code))
+
+
+try:
+    from xdist.workermanage import WorkerController  # pylint: disable=unused-import
+except ImportError:  # pragma: no cover
+    WorkerController = None
+
+
+class XdistOpenTelemetryPlugin(OpenTelemetryPlugin):
+    """An xdist-aware version of the OpenTelemetryPlugin"""
+
+    @classmethod
+    def get_trace_parent(cls, config: Config) -> Optional[Context]:
+        if workerinput := getattr(config, 'workerinput', None):
+            return propagate.extract(workerinput)
+
+        return super().get_trace_parent(config)
+
+    def pytest_configure(self, config: Config) -> None:
+        super().pytest_configure(config)
+        worker_id = getattr(config, 'workerinput', {}).get('workerid')
+        self.session_name = (
+            f'test worker {worker_id}' if worker_id else self.session_name
+        )
+
+    def pytest_configure_node(self, node: WorkerController) -> None:  # pragma: no cover
+        with trace.use_span(self.session_span, end_on_exit=False):
+            propagate.inject(node.workerinput)

src/pytest_opentelemetry/plugin.py

@@ -30,6 +30,12 @@ def pytest_addoption(parser: Parser) -> None:
 
 def pytest_configure(config: Config) -> None:
     # pylint: disable=import-outside-toplevel
-    from pytest_opentelemetry.instrumentation import OpenTelemetryPlugin
+    from pytest_opentelemetry.instrumentation import (
+        OpenTelemetryPlugin,
+        XdistOpenTelemetryPlugin,
+    )
 
-    config.pluginmanager.register(OpenTelemetryPlugin())
+    if config.pluginmanager.has_plugin("xdist"):
+        config.pluginmanager.register(XdistOpenTelemetryPlugin())
+    else:
+        config.pluginmanager.register(OpenTelemetryPlugin())

tests/test_sessions.py

@@ -1,7 +1,10 @@
 from _pytest.pytester import Pytester
 from opentelemetry import trace
 
-from pytest_opentelemetry.instrumentation import OpenTelemetryPlugin
+from pytest_opentelemetry.instrumentation import (
+    OpenTelemetryPlugin,
+    XdistOpenTelemetryPlugin,
+)
 
 from . import SpanRecorder
 
@@ -35,7 +38,7 @@ def test_getting_trace_id_from_worker_input(pytester: Pytester) -> None:
         'workerinput',
         {'traceparent': '00-1234567890abcdef1234567890abcdef-fedcba0987654321-01'},
     )
-    context = OpenTelemetryPlugin.get_trace_parent(config)
+    context = XdistOpenTelemetryPlugin.get_trace_parent(config)
     assert context
 
     parent_span = next(iter(context.values()))
@@ -86,15 +89,13 @@ def test_one(worker_id):
 
             span = trace.get_current_span()
             assert span.context.trace_id == 0x1234567890abcdef1234567890abcdef
-            assert span.context.span_id != 0xfedcba0987654321
 
         def test_two(worker_id):
             # confirm that this is an xdist worker
             assert worker_id in {'gw0', 'gw1'}
 
             span = trace.get_current_span()
             assert span.context.trace_id == 0x1234567890abcdef1234567890abcdef
-            assert span.context.span_id != 0xfedcba0987654321
     """
     )
     result = pytester.runpytest_subprocess(
@@ -104,3 +105,26 @@ def test_two(worker_id):
         '00-1234567890abcdef1234567890abcdef-fedcba0987654321-01',
     )
     result.assert_outcomes(passed=2)
+
+
+def test_works_without_xdist(pytester: Pytester, span_recorder: SpanRecorder) -> None:
+    pytester.makepyfile(
+        """
+        from opentelemetry import trace
+
+        def test_one():
+            span = trace.get_current_span()
+            assert span.context.trace_id == 0x1234567890abcdef1234567890abcdef
+
+        def test_two():
+            span = trace.get_current_span()
+            assert span.context.trace_id == 0x1234567890abcdef1234567890abcdef
+    """
+    )
+    result = pytester.runpytest_subprocess(
+        '-p',
+        'no:xdist',
+        '--trace-parent',
+        '00-1234567890abcdef1234567890abcdef-fedcba0987654321-01',
+    )
+    result.assert_outcomes(passed=2)