Supporting the emerging standard TRACEPARENT variable. (#21)

open-telemetry/opentelemetry-specification#740

Closes #20

Author: chrisguidry

src/pytest_opentelemetry/instrumentation.py

@@ -42,8 +42,10 @@ def session_name(self, name):
     @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)
+            return propagate.extract({'traceparent': trace_parent})
+
+        if trace_parent := os.environ.get('TRACEPARENT'):
+            return propagate.extract({'traceparent': trace_parent})
 
         return None
 

tests/test_sessions.py

@@ -1,3 +1,7 @@
+import os
+from contextlib import contextmanager
+from typing import Dict, Generator, Optional
+
 from _pytest.pytester import Pytester
 from opentelemetry import trace
 
@@ -9,6 +13,40 @@
 from . import SpanRecorder
 
 
+@contextmanager
+def environment(**overrides: Optional[str]) -> Generator[None, None, None]:
+    original: Dict[str, Optional[str]] = {}
+    for key, value in overrides.items():
+        original[key] = os.environ.pop(key, None)
+        if value is not None:
+            os.environ[key] = value
+
+    try:
+        yield
+    finally:
+        for key, value in original.items():
+            if value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = value
+
+
+def test_environment_manipulation():
+    assert 'VALUE' not in os.environ
+    with environment(VALUE='outer'):
+        assert os.environ['VALUE'] == 'outer'
+        with environment(VALUE='inner'):
+            assert os.environ['VALUE'] == 'inner'
+            with environment(VALUE=None):
+                assert 'VALUE' not in os.environ
+                with environment(VALUE='once more'):
+                    assert os.environ['VALUE'] == 'once more'
+                assert 'VALUE' not in os.environ
+            assert os.environ['VALUE'] == 'inner'
+        assert os.environ['VALUE'] == 'outer'
+    assert 'VALUE' not in os.environ
+
+
 def test_getting_no_trace_id(pytester: Pytester) -> None:
     config = pytester.parseconfig()
     context = OpenTelemetryPlugin.get_trace_parent(config)
@@ -31,6 +69,24 @@ def test_getting_trace_id_from_command_line(pytester: Pytester) -> None:
     assert parent.span_id == 0xFEDCBA0987654321
 
 
+def test_getting_trace_id_from_environment_variable(pytester: Pytester) -> None:
+    config = pytester.parseconfig()
+
+    with environment(
+        TRACEPARENT='00-1234567890abcdef1234567890abcdef-fedcba0987654321-01'
+    ):
+        context = OpenTelemetryPlugin.get_trace_parent(config)
+
+    assert context
+
+    parent_span = next(iter(context.values()))
+    assert isinstance(parent_span, trace.Span)
+
+    parent = parent_span.get_span_context()
+    assert parent.trace_id == 0x1234567890ABCDEF1234567890ABCDEF
+    assert parent.span_id == 0xFEDCBA0987654321
+
+
 def test_getting_trace_id_from_worker_input(pytester: Pytester) -> None:
     config = pytester.parseconfig()
     setattr(