added unit tests

Author: AsakerMohd

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py

@@ -13,6 +13,7 @@
 AWS_SDK_DESCENDANT: str = "aws.sdk.descendant"
 AWS_CONSUMER_PARENT_SPAN_KIND: str = "aws.consumer.parent.span.kind"
 AWS_TRACE_FLAG_SAMPLED: str = "aws.trace.flag.sampled"
+AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER: str = "aws.trace.lambda.multiple-server"
 AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER: str = "aws.remote.resource.cfn.primary.identifier"
 
 # AWS_#_NAME attributes are not supported in python as they are not part of the Semantic Conventions.

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_lambda_span_processor.py

@@ -4,17 +4,14 @@
 
 from typing_extensions import override
 
+from amazon.opentelemetry.distro._aws_attribute_keys import AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER
 from opentelemetry.context import Context
 from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
 from opentelemetry.trace import SpanKind
-from amazon.opentelemetry.distro._aws_attribute_keys import AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER
+
 
 class AwsLambdaSpanProcessor(SpanProcessor):
     def __init__(self, instrumentation_names=None):
-        """
-        :param instrumentation_names: A set or list of instrumentation scope names
-            for which we want to mark as SERVER spans if they are INTERNAL.
-        """
         self.instrumentation_names = set(instrumentation_names or ["opentelemetry.instrumentation.flask"])
         self.parent_lambda_span = None
 
@@ -23,22 +20,25 @@ def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None
         scope = getattr(span, "instrumentation_scope", None)
         if span.kind == SpanKind.SERVER and scope.name == "opentelemetry.instrumentation.aws_lambda":
             self.parent_lambda_span = span
-        
-        if span.kind == SpanKind.INTERNAL and scope.name in self.instrumentation_names:
-            span._kind = SpanKind.SERVER
+
+        if (
+            scope.name in self.instrumentation_names
+            and self.parent_lambda_span.get_span_context().span_id == span.parent.span_id
+        ):
+            if span.kind == SpanKind.INTERNAL:
+                span._kind = SpanKind.SERVER
             self.parent_lambda_span.set_attribute(AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER, True)
         return
 
     @override
     def on_end(self, span: ReadableSpan) -> None:
         return
 
-    @override
-    def force_flush(self, timeout_millis: int = 30000) -> bool:
-        """Flush any buffered data."""
-        return True
-
     @override
     def shutdown(self) -> None:
-        """Clean up."""
         self.force_flush()
+
+    # pylint: disable=no-self-use
+    @override
+    def force_flush(self, timeout_millis: int = None) -> bool:
+        return True

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py

@@ -16,6 +16,7 @@
     AttributePropagatingSpanProcessorBuilder,
 )
 from amazon.opentelemetry.distro.aws_batch_unsampled_span_processor import BatchUnsampledSpanProcessor
+from amazon.opentelemetry.distro.aws_lambda_span_processor import AwsLambdaSpanProcessor
 from amazon.opentelemetry.distro.aws_metric_attributes_span_exporter_builder import (
     AwsMetricAttributesSpanExporterBuilder,
 )
@@ -24,7 +25,6 @@
 from amazon.opentelemetry.distro.otlp_udp_exporter import OTLPUdpSpanExporter
 from amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler import AwsXRayRemoteSampler
 from amazon.opentelemetry.distro.scope_based_exporter import ScopeBasedPeriodicExportingMetricReader
-from amazon.opentelemetry.distro.aws_lambda_span_processor import AwsLambdaSpanProcessor
 from amazon.opentelemetry.distro.scope_based_filtering_view import ScopeBasedRetainingView
 from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter as OTLPHttpOTLPMetricExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_lambda_span_processor.py

