|
39 | 39 | from opentelemetry.semconv.trace import SpanAttributes
|
40 | 40 | from opentelemetry.test.mock_textmap import MockTextMapPropagator
|
41 | 41 | from opentelemetry.test.test_base import TestBase
|
| 42 | +from opentelemetry.trace.span import Span, format_span_id, format_trace_id |
| 43 | +from opentelemetry.propagators.aws.aws_xray_propagator import ( |
| 44 | + TRACE_HEADER_KEY, |
| 45 | +) |
42 | 46 |
|
43 | 47 | _REQUEST_ID_REGEX_MATCH = r"[A-Z0-9]{52}"
|
44 | 48 |
|
@@ -225,27 +229,21 @@ def test_unpatch(self):
|
225 | 229 | @mock_ec2
|
226 | 230 | def test_uninstrument_does_not_inject_headers(self):
|
227 | 231 | headers = {}
|
228 |
| - previous_propagator = get_global_textmap() |
229 |
| - try: |
230 |
| - set_global_textmap(MockTextMapPropagator()) |
231 | 232 |
|
232 |
| - def intercept_headers(**kwargs): |
233 |
| - headers.update(kwargs["request"].headers) |
| 233 | + def intercept_headers(**kwargs): |
| 234 | + headers.update(kwargs["request"].headers) |
234 | 235 |
|
235 |
| - ec2 = self._make_client("ec2") |
| 236 | + ec2 = self._make_client("ec2") |
236 | 237 |
|
237 |
| - BotocoreInstrumentor().uninstrument() |
| 238 | + BotocoreInstrumentor().uninstrument() |
238 | 239 |
|
239 |
| - ec2.meta.events.register_first( |
240 |
| - "before-send.ec2.DescribeInstances", intercept_headers |
241 |
| - ) |
242 |
| - with self.tracer_provider.get_tracer("test").start_span("parent"): |
243 |
| - ec2.describe_instances() |
| 240 | + ec2.meta.events.register_first( |
| 241 | + "before-send.ec2.DescribeInstances", intercept_headers |
| 242 | + ) |
| 243 | + with self.tracer_provider.get_tracer("test").start_span("parent"): |
| 244 | + ec2.describe_instances() |
244 | 245 |
|
245 |
| - self.assertNotIn(MockTextMapPropagator.TRACE_ID_KEY, headers) |
246 |
| - self.assertNotIn(MockTextMapPropagator.SPAN_ID_KEY, headers) |
247 |
| - finally: |
248 |
| - set_global_textmap(previous_propagator) |
| 246 | + self.assertNotIn(TRACE_HEADER_KEY, headers) |
249 | 247 |
|
250 | 248 | @mock_sqs
|
251 | 249 | def test_double_patch(self):
|
@@ -306,20 +304,47 @@ def check_headers(**kwargs):
|
306 | 304 | "EC2", "DescribeInstances", request_id=request_id
|
307 | 305 | )
|
308 | 306 |
|
309 |
| - self.assertIn(MockTextMapPropagator.TRACE_ID_KEY, headers) |
310 |
| - self.assertEqual( |
311 |
| - str(span.get_span_context().trace_id), |
312 |
| - headers[MockTextMapPropagator.TRACE_ID_KEY], |
| 307 | + # only x-ray propagation is used in HTTP requests |
| 308 | + self.assertIn(TRACE_HEADER_KEY, headers) |
| 309 | + xray_context = headers[TRACE_HEADER_KEY] |
| 310 | + formated_trace_id = format_trace_id( |
| 311 | + span.get_span_context().trace_id |
313 | 312 | )
|
314 |
| - self.assertIn(MockTextMapPropagator.SPAN_ID_KEY, headers) |
315 |
| - self.assertEqual( |
316 |
| - str(span.get_span_context().span_id), |
317 |
| - headers[MockTextMapPropagator.SPAN_ID_KEY], |
| 313 | + formated_trace_id = ( |
| 314 | + formated_trace_id[:8] + "-" + formated_trace_id[8:] |
318 | 315 | )
|
319 | 316 |
|
| 317 | + self.assertEqual( |
| 318 | + xray_context.lower(), |
| 319 | + f"root=1-{formated_trace_id};parent={format_span_id(span.get_span_context().span_id)};sampled=1".lower(), |
| 320 | + ) |
320 | 321 | finally:
|
321 | 322 | set_global_textmap(previous_propagator)
|
322 | 323 |
|
| 324 | + @mock_ec2 |
| 325 | + def test_override_xray_propagator_injects_into_request(self): |
| 326 | + headers = {} |
| 327 | + |
| 328 | + def check_headers(**kwargs): |
| 329 | + nonlocal headers |
| 330 | + headers = kwargs["request"].headers |
| 331 | + |
| 332 | + BotocoreInstrumentor().instrument() |
| 333 | + |
| 334 | + ec2 = self._make_client("ec2") |
| 335 | + ec2.meta.events.register_first( |
| 336 | + "before-send.ec2.DescribeInstances", check_headers |
| 337 | + ) |
| 338 | + ec2.describe_instances() |
| 339 | + |
| 340 | + request_id = "fdcdcab1-ae5c-489e-9c33-4637c5dda355" |
| 341 | + span = self.assert_span( |
| 342 | + "EC2", "DescribeInstances", request_id=request_id |
| 343 | + ) |
| 344 | + |
| 345 | + self.assertNotIn(MockTextMapPropagator.TRACE_ID_KEY, headers) |
| 346 | + self.assertNotIn(MockTextMapPropagator.SPAN_ID_KEY, headers) |
| 347 | + |
323 | 348 | @mock_xray
|
324 | 349 | def test_suppress_instrumentation_xray_client(self):
|
325 | 350 | xray_client = self._make_client("xray")
|
|
0 commit comments