update unit tests after merging in baggage processor support changes

Author: yiyuan-he

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

@@ -362,24 +362,21 @@ def test_customize_span_exporter_with_agent_observability(self):
     def test_customize_span_processors_with_agent_observability(self):
         mock_tracer_provider: TracerProvider = MagicMock()
 
-        # Test that BaggageSpanProcessor is not added when agent observability is disabled
         os.environ.pop("AGENT_OBSERVABILITY_ENABLED", None)
         _customize_span_processors(mock_tracer_provider, Resource.get_empty())
         self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 0)
 
-        # Reset mock for next test
         mock_tracer_provider.reset_mock()
 
-        # Test that BaggageSpanProcessor is added when agent observability is enabled
         os.environ["AGENT_OBSERVABILITY_ENABLED"] = "true"
         _customize_span_processors(mock_tracer_provider, Resource.get_empty())
-        self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 1)
+        self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 2)
 
-        # Verify the added processor is BaggageSpanProcessor
-        added_processor = mock_tracer_provider.add_span_processor.call_args_list[0].args[0]
-        self.assertIsInstance(added_processor, BaggageSpanProcessor)
+        first_processor = mock_tracer_provider.add_span_processor.call_args_list[0].args[0]
+        self.assertIsInstance(first_processor, BatchUnsampledSpanProcessor)
+        second_processor = mock_tracer_provider.add_span_processor.call_args_list[1].args[0]
+        self.assertIsInstance(second_processor, BaggageSpanProcessor)
 
-        # Clean up
         os.environ.pop("AGENT_OBSERVABILITY_ENABLED", None)
 
     def test_baggage_span_processor_session_id_filtering(self):
@@ -660,18 +657,15 @@ def capture_exporter(*args, **kwargs):
 
     def test_customize_span_processors(self):
         mock_tracer_provider: TracerProvider = MagicMock()
-        # Clean up environment to ensure consistent test state
         os.environ.pop("AGENT_OBSERVABILITY_ENABLED", None)
         os.environ.pop("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", None)
         os.environ.pop("OTEL_AWS_APPLICATION_SIGNALS_RUNTIME_ENABLED", None)
 
         _customize_span_processors(mock_tracer_provider, Resource.get_empty())
         self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 0)
 
-        # Reset mock for next test
         mock_tracer_provider.reset_mock()
 
-        # Test application signals only
         os.environ.setdefault("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", "True")
         os.environ.setdefault("OTEL_AWS_APPLICATION_SIGNALS_RUNTIME_ENABLED", "False")
         _customize_span_processors(mock_tracer_provider, Resource.get_empty())
@@ -681,19 +675,17 @@ def test_customize_span_processors(self):
         second_processor: SpanProcessor = mock_tracer_provider.add_span_processor.call_args_list[1].args[0]
         self.assertIsInstance(second_processor, AwsSpanMetricsProcessor)
 
-        # Reset mock for next test
         mock_tracer_provider.reset_mock()
 
-        # Test both agent observability and application signals enabled
         os.environ.setdefault("AGENT_OBSERVABILITY_ENABLED", "true")
         _customize_span_processors(mock_tracer_provider, Resource.get_empty())
-        self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 3)
+        self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 4)
 
-        # Verify processors are added in the expected order
         processors = [call.args[0] for call in mock_tracer_provider.add_span_processor.call_args_list]
-        self.assertIsInstance(processors[0], BaggageSpanProcessor)  # Agent observability processor added first
-        self.assertIsInstance(processors[1], AttributePropagatingSpanProcessor)  # Application signals processors
-        self.assertIsInstance(processors[2], AwsSpanMetricsProcessor)
+        self.assertIsInstance(processors[0], BatchUnsampledSpanProcessor)
+        self.assertIsInstance(processors[1], BaggageSpanProcessor)
+        self.assertIsInstance(processors[2], AttributePropagatingSpanProcessor)
+        self.assertIsInstance(processors[3], AwsSpanMetricsProcessor)
 
     def test_customize_span_processors_lambda(self):
         mock_tracer_provider: TracerProvider = MagicMock()
@@ -800,19 +792,18 @@ def test_export_unsampled_span_for_lambda(self):
     def test_export_unsampled_span_for_agent_observability(self):
         mock_tracer_provider: TracerProvider = MagicMock()
 
-        # Test when agent observability is disabled (default)
         _export_unsampled_span_for_agent_observability(mock_tracer_provider, Resource.get_empty())
         self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 0)
 
-        # Test when agent observability is enabled with AWS endpoint (the default case)
+        mock_tracer_provider.reset_mock()
+
         os.environ["AGENT_OBSERVABILITY_ENABLED"] = "true"
         os.environ["OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"] = "https://xray.us-east-1.amazonaws.com/v1/traces"
         _export_unsampled_span_for_agent_observability(mock_tracer_provider, Resource.get_empty())
         self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 1)
         processor: SpanProcessor = mock_tracer_provider.add_span_processor.call_args_list[0].args[0]
         self.assertIsInstance(processor, BatchUnsampledSpanProcessor)
 
-        # Clean up
         os.environ.pop("AGENT_OBSERVABILITY_ENABLED", None)
         os.environ.pop("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", None)