fix unit test coverage

Author: wangzlei

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

@@ -1,3 +1,5 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
 import logging
 from typing import Optional
 
@@ -13,6 +15,7 @@
 
 class BatchUnsampledSpanProcessor(BaseBatchSpanProcessor):
 
+    # pylint: disable=no-self-use
     def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
         if not span.context.trace_flags.sampled:
             span.set_attribute(AWS_TRACE_FLAG_UNSAMPLED, "True")
@@ -26,6 +29,7 @@ def on_end(self, span: ReadableSpan) -> None:
             return
 
         if len(self.queue) == self.max_queue_size:
+            # pylint: disable=access-member-before-definition
             if not self._spans_dropped:
                 logger.warning("Queue is full, likely spans will be dropped.")
                 self._spans_dropped = True

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

@@ -154,8 +154,8 @@ def _init_tracing(
         span_exporter: SpanExporter = exporter_class(**exporter_args)
         span_exporter = _customize_exporter(span_exporter, resource)
         trace_provider.add_span_processor(BatchSpanProcessor(span_exporter))
-        _export_unsampled_span_for_lambda(trace_provider, resource)
 
+    _export_unsampled_span_for_lambda(trace_provider, resource)
     _customize_span_processors(trace_provider, resource)
 
     set_tracer_provider(trace_provider)

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

@@ -1,3 +1,5 @@
+# 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
 

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

@@ -7,6 +7,7 @@
 
 from amazon.opentelemetry.distro.always_record_sampler import AlwaysRecordSampler
 from amazon.opentelemetry.distro.attribute_propagating_span_processor import AttributePropagatingSpanProcessor
+from amazon.opentelemetry.distro.aws_batch_unsampled_span_processor import BatchUnsampledSpanProcessor
 from amazon.opentelemetry.distro.aws_metric_attributes_span_exporter import AwsMetricAttributesSpanExporter
 from amazon.opentelemetry.distro.aws_opentelemetry_configurator import (
     ApplicationSignalsExporterProvider,
@@ -15,6 +16,7 @@
     _customize_exporter,
     _customize_sampler,
     _customize_span_processors,
+    _export_unsampled_span_for_lambda,
     _is_application_signals_enabled,
     _is_defer_to_workers_enabled,
     _is_wsgi_master_process,
@@ -351,6 +353,19 @@ def test_initialize_components_called_when_deferred_disabled(self, mock_initiali
         mock_initialize_components.assert_called_once()
         os.environ.pop("IS_WSGI_MASTER_PROCESS_ALREADY_SEEN", None)
 
+    def test_export_unsampled_span_for_lambda(self):
+        mock_tracer_provider: TracerProvider = MagicMock()
+        _export_unsampled_span_for_lambda(mock_tracer_provider, Resource.get_empty())
+        self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 0)
+
+        os.environ.setdefault("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", "True")
+        os.environ.setdefault("AWS_LAMBDA_FUNCTION_NAME", "myfunction")
+        _export_unsampled_span_for_lambda(mock_tracer_provider, Resource.get_empty())
+        self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 1)
+        first_processor: SpanProcessor = mock_tracer_provider.add_span_processor.call_args_list[0].args[0]
+        self.assertIsInstance(first_processor, BatchUnsampledSpanProcessor)
+        os.environ.pop("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", None)
+
 
 def validate_distro_environ():
     tc: TestCase = TestCase()