Skip to content

Commit

Permalink
Support for running with xdist disabled or not installed (#10)
Browse files Browse the repository at this point in the history
Fixes #9
  • Loading branch information
chrisguidry authored Sep 6, 2022
1 parent a7a700a commit 0c02220
Show file tree
Hide file tree
Showing 6 changed files with 77 additions and 24 deletions.
3 changes: 3 additions & 0 deletions .flake8
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
[flake8]
max-line-length = 88
extend-ignore = E203
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -6,3 +6,4 @@ dist
__pycache__
.python-version
*.egg-info
.vscode
2 changes: 1 addition & 1 deletion setup.cfg
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[metadata]
name = pytest-opentelemetry
version = 0.3.2
version = 0.4.0
author = Chris Guidry
author_email = chris@theguidrys.us
description = A pytest plugin for instrumenting test runs via OpenTelemetry
Expand Down
53 changes: 36 additions & 17 deletions src/pytest_opentelemetry/instrumentation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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()
Expand Down Expand Up @@ -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)
10 changes: 8 additions & 2 deletions src/pytest_opentelemetry/plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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())
32 changes: 28 additions & 4 deletions tests/test_sessions.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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()))
Expand Down Expand Up @@ -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(
Expand All @@ -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)

0 comments on commit 0c02220

Please sign in to comment.