games-development-ai/cheap-worker/tests/test_trace_propagation.py

"""test_trace_propagation.py — 便宜档面四三跳 W3C traceparent 传播单测(阶段四观测)。

守的不变量(把一次便宜档生成串成 Java→worker→cheap-service 单条 trace 树的三个断点):
  · 断点①提取:propagator 从构造的 traceparent header 提取出正确的 trace_id/span_id(is_remote)。
  · 断点①建 span:OTLP 开时 worker_span_scope 在入站 context 下开 cheap.worker.generate SERVER span 当子,
    业务 traceId 并存(§10 决策5)作属性;OTLP 关时只透传入站 context(不依赖开关)。
  · 断点②注入:worker_span_scope 内 current_traceparent_carrier() 注入出的 traceparent 与入站同 trace_id(往返一致)。
  · 断点③挂靠:CheapOtlpSink(parent_carrier=...) 让生成 span 挂到入站 context 下当子 span(不再私有 root);
    缺省 parent_carrier=None 时保持波③ 私有 root 行为(默认关字节不变)。
  · best-effort(设计 §8):坏 carrier 不抛;生成异常原样穿出 worker_span_scope(worker loop 据此发兜底 failed)。

只本地跑、零网络:用 InMemorySpanExporter / 手工 traceparent 构造断言,不发真 Collector。
跑:cheap-worker/.venv/bin/python -m pytest cheap-worker/tests/test_trace_propagation.py -q
"""
import os
import sys
from pathlib import Path

# sys.path:cheap-worker/ + tier2/gen-worker/(observability / service / worker 包;与 test_cheap_otlp_sink 同范式)。
_HERE = os.path.dirname(os.path.abspath(__file__))
_CW = os.path.dirname(_HERE)
_REPO_ROOT = os.path.dirname(_CW)
_GEN_WORKER = os.path.join(_REPO_ROOT, "tier2", "gen-worker")
for _p in (_CW, _GEN_WORKER):
    if _p not in sys.path:
        sys.path.insert(0, _p)

import pytest  # noqa: E402
import cheap_otlp_sink as S  # noqa: E402
from opentelemetry import trace as _otel_trace  # noqa: E402
from opentelemetry.trace import SpanKind  # noqa: E402
from opentelemetry.sdk.trace.export.in_memory_span_exporter import (  # noqa: E402
    InMemorySpanExporter,
)

# 固定构造的入站 W3C 值(模拟 Java OTel agent 出网注入的 traceparent)。
_TRACE_HEX = "11112222333344445555666677778888"
_SPAN_HEX = "1111222233334444"
_TP = f"00-{_TRACE_HEX}-{_SPAN_HEX}-01"


def _span_ctx(ctx):
    """从 extract 出的 Context 取当前 span 的 SpanContext。"""
    return _otel_trace.get_current_span(ctx).get_span_context()


# ── 断点①提取:propagator 从 traceparent header 提取正确 trace_id/span_id ─────────────
def test_extract_from_traceparent_header_yields_correct_ids():
    ctx = S._extract_remote_context({"traceparent": _TP})
    assert ctx is not None
    sc = _span_ctx(ctx)
    assert format(sc.trace_id, "032x") == _TRACE_HEX
    assert format(sc.span_id, "016x") == _SPAN_HEX
    assert sc.is_remote is True  # 入站(远端)context


def test_extract_none_or_missing_traceparent_returns_none():
    assert S._extract_remote_context(None) is None
    assert S._extract_remote_context({}) is None
    assert S._extract_remote_context({"tracestate": "x=1"}) is None  # 无 traceparent 键


# ── 断点②注入:worker_span_scope(OTLP 关)内注入的 traceparent 与入站同 trace_id(往返一致)──
def test_inject_roundtrip_trace_id_consistent_when_otlp_off(monkeypatch):
    monkeypatch.setattr(S, "resolve_otlp_endpoint", lambda explicit=None: None)  # 强制 OTLP 关
    S.reset_shared_tracer_for_test()
    with S.worker_span_scope({"traceparent": _TP}, conversation_id="w1", business_trace_id="tp1"):
        out = S.current_traceparent_carrier()
    assert "traceparent" in out
    parts = out["traceparent"].split("-")
    assert parts[1] == _TRACE_HEX, "出站 trace_id 必与入站一致(往返一致)"
    # OTLP 关:无 worker span,当前 = 入站 remote context → span_id 段透传入站 Java span(直挂 Java 下)。
    assert parts[2] == _SPAN_HEX


