|
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