#![cfg(feature = "observability")]
use axum::{Router, routing::get};
use rs_zero::observability::{
CorrelationContext, MetricsRegistry, current_trace_id, metrics_router, request_id_from_headers,
span_id_from_traceparent, trace_id_from_traceparent, traceparent_from_headers,
};
use rs_zero::rest::{ApiResponse, RestConfig, RestMetricsConfig, RestMiddlewareConfig, RestServer};
use tower::ServiceExt;
#[tokio::test]
async fn http_request_id_is_available_for_trace_correlation() {
let registry = MetricsRegistry::new();
let config = RestConfig {
metrics_registry: Some(registry.clone()),
middlewares: RestMiddlewareConfig {
metrics: RestMetricsConfig { enabled: true },
..RestMiddlewareConfig::default()
},
..RestConfig::default()
};
let app = RestServer::new(
config,
Router::new()
.route("/ready", get(|| async { ApiResponse::success("ok") }))
.merge(metrics_router(registry.clone())),
)
.into_router();
let response = app
.oneshot(
axum::http::Request::builder()
.uri("/ready")
.header("x-request-id", "trace-req-1")
.header(
"traceparent",
"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
)
.body(axum::body::Body::empty())
.expect("request"),
)
.await
.expect("response");
assert_eq!(
response.headers().get("x-request-id").expect("request id"),
"trace-req-1"
);
assert_eq!(
request_id_from_headers(response.headers()).as_deref(),
Some("trace-req-1")
);
assert!(registry.render_prometheus().contains("route=\"/ready\""));
assert!(current_trace_id().is_none());
}
#[tokio::test]
async fn request_id_extension_is_available_to_handlers() {
use axum::extract::Extension;
use rs_zero::observability::CurrentRequestId;
async fn handler(Extension(request_id): Extension<CurrentRequestId>) -> String {
request_id.0
}
let registry = MetricsRegistry::new();
let config = RestConfig {
metrics_registry: Some(registry.clone()),
middlewares: RestMiddlewareConfig {
metrics: RestMetricsConfig { enabled: true },
..RestMiddlewareConfig::default()
},
..RestConfig::default()
};
let app = RestServer::new(config, Router::new().route("/ready", get(handler))).into_router();
let response = app
.oneshot(
axum::http::Request::builder()
.uri("/ready")
.header("x-request-id", "req-handler-1")
.body(axum::body::Body::empty())
.expect("request"),
)
.await
.expect("response");
let body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.expect("body");
assert_eq!(&body[..], b"req-handler-1");
}
#[cfg(feature = "rpc")]
#[tokio::test]
async fn http_request_id_scope_is_available_to_rpc_interceptor_in_handlers() {
use rs_zero::rpc::{REQUEST_ID_METADATA, request_id_interceptor};
use tonic::{Request as TonicRequest, service::Interceptor};
async fn handler() -> String {
let mut interceptor = request_id_interceptor();
let request = interceptor.call(TonicRequest::new(())).expect("request");
request
.metadata()
.get(REQUEST_ID_METADATA)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string()
}
let registry = MetricsRegistry::new();
let config = RestConfig {
metrics_registry: Some(registry),
middlewares: RestMiddlewareConfig {
metrics: RestMetricsConfig { enabled: true },
..RestMiddlewareConfig::default()
},
..RestConfig::default()
};
let app = RestServer::new(config, Router::new().route("/ready", get(handler))).into_router();
let response = app
.oneshot(
axum::http::Request::builder()
.uri("/ready")
.header("x-request-id", "req-rpc-scope-1")
.body(axum::body::Body::empty())
.expect("request"),
)
.await
.expect("response");
let body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.expect("body");
assert_eq!(&body[..], b"req-rpc-scope-1");
}
#[test]
fn http_traceparent_is_parsed_without_forging_trace_context() {
let request = axum::http::Request::builder()
.uri("/ready")
.header(
"traceparent",
"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
)
.body(())
.expect("request");
let traceparent = traceparent_from_headers(request.headers()).expect("traceparent");
assert_eq!(
trace_id_from_traceparent(&traceparent),
Some("4bf92f3577b34da6a3ce929d0e0e4736")
);
assert!(current_trace_id().is_none());
}
#[test]
fn http_correlation_context_uses_low_cardinality_route_fields() {
let request = axum::http::Request::builder()
.uri("/users/42?token=raw-secret")
.header("x-request-id", "req-http-1")
.header(
"traceparent",
"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
)
.body(())
.expect("request");
let context = CorrelationContext::from_http_headers(
Some("users"),
"GET",
Some("/users/:id"),
request.headers(),
);
let pairs = context.as_pairs();
assert!(pairs.contains(&("service".to_string(), "users".to_string())));
assert!(pairs.contains(&("transport".to_string(), "http".to_string())));
assert!(pairs.contains(&("route".to_string(), "/users/:id".to_string())));
assert!(pairs.contains(&("method".to_string(), "GET".to_string())));
assert!(pairs.contains(&("request_id".to_string(), "req-http-1".to_string())));
assert!(pairs.contains(&(
"trace_id".to_string(),
"4bf92f3577b34da6a3ce929d0e0e4736".to_string()
)));
assert!(pairs.contains(&("span_id".to_string(), "00f067aa0ba902b7".to_string())));
assert!(!pairs.iter().any(|(_, value)| value.contains("42")));
assert!(!pairs.iter().any(|(_, value)| value.contains("raw-secret")));
}
#[test]
fn raw_http_paths_are_not_used_as_correlation_routes() {
let request = axum::http::Request::builder()
.uri("/users/42?api_key=raw-secret")
.body(())
.expect("request");
let context =
CorrelationContext::from_http_headers(Some("users"), "GET", None, request.headers());
assert_eq!(context.route(), "unknown");
assert!(
!context
.as_pairs()
.iter()
.any(|(_, value)| value == "/users/42")
);
}
#[cfg(feature = "otlp")]
#[test]
fn http_traceparent_extracts_valid_opentelemetry_parent() {
use opentelemetry::global;
use opentelemetry::trace::TraceContextExt;
use opentelemetry_sdk::propagation::TraceContextPropagator;
use rs_zero::observability::opentelemetry_context_from_headers;
global::set_text_map_propagator(TraceContextPropagator::new());
let request = axum::http::Request::builder()
.uri("/ready")
.header(
"traceparent",
"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
)
.body(())
.expect("request");
let context = opentelemetry_context_from_headers(request.headers()).expect("context");
assert_eq!(
context.span().span_context().trace_id().to_string(),
"4bf92f3577b34da6a3ce929d0e0e4736"
);
}
#[cfg(feature = "rpc")]
#[test]
fn rpc_request_id_metadata_is_available_for_trace_correlation() {
use rs_zero::observability::{
insert_traceparent_metadata, request_id_from_metadata, traceparent_from_metadata,
};
use rs_zero::rpc::request_id_interceptor;
use tonic::{Request, service::Interceptor};
let mut interceptor = request_id_interceptor();
let request = interceptor.call(Request::new(())).expect("request");
assert!(request_id_from_metadata(request.metadata()).is_some());
let mut request = Request::new(());
insert_traceparent_metadata(
request.metadata_mut(),
"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
)
.expect("traceparent");
assert!(traceparent_from_metadata(request.metadata()).is_some());
}
#[cfg(feature = "rpc")]
#[test]
fn rpc_correlation_context_uses_metadata_and_method_pattern() {
use rs_zero::observability::insert_traceparent_metadata;
use tonic::Request;
let mut request = Request::new(());
request
.metadata_mut()
.insert("x-request-id", "req-rpc-1".parse().expect("request id"));
insert_traceparent_metadata(
request.metadata_mut(),
"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
)
.expect("traceparent");
let context =
CorrelationContext::from_rpc_metadata("billing.Billing", "GetInvoice", request.metadata());
let pairs = context.as_pairs();
assert!(pairs.contains(&("service".to_string(), "billing.Billing".to_string())));
assert!(pairs.contains(&("transport".to_string(), "grpc".to_string())));
assert!(pairs.contains(&("route".to_string(), "GetInvoice".to_string())));
assert!(pairs.contains(&("method".to_string(), "GetInvoice".to_string())));
assert!(pairs.contains(&("request_id".to_string(), "req-rpc-1".to_string())));
assert_eq!(
span_id_from_traceparent(context.traceparent().expect("traceparent")),
Some("00f067aa0ba902b7")
);
assert!(!pairs.iter().any(|(_, value)| value.contains("token")));
}
#[cfg(feature = "rpc")]
#[tokio::test]
async fn rpc_observe_helper_reads_correlation_from_metadata() {
use rs_zero::observability::{insert_traceparent_metadata, observe_rpc_unary_with_metadata};
use tonic::Request;
let registry = MetricsRegistry::new();
let mut request = Request::new(());
request.metadata_mut().insert(
"x-request-id",
"req-rpc-metadata-1".parse().expect("request id"),
);
insert_traceparent_metadata(
request.metadata_mut(),
"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
)
.expect("traceparent");
let value = observe_rpc_unary_with_metadata(
Some(®istry),
"hello.Hello",
"GetHello",
request.metadata(),
async { Ok::<_, tonic::Status>("ok") },
)
.await
.expect("rpc");
assert_eq!(value, "ok");
let text = registry.render_prometheus();
assert!(text.contains("service=\"hello.Hello\",method=\"GetHello\",code=\"Ok\""));
}
#[cfg(all(feature = "otlp", feature = "rpc"))]
#[test]
fn rpc_traceparent_extracts_valid_opentelemetry_parent() {
use opentelemetry::global;
use opentelemetry::trace::TraceContextExt;
use opentelemetry_sdk::propagation::TraceContextPropagator;
use rs_zero::observability::{
insert_traceparent_metadata, opentelemetry_context_from_metadata,
};
use tonic::Request;
global::set_text_map_propagator(TraceContextPropagator::new());
let mut request = Request::new(());
insert_traceparent_metadata(
request.metadata_mut(),
"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
)
.expect("traceparent");
let context = opentelemetry_context_from_metadata(request.metadata()).expect("context");
assert_eq!(
context.span().span_context().trace_id().to_string(),
"4bf92f3577b34da6a3ce929d0e0e4736"
);
}
#[cfg(feature = "rpc")]
#[test]
fn request_context_injects_metadata_without_manual_interceptor() {
use rs_zero::layer::context::RequestContext;
use rs_zero::rpc::REQUEST_ID_METADATA;
use tonic::Request;
let context = RequestContext::new("hello", "grpc", "SayHello", "SayHello")
.with_request_id("ctx-req-1")
.with_traceparent("00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01");
let mut request = Request::new(());
context
.inject_tonic_metadata(request.metadata_mut())
.expect("inject");
assert_eq!(
request
.metadata()
.get(REQUEST_ID_METADATA)
.and_then(|value| value.to_str().ok()),
Some("ctx-req-1")
);
assert_eq!(
request
.metadata()
.get("traceparent")
.and_then(|value| value.to_str().ok()),
Some("00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01")
);
}
#[cfg(feature = "rpc")]
#[test]
fn request_context_reads_tonic_metadata_for_server_layers() {
use rs_zero::{layer::context::RequestContext, observability::insert_traceparent_metadata};
use tonic::Request;
let mut request = Request::new(());
request
.metadata_mut()
.insert("x-request-id", "ctx-rpc-1".parse().expect("request id"));
insert_traceparent_metadata(
request.metadata_mut(),
"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
)
.expect("traceparent");
let context =
RequestContext::from_tonic_metadata("hello.Hello", "SayHello", request.metadata());
assert_eq!(context.request_id(), Some("ctx-rpc-1"));
assert_eq!(context.trace_id(), Some("4bf92f3577b34da6a3ce929d0e0e4736"));
assert_eq!(context.span_id(), Some("00f067aa0ba902b7"));
assert_eq!(context.route(), "SayHello");
}
#[cfg(feature = "rpc")]
#[tokio::test]
async fn rpc_client_builder_scopes_request_id_for_interceptor() {
use rs_zero::rpc::{
REQUEST_ID_METADATA, RpcClientBuilder, RpcClientConfig, request_id_interceptor,
};
use tonic::{Request, service::Interceptor};
let builder = RpcClientBuilder::new(RpcClientConfig::new("http://127.0.0.1:50051"));
let request_id = builder
.scope_request_id("builder-scope-req-1", async {
let mut interceptor = request_id_interceptor();
let request = interceptor.call(Request::new(())).expect("request");
request
.metadata()
.get(REQUEST_ID_METADATA)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string()
})
.await;
assert_eq!(request_id, "builder-scope-req-1");
}