def test_current_carrier_empty_when_no_context():
    """无任何当前 context → 注入空 carrier(出站少带 header,不抛)。"""
    S.reset_shared_tracer_for_test()
    out = S.current_traceparent_carrier()
    assert out == {} or "traceparent" not in out


# ── 断点①建 span:OTLP 开 → worker span 在入站 context 下当子 + 业务 traceId 并存 ────────
def test_worker_span_scope_creates_child_server_span_when_otlp_on(monkeypatch):
    S.reset_shared_tracer_for_test()
    mem = InMemorySpanExporter()
    monkeypatch.setattr(S, "resolve_otlp_endpoint", lambda explicit=None: "http://collector.test:4318")
    # 预置进程级共享 tracer 用内存 exporter(worker_span_scope 复用同一单例,无需真网络)。
    S._get_shared_tracer("http://collector.test:4318", span_exporter=mem)

    with S.worker_span_scope({"traceparent": _TP}, conversation_id="w1", business_trace_id="tp1"):
        inj = S.current_traceparent_carrier()  # 作用域内当前 = worker span

    spans = mem.get_finished_spans()
    assert len(spans) == 1
    ws = spans[0]
    assert ws.name == "cheap.worker.generate"
    assert ws.kind == SpanKind.SERVER
    # worker span 在入站 trace 下(同 trace_id),父 = 入站 Java span(不新造 root)。
    assert format(ws.context.trace_id, "032x") == _TRACE_HEX
    assert ws.parent is not None and format(ws.parent.span_id, "016x") == _SPAN_HEX
    # 断点②:作用域内注入的 traceparent = worker span 自身(同 trace,span_id = worker span)。
    p = inj["traceparent"].split("-")
    assert p[1] == _TRACE_HEX and p[2] == format(ws.context.span_id, "016x")
    # §10 决策5 业务 traceId 并存:gameId 作 conversation.id、业务 traceId 另存属性(都不改格式)。
    attrs = dict(ws.attributes)
    assert attrs.get("gen_ai.conversation.id") == "w1"
    assert attrs.get("agentscope.cheap.business_trace_id") == "tp1"
    S.reset_shared_tracer_for_test()


# ── 断点③挂靠:CheapOtlpSink(parent_carrier) 让生成 span 挂到入站 context 下当子 span ────
def test_cheap_service_gen_span_parents_under_inbound_carrier():
    S.reset_shared_tracer_for_test()
    mem = InMemorySpanExporter()
    sink = S.make_cheap_otlp_sink(trace_id="conv-x", endpoint="http://collector.test:4318",
                                  span_exporter=mem, parent_carrier={"traceparent": _TP})
    sink({"traceId": "conv-x", "step": 0, "cost": None, "verdict": None,
          "timestamp": "2026-07-05T10:00:00", "ext": {"raw": {"event": "ReplyStartEvent"}}})
    spans = mem.get_finished_spans()
    assert len(spans) == 1
    sp = spans[0]
    # 生成 span 挂在入站 trace 下(同 trace_id、父 = 入站 span);不再私有 root。
    assert format(sp.context.trace_id, "032x") == _TRACE_HEX
    assert sp.parent is not None and format(sp.parent.span_id, "016x") == _SPAN_HEX
    # 业务 traceId 并存:conversation.id 仍是业务 traceId(承担产物归位/回调关联),不被 W3C 贯穿覆盖。
    assert dict(sp.attributes).get("gen_ai.conversation.id") == "conv-x"
    S.reset_shared_tracer_for_test()


def test_cheap_service_gen_span_is_root_when_no_parent_carrier():
    """缺省 parent_carrier=None(未开传播)→ 生成 span 仍私有 root(波③ 行为字节不变)。"""
    S.reset_shared_tracer_for_test()
    mem = InMemorySpanExporter()
    sink = S.make_cheap_otlp_sink(trace_id="conv-y", endpoint="http://collector.test:4318",
                                  span_exporter=mem)  # parent_carrier 缺省 None
    sink({"traceId": "conv-y", "step": 0, "ext": {"raw": {"event": "ReplyStartEvent"}}})
    sp = mem.get_finished_spans()[0]
    assert sp.parent is None, "未开传播时保持私有 root(默认关字节不变)"
    S.reset_shared_tracer_for_test()


