change reflection comment

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

@@ -1,8 +1,10 @@
 # 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
+
+from opentelemetry.sdk.trace import ReadableSpan
 
-from amazon.opentelemetry.distro._aws_span_processing_util import get_ingress_operation
 from amazon.opentelemetry.distro.attribute_propagating_span_processor import AttributePropagatingSpanProcessor
 from amazon.opentelemetry.distro.attribute_propagating_span_processor_builder import (
     AttributePropagatingSpanProcessorBuilder,
@@ -13,9 +15,13 @@ class TestAttributePropagatingSpanProcessorBuilder(TestCase):
     def test_basic(self):
         builder: AttributePropagatingSpanProcessorBuilder = AttributePropagatingSpanProcessorBuilder()
         self.assertIs(builder.set_propagation_data_key("test"), builder)
-        self.assertIs(builder.set_propagation_data_extractor(get_ingress_operation), builder)
+
+        def mock_extractor(_: ReadableSpan) -> str:
+            return "test"
+
+        self.assertIs(builder.set_propagation_data_extractor(mock_extractor), builder)
         self.assertIs(builder.set_attributes_keys_to_propagate(["test"]), builder)
         span_processor: AttributePropagatingSpanProcessor = builder.build()
         self.assertIs(span_processor._propagation_data_key, "test")
-        self.assertIs(span_processor._propagation_data_extractor, get_ingress_operation)
+        self.assertEqual(span_processor._propagation_data_extractor(MagicMock()), "test")
         self.assertEqual(span_processor._attribute_keys_to_propagate, tuple(["test"]))

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

@@ -1,6 +1,7 @@
 # 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
 
 from amazon.opentelemetry.distro._aws_metric_attribute_generator import _AwsMetricAttributeGenerator
 from amazon.opentelemetry.distro.aws_metric_attributes_span_exporter import AwsMetricAttributesSpanExporter
@@ -11,8 +12,9 @@
 
 class TestAwsMetricAttributesSpanExporterBuilder(TestCase):
     def test_basic(self):
-        generator: _AwsMetricAttributeGenerator = _AwsMetricAttributeGenerator()
+        generator: _AwsMetricAttributeGenerator = MagicMock()
+        generator.test_key = "test"
         builder: AwsMetricAttributesSpanExporterBuilder = AwsMetricAttributesSpanExporterBuilder(None, None)
         self.assertIs(builder.set_generator(generator), builder)
         exporter: AwsMetricAttributesSpanExporter = builder.build()
-        self.assertIs(exporter._generator, generator)
+        self.assertIs(exporter._generator.test_key, "test")