crank/crates/crank-observability/tests/propagation.rs

use axum::http::{HeaderMap, HeaderValue};
use crank_observability::{inject_current_trace_context, set_remote_trace_parent};
use opentelemetry::{global, trace::TracerProvider as _};
use opentelemetry_sdk::{propagation::TraceContextPropagator, trace::SdkTracerProvider};
use tracing::info_span;
use tracing_subscriber::layer::SubscriberExt;

const REMOTE_TRACE_ID: &str = "0af7651916cd43dd8448eb211c80319c";

#[test]
fn valid_remote_parent_is_continued_and_request_id_is_unrelated() {
    with_trace_dispatch(|| {
        let mut incoming = HeaderMap::new();
        incoming.insert(
            "traceparent",
            HeaderValue::from_static("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"),
        );
        incoming.insert("x-request-id", HeaderValue::from_static("req-unrelated"));

        let span = info_span!("http.request", request_id = "req-unrelated");
        assert!(set_remote_trace_parent(&span, &incoming));
        let _guard = span.enter();

        let mut outgoing = HeaderMap::new();
        assert!(inject_current_trace_context(&mut outgoing));
        let propagated = outgoing["traceparent"].to_str().unwrap();
        assert_eq!(&propagated[3..35], REMOTE_TRACE_ID);
        assert!(!propagated.contains("req-unrelated"));
        assert!(!outgoing.contains_key("baggage"));
    });
}

#[test]
fn invalid_parent_is_ignored_and_a_new_trace_is_created() {
    with_trace_dispatch(|| {
        let mut incoming = HeaderMap::new();
        incoming.insert(
            "traceparent",
            HeaderValue::from_static("canary-invalid-traceparent"),
        );

        let span = info_span!("mcp.request");
        assert!(!set_remote_trace_parent(&span, &incoming));
        let _guard = span.enter();

        let mut outgoing = HeaderMap::new();
        assert!(inject_current_trace_context(&mut outgoing));
        let propagated = outgoing["traceparent"].to_str().unwrap();
        assert!(propagated.starts_with("00-"));
        assert!(!propagated.contains("canary-invalid-traceparent"));
    });
}

fn with_trace_dispatch(test: impl FnOnce()) {
    global::set_text_map_propagator(TraceContextPropagator::new());
    let provider = SdkTracerProvider::builder().build();
    let tracer = provider.tracer("propagation-test");
    let subscriber =
        tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer));
    let dispatch = tracing::Dispatch::new(subscriber);
    let _guard = tracing::dispatcher::set_default(&dispatch);

    test();
    provider.shutdown().unwrap();
}