|
2 | 2 | import logging |
3 | 3 |
|
4 | 4 | import pytest |
| 5 | +from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags, use_span |
5 | 6 |
|
6 | 7 | from posthog import AsyncPosthog, new_context, set_context_option, tag |
7 | 8 | from posthog.ai.utils import _capture_ai_event |
@@ -475,6 +476,119 @@ async def test_async_caller_layers_beat_flag_properties(source): |
475 | 476 | _assert_flag_layers(events, source) |
476 | 477 |
|
477 | 478 |
|
| 479 | +SPAN_IDS = { |
| 480 | + "$trace_id": "00000000000000000000000000000123", |
| 481 | + "$span_id": "0000000000000456", |
| 482 | +} |
| 483 | + |
| 484 | +# Each caller layer sets one key of the pair, so the SDK adds neither span id. |
| 485 | +TRACE_ID_CASES = [ |
| 486 | + pytest.param(None, SPAN_IDS, id="no_layer"), |
| 487 | + pytest.param("event", {"$trace_id": "caller"}, id="event_trace_id"), |
| 488 | + pytest.param("context", {"$trace_id": "caller"}, id="context_trace_id"), |
| 489 | + pytest.param("super", {"$span_id": "caller"}, id="super_span_id"), |
| 490 | +] |
| 491 | + |
| 492 | +TRACE_CALLS = { |
| 493 | + "capture": lambda c, p: c.capture("e", distinct_id="u", properties=p), |
| 494 | + "capture_ai": lambda c, p: c.capture_ai( |
| 495 | + "$ai_generation", distinct_id="u", properties=p |
| 496 | + ), |
| 497 | + "capture_exception": lambda c, p: c.capture_exception( |
| 498 | + ValueError("boom"), distinct_id="u", properties=p |
| 499 | + ), |
| 500 | +} |
| 501 | + |
| 502 | + |
| 503 | +def _capture_in_span(method, source=None): |
| 504 | + span = NonRecordingSpan( |
| 505 | + SpanContext( |
| 506 | + trace_id=0x123, |
| 507 | + span_id=0x456, |
| 508 | + is_remote=False, |
| 509 | + trace_flags=TraceFlags.SAMPLED, |
| 510 | + ) |
| 511 | + ) |
| 512 | + |
| 513 | + def call(client): |
| 514 | + with new_context(fresh=True), use_span(span): |
| 515 | + if source == "context": |
| 516 | + tag("$trace_id", "caller") |
| 517 | + properties = {"$trace_id": "caller"} if source == "event" else None |
| 518 | + return TRACE_CALLS[method](client, properties) |
| 519 | + |
| 520 | + return call |
| 521 | + |
| 522 | + |
| 523 | +def _trace_config(source, capture_trace_context=True): |
| 524 | + config: dict = {"capture_trace_context": capture_trace_context} |
| 525 | + if source == "super": |
| 526 | + config["super_properties"] = {"$span_id": "caller"} |
| 527 | + return config |
| 528 | + |
| 529 | + |
| 530 | +def _span_ids(properties): |
| 531 | + return {k: properties[k] for k in SPAN_IDS if k in properties} |
| 532 | + |
| 533 | + |
| 534 | +def _sync_trace_context(method): |
| 535 | + # Sync capture_exception adds the span ids even with capture_trace_context off. |
| 536 | + return method != "capture_exception" |
| 537 | + |
| 538 | + |
| 539 | +@pytest.mark.parametrize("method", list(TRACE_CALLS)) |
| 540 | +@pytest.mark.parametrize("source,expected", TRACE_ID_CASES) |
| 541 | +def test_sync_span_ids_fill_as_a_pair_below_caller_layers(method, source, expected): |
| 542 | + events = _sync_wire_events( |
| 543 | + _capture_in_span(method, source), |
| 544 | + **_trace_config(source, _sync_trace_context(method)), |
| 545 | + ) |
| 546 | + assert _span_ids(events[0]["properties"]) == expected |
| 547 | + |
| 548 | + |
| 549 | +@pytest.mark.asyncio |
| 550 | +@pytest.mark.parametrize("method", list(TRACE_CALLS)) |
| 551 | +@pytest.mark.parametrize("source,expected", TRACE_ID_CASES) |
| 552 | +async def test_async_span_ids_fill_as_a_pair_below_caller_layers( |
| 553 | + method, source, expected |
| 554 | +): |
| 555 | + events = await _async_wire_events( |
| 556 | + _capture_in_span(method, source), **_trace_config(source) |
| 557 | + ) |
| 558 | + assert _span_ids(events[0]["properties"]) == expected |
| 559 | + |
| 560 | + |
| 561 | +def _span_id_hook(seen): |
| 562 | + def hook(msg): |
| 563 | + seen.append(_span_ids(msg["properties"])) |
| 564 | + return {**msg, "properties": {**msg["properties"], "$span_id": "hook"}} |
| 565 | + |
| 566 | + return hook |
| 567 | + |
| 568 | + |
| 569 | +@pytest.mark.parametrize("method", list(TRACE_CALLS)) |
| 570 | +def test_sync_before_send_sees_and_changes_span_ids(method): |
| 571 | + seen: list = [] |
| 572 | + events = _sync_wire_events( |
| 573 | + _capture_in_span(method), |
| 574 | + before_send=_span_id_hook(seen), |
| 575 | + **_trace_config(None, _sync_trace_context(method)), |
| 576 | + ) |
| 577 | + assert seen == [SPAN_IDS] |
| 578 | + assert _span_ids(events[0]["properties"]) == {**SPAN_IDS, "$span_id": "hook"} |
| 579 | + |
| 580 | + |
| 581 | +@pytest.mark.asyncio |
| 582 | +@pytest.mark.parametrize("method", list(TRACE_CALLS)) |
| 583 | +async def test_async_before_send_sees_and_changes_span_ids(method): |
| 584 | + seen: list = [] |
| 585 | + events = await _async_wire_events( |
| 586 | + _capture_in_span(method), before_send=_span_id_hook(seen), **_trace_config(None) |
| 587 | + ) |
| 588 | + assert seen == [SPAN_IDS] |
| 589 | + assert _span_ids(events[0]["properties"]) == {**SPAN_IDS, "$span_id": "hook"} |
| 590 | + |
| 591 | + |
478 | 592 | GEOIP_CALL_CASES = [ |
479 | 593 | pytest.param(True, "context", False, True, id="call_beats_context"), |
480 | 594 | pytest.param(True, "super", False, True, id="call_beats_super"), |
|
0 commit comments