def test_build_trace_sink_threads_parent_carrier_to_child_span():
    """build_trace_sink(parent_carrier=...) 开时:fan-out 的 OTLP 侧生成 span 也挂到入站 context 下当子。"""
    S.reset_shared_tracer_for_test()
    mem = InMemorySpanExporter()
    seen = []
    fan = S.build_trace_sink(lambda step: seen.append(step), trace_id="g-9",
                             endpoint="http://collector.test:4318", span_exporter=mem,
                             parent_carrier={"traceparent": _TP})
    fan({"traceId": "g-9", "step": 0, "ext": {"raw": {"event": "ReplyStartEvent"}}})
    assert len(seen) == 1  # jsonl 侧照落不被吞
    sp = mem.get_finished_spans()[0]
    assert format(sp.context.trace_id, "032x") == _TRACE_HEX  # OTLP 侧挂到入站 trace 下
    assert sp.parent is not None and format(sp.parent.span_id, "016x") == _SPAN_HEX
    S.reset_shared_tracer_for_test()


# ── 三跳串成单条 trace 树(端到端:入站 Java → worker span → cheap-service 生成 span)──────
def test_three_hops_form_single_trace_tree(monkeypatch):
    """把三个断点串起来:worker span 挂在入站 Java span 下,生成 span 又挂在 worker span 下 —— 同一 trace_id、
    父子链完整(正是 dev 真机在 Tempo 里要看到的 Java→worker→cheap-service 单条 trace 树)。"""
    S.reset_shared_tracer_for_test()
    mem = InMemorySpanExporter()
    monkeypatch.setattr(S, "resolve_otlp_endpoint", lambda explicit=None: "http://collector.test:4318")
    S._get_shared_tracer("http://collector.test:4318", span_exporter=mem)

    # 跳①:worker 线程在入站 Java context 下开 worker span,取其出站 carrier(= driver 会注入 header/写 sidecar 的)。
    with S.worker_span_scope({"traceparent": _TP}, conversation_id="w1", business_trace_id="tp1"):
        worker_carrier = S.current_traceparent_carrier()
    # 跳③:cheap-service 端拿 worker 的 carrier 当 parent(经 sidecar 桥),生成 span 挂到 worker span 下。
    gen_sink = S.make_cheap_otlp_sink(trace_id="w1", endpoint="http://collector.test:4318",
                                      span_exporter=mem, parent_carrier=worker_carrier)
    gen_sink({"traceId": "w1", "step": 0, "ext": {"raw": {"event": "ReplyStartEvent"}}})

    spans = {sp.name: sp for sp in mem.get_finished_spans()}
    worker_sp = spans["cheap.worker.generate"]
    gen_sp = spans["step 0 · ReplyStartEvent"]
    # 三段同一 trace_id(单条 trace)。
    assert format(worker_sp.context.trace_id, "032x") == _TRACE_HEX
    assert format(gen_sp.context.trace_id, "032x") == _TRACE_HEX
    # 父子链:worker span 父 = 入站 Java span;生成 span 父 = worker span。
    assert format(worker_sp.parent.span_id, "016x") == _SPAN_HEX
    assert gen_sp.parent.span_id == worker_sp.context.span_id
    S.reset_shared_tracer_for_test()


# ── best-effort 铁律(设计 §8):坏 carrier 不抛;生成异常原样穿出 ────────────────────────
def test_worker_span_scope_bad_carrier_does_not_raise(monkeypatch):
    monkeypatch.setattr(S, "resolve_otlp_endpoint", lambda explicit=None: None)
    with S.worker_span_scope({"traceparent": "garbage-not-w3c"}, conversation_id="w"):
        pass  # 坏 traceparent 只导致无父,不抛


def test_worker_span_scope_does_not_swallow_body_exception(monkeypatch):
    """生成异常必须原样穿出 worker_span_scope —— worker loop 据此发兜底 failed 回调,绝不被观测吞掉。"""
    monkeypatch.setattr(S, "resolve_otlp_endpoint", lambda explicit=None: None)
    with pytest.raises(ValueError):
        with S.worker_span_scope({"traceparent": _TP}):
            raise ValueError("boom")


def test_worker_span_scope_body_exception_ends_span_when_otlp_on(monkeypatch):
    """OTLP 开时生成异常:worker span 仍被 end(不漏 span)、异常照常穿出。"""
    S.reset_shared_tracer_for_test()
    mem = InMemorySpanExporter()
    monkeypatch.setattr(S, "resolve_otlp_endpoint", lambda explicit=None: "http://collector.test:4318")
    S._get_shared_tracer("http://collector.test:4318", span_exporter=mem)
    with pytest.raises(RuntimeError):
        with S.worker_span_scope({"traceparent": _TP}, conversation_id="w1"):
            raise RuntimeError("gen crashed")
    spans = mem.get_finished_spans()
    assert len(spans) == 1 and spans[0].name == "cheap.worker.generate"  # 异常路 span 照样收尾导出
    S.reset_shared_tracer_for_test()