WIP

open-telemetry · Jun 10, 2024 · 34ab8b3 · 34ab8b3
1 parent b6be46b
commit 34ab8b3
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 62 deletions.
diff --git a/...pentelemetry-propagator-aws-xray/src/opentelemetry/propagators/aws/aws_xray_propagator.py b/...pentelemetry-propagator-aws-xray/src/opentelemetry/propagators/aws/aws_xray_propagator.py
@@ -58,6 +58,7 @@
 
 import logging
 import typing
+from os import environ
 
 from opentelemetry import trace
 from opentelemetry.context import Context
@@ -69,7 +70,6 @@
     default_getter,
     default_setter,
 )
-from os import environ
 
 TRACE_HEADER_KEY = "X-Amzn-Trace-Id"
 AWS_TRACE_HEADER_PROP = "com.amazonaws.xray.traceHeader"
@@ -351,9 +351,8 @@ def extract(
         if trace.get_current_span(context=context).get_span_context().is_valid:
             return xray_context
 
-        trace_header = (
-            environ.get(AWS_TRACE_HEADER_PROP) or
-            environ.get(AWS_TRACE_HEADER_ENV_KEY)
+        trace_header = environ.get(AWS_TRACE_HEADER_PROP) or environ.get(
+            AWS_TRACE_HEADER_ENV_KEY
         )
 
         if trace_header is None:
@@ -362,7 +361,7 @@ def extract(
         result = super().extract(
             {TRACE_HEADER_KEY: trace_header},
             context=xray_context,
-            getter=getter
+            getter=getter,
         )
 
         return result
diff --git a/propagator/opentelemetry-propagator-aws-xray/tests/test_aws_xray_lambda_propagator.py b/propagator/opentelemetry-propagator-aws-xray/tests/test_aws_xray_lambda_propagator.py
@@ -12,38 +12,32 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from os import environ
 from unittest import TestCase
-from requests.structures import CaseInsensitiveDict
 from unittest.mock import patch
-from os import environ
-from opentelemetry.sdk.trace import ReadableSpan
-from opentelemetry.propagators.textmap import (
-    DefaultGetter,
-)
-from opentelemetry.trace import (
-    Link,
-    get_current_span,
-    TraceState
-)
-from opentelemetry.context import get_current
 
+from requests.structures import CaseInsensitiveDict
+
+from opentelemetry.context import get_current
 from opentelemetry.propagators.aws.aws_xray_propagator import (
+    TRACE_HEADER_KEY,
     AwsXrayLambdaPropagator,
-    TRACE_HEADER_KEY
 )
+from opentelemetry.propagators.textmap import DefaultGetter
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.trace import Link, TraceState, get_current_span
 
 
 class AwsXRayLambdaPropagatorTest(TestCase):
 
     @patch.dict(
         environ,
         {
-            "_X_AMZN_TRACE_ID":
-            (
+            "_X_AMZN_TRACE_ID": (
                 "Root=1-00000001-d188f8fa79d48a391a778fa6;"
                 "Parent=53995c3f42cd8ad8;Sampled=1;Foo=Bar"
             )
-        }
+        },
     )
     def test_extract_from_environment_variable(self):
         propagator = AwsXrayLambdaPropagator()
@@ -52,37 +46,27 @@ def test_extract_from_environment_variable(self):
 
         actual_context = get_current_span(
             propagator.extract(
-                {},
-                context=get_current(),
-                getter=default_getter
+                {}, context=get_current(), getter=default_getter
             )
         ).get_span_context()
 
         self.assertEqual(
-            hex(actual_context.trace_id),
-            "0x1d188f8fa79d48a391a778fa6"
-        )
-        self.assertEqual(
-            hex(actual_context.span_id),
-            "0x53995c3f42cd8ad8"
+            hex(actual_context.trace_id), "0x1d188f8fa79d48a391a778fa6"
         )
+        self.assertEqual(hex(actual_context.span_id), "0x53995c3f42cd8ad8")
         self.assertTrue(
             actual_context.trace_flags.sampled,
         )
-        self.assertEqual(
-            actual_context.trace_state,
-            TraceState.get_default()
-        )
+        self.assertEqual(actual_context.trace_state, TraceState.get_default())
 
     @patch.dict(
         environ,
         {
-            "_X_AMZN_TRACE_ID":
-            (
+            "_X_AMZN_TRACE_ID": (
                 "Root=1-00000002-240000000000000000000002;"
                 "Parent=1600000000000002;Sampled=1;Foo=Bar"
             )
-        }
+        },
     )
     def test_add_link_from_environment_variable(self):
 
@@ -92,24 +76,19 @@ def test_add_link_from_environment_variable(self):
 
         carrier = CaseInsensitiveDict(
             {
-                TRACE_HEADER_KEY:
-                (
+                TRACE_HEADER_KEY: (
                     "Root=1-00000001-240000000000000000000001;"
                     "Parent=1600000000000001;Sampled=1"
                 )
             }
         )
 
         extracted_context = propagator.extract(
-            carrier,
-            context=get_current(),
-            getter=default_getter
+            carrier, context=get_current(), getter=default_getter
         )
 
         link_context = propagator.extract(
-            carrier,
-            context=extracted_context,
-            getter=default_getter
+            carrier, context=extracted_context, getter=default_getter
         )
 
         span = ReadableSpan(
@@ -119,37 +98,29 @@ def test_add_link_from_environment_variable(self):
         span_parent_context = get_current_span(span.parent).get_span_context()
 
         self.assertEqual(
-            hex(span_parent_context.trace_id),
-            "0x2240000000000000000000002"
+            hex(span_parent_context.trace_id), "0x2240000000000000000000002"
         )
         self.assertEqual(
-            hex(span_parent_context.span_id),
-            "0x1600000000000002"
+            hex(span_parent_context.span_id), "0x1600000000000002"
         )
         self.assertTrue(
             span_parent_context.trace_flags.sampled,
         )
         self.assertEqual(
-            span_parent_context.trace_state,
-            TraceState.get_default()
+            span_parent_context.trace_state, TraceState.get_default()
         )
 
-        span_link_context = (
-            get_current_span(span.links[0].context).get_span_context()
-        )
+        span_link_context = get_current_span(
+            span.links[0].context
+        ).get_span_context()
 
         self.assertEqual(
-            hex(span_link_context.trace_id),
-            "0x1240000000000000000000001"
-        )
-        self.assertEqual(
-            hex(span_link_context.span_id),
-            "0x1600000000000001"
+            hex(span_link_context.trace_id), "0x1240000000000000000000001"
         )
+        self.assertEqual(hex(span_link_context.span_id), "0x1600000000000001")
         self.assertTrue(
             span_link_context.trace_flags.sampled,
         )
         self.assertEqual(
-            span_link_context.trace_state,
-            TraceState.get_default()
+            span_link_context.trace_state, TraceState.get_default()
         )