@@ -0,0 +1,82 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from unittest import TestCase
+from unittest.mock import MagicMock, patch
+
+from amazon.opentelemetry.distro._aws_attribute_keys import AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER
+from amazon.opentelemetry.distro.aws_lambda_span_processor import AwsLambdaSpanProcessor
+from opentelemetry.trace import Span, SpanContext, SpanKind, TraceFlags
+
+
+class TestAwsLambdaSpanProcessor(TestCase):
+
+    def setUp(self):
+        self.processor = AwsLambdaSpanProcessor()
+        self.lambda_span: Span = MagicMock()
+        self.lambda_span.instrumentation_scope.name = "opentelemetry.instrumentation.aws_lambda"
+        self.lambda_span.kind = SpanKind.SERVER
+
+        self.lambda_span_context: SpanContext = MagicMock()
+        self.lambda_span_context.trace_id = "ABC"
+        self.lambda_span_context.span_id = "lambda_id"
+
+        self.lambda_span.get_span_context.return_value = self.lambda_span_context
+        self.processor.on_start(self.lambda_span)
+
+    def tearDown(self):
+        self.processor.on_end(self.lambda_span)
+        self.processor.shutdown()
+
+    @patch("opentelemetry.sdk.trace.Span")
+    def test_lambda_span_multiple_server_flag_internal_api(self, mock_span_class):
+
+        flask_span = mock_span_class.return_value
+        flask_span.instrumentation_scope.name = "opentelemetry.instrumentation.flask"
+        flask_span.kind = SpanKind.INTERNAL
+        flask_span.parent = self.lambda_span_context
+
+        self.processor.on_start(flask_span)
+
+        self.assertEqual(flask_span._kind, SpanKind.SERVER)
+        self.assertIn(AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER, self.lambda_span.set_attribute.call_args_list[0][0][0])
+
+        self.processor.on_end(flask_span)
+        self.processor.on_end(self.lambda_span)
+
+        self.processor.shutdown()
+
+    @patch("opentelemetry.sdk.trace.Span")
+    def test_lambda_span_multiple_server_flag_server_api(self, mock_span_class):
+
+        flask_span = mock_span_class.return_value
+        flask_span.instrumentation_scope.name = "opentelemetry.instrumentation.flask"
+        flask_span.kind = SpanKind.SERVER
+        flask_span.parent = self.lambda_span_context
+
+        self.processor.on_start(flask_span)
+
+        self.assertEqual(flask_span.kind, SpanKind.SERVER)
+        self.assertIn(AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER, self.lambda_span.set_attribute.call_args_list[0][0][0])
+
+        self.processor.on_end(flask_span)
+        self.processor.on_end(self.lambda_span)
+
+        self.processor.shutdown()
+
+    @patch("opentelemetry.sdk.trace.Span")
+    def test_lambda_span_single_server_span(self, mock_span_class):
+
+        flask_span = mock_span_class.return_value
+        flask_span.instrumentation_scope.name = "opentelemetry.instrumentation.http"
+        flask_span.kind = SpanKind.CLIENT
+        flask_span.parent = self.lambda_span_context
+
+        self.processor.on_start(flask_span, self.lambda_span_context)
+
+        self.assertEqual(flask_span.kind, SpanKind.CLIENT)
+        flask_span.set_attribute.assert_not_called()
+
+        self.processor.on_end(flask_span)
+        self.processor.on_end(self.lambda_span)
+
+        self.processor.shutdown()

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py

@@ -318,7 +318,7 @@ def test_customize_span_processors_lambda(self):
         os.environ.setdefault("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", "True")
         os.environ.setdefault("AWS_LAMBDA_FUNCTION_NAME", "myLambdaFunc")
         _customize_span_processors(mock_tracer_provider, Resource.get_empty())
-        self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 2)
+        self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 3)
         first_processor: SpanProcessor = mock_tracer_provider.add_span_processor.call_args_list[0].args[0]
         self.assertIsInstance(first_processor, AttributePropagatingSpanProcessor)
         second_processor: SpanProcessor = mock_tracer_provider.add_span_processor.call_args_list[1].args[0]