| 
10 | 10 | 
 
  | 
11 | 11 | from requests import Session  | 
12 | 12 | 
 
  | 
 | 13 | +from amazon.opentelemetry.distro._aws_attribute_keys import AWS_AI_AGENT_TYPE, AWS_LOCAL_SERVICE  | 
13 | 14 | from amazon.opentelemetry.distro.always_record_sampler import AlwaysRecordSampler  | 
14 | 15 | from amazon.opentelemetry.distro.attribute_propagating_span_processor import AttributePropagatingSpanProcessor  | 
15 | 16 | from amazon.opentelemetry.distro.aws_batch_unsampled_span_processor import BatchUnsampledSpanProcessor  | 
 | 
27 | 28 |     _custom_import_sampler,  | 
28 | 29 |     _customize_logs_exporter,  | 
29 | 30 |     _customize_metric_exporters,  | 
 | 31 | +    _customize_resource,  | 
30 | 32 |     _customize_sampler,  | 
31 | 33 |     _customize_span_exporter,  | 
32 | 34 |     _customize_span_processors,  | 
 | 
67 | 69 | from opentelemetry.sdk.trace import Span, SpanProcessor, Tracer, TracerProvider  | 
68 | 70 | from opentelemetry.sdk.trace.export import SpanExporter  | 
69 | 71 | from opentelemetry.sdk.trace.sampling import DEFAULT_ON, Sampler  | 
 | 72 | +from opentelemetry.semconv.resource import ResourceAttributes  | 
70 | 73 | from opentelemetry.trace import get_tracer_provider  | 
71 | 74 | 
 
  | 
72 | 75 | 
 
  | 
@@ -1091,6 +1094,53 @@ def test_customize_metric_exporters_with_emf(self):  | 
1091 | 1094 |             self.assertEqual(len(metric_readers), 1)  | 
1092 | 1095 |             self.assertIsInstance(metric_readers[0], PeriodicExportingMetricReader)  | 
1093 | 1096 | 
 
  | 
 | 1097 | +    @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.is_agent_observability_enabled")  | 
 | 1098 | +    @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.get_service_attribute")  | 
 | 1099 | +    def test_customize_resource_without_agent_observability(self, mock_get_service_attribute, mock_is_agent_enabled):  | 
 | 1100 | +        """Test _customize_resource when agent observability is disabled"""  | 
 | 1101 | +        mock_is_agent_enabled.return_value = False  | 
 | 1102 | +        mock_get_service_attribute.return_value = ("test-service", False)  | 
 | 1103 | + | 
 | 1104 | +        resource = Resource.create({ResourceAttributes.SERVICE_NAME: "test-service"})  | 
 | 1105 | +        result = _customize_resource(resource)  | 
 | 1106 | + | 
 | 1107 | +        # Should only have AWS_LOCAL_SERVICE added  | 
 | 1108 | +        self.assertEqual(result.attributes[AWS_LOCAL_SERVICE], "test-service")  | 
 | 1109 | +        self.assertNotIn(AWS_AI_AGENT_TYPE, result.attributes)  | 
 | 1110 | + | 
 | 1111 | +    @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.is_agent_observability_enabled")  | 
 | 1112 | +    @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.get_service_attribute")  | 
 | 1113 | +    def test_customize_resource_with_agent_observability_default(  | 
 | 1114 | +        self, mock_get_service_attribute, mock_is_agent_enabled  | 
 | 1115 | +    ):  | 
 | 1116 | +        """Test _customize_resource when agent observability is enabled with default agent type"""  | 
 | 1117 | +        mock_is_agent_enabled.return_value = True  | 
 | 1118 | +        mock_get_service_attribute.return_value = ("test-service", False)  | 
 | 1119 | + | 
 | 1120 | +        resource = Resource.create({ResourceAttributes.SERVICE_NAME: "test-service"})  | 
 | 1121 | +        result = _customize_resource(resource)  | 
 | 1122 | + | 
 | 1123 | +        # Should have both AWS_LOCAL_SERVICE and AWS_AI_AGENT_TYPE with default value  | 
 | 1124 | +        self.assertEqual(result.attributes[AWS_LOCAL_SERVICE], "test-service")  | 
 | 1125 | +        self.assertEqual(result.attributes[AWS_AI_AGENT_TYPE], "default")  | 
 | 1126 | + | 
 | 1127 | +    @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.is_agent_observability_enabled")  | 
 | 1128 | +    @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.get_service_attribute")  | 
 | 1129 | +    def test_customize_resource_with_existing_agent_type(self, mock_get_service_attribute, mock_is_agent_enabled):  | 
 | 1130 | +        """Test _customize_resource when agent type already exists in resource"""  | 
 | 1131 | +        mock_is_agent_enabled.return_value = True  | 
 | 1132 | +        mock_get_service_attribute.return_value = ("test-service", False)  | 
 | 1133 | + | 
 | 1134 | +        # Create resource with existing agent type  | 
 | 1135 | +        resource = Resource.create(  | 
 | 1136 | +            {ResourceAttributes.SERVICE_NAME: "test-service", AWS_AI_AGENT_TYPE: "existing-agent"}  | 
 | 1137 | +        )  | 
 | 1138 | +        result = _customize_resource(resource)  | 
 | 1139 | + | 
 | 1140 | +        # Should preserve existing agent type and not override it  | 
 | 1141 | +        self.assertEqual(result.attributes[AWS_LOCAL_SERVICE], "test-service")  | 
 | 1142 | +        self.assertEqual(result.attributes[AWS_AI_AGENT_TYPE], "existing-agent")  | 
 | 1143 | + | 
1094 | 1144 | 
 
  | 
1095 | 1145 | def validate_distro_environ():  | 
1096 | 1146 |     tc: TestCase = TestCase()  | 
 | 
0 commit comments