klieo-otel 3.5.0

//! Test utilities for asserting against emitted `tracing` spans
//! without an OTLP collector.
//!
//! The flagship type is [`InMemorySpans`] — a `tracing_subscriber::Layer`
//! that captures every span the subscriber receives into an
//! `Arc<Mutex<Vec<SpanRecord>>>` ring buffer. Tests can then assert that
//! a specific span name fired, that an attribute matched a value, or
//! that a sequence of sibling spans appeared in order.
//!
//! Without this layer, asserting that `gen_ai.system="ollama"` reached
//! the OTel layer requires a docker Jaeger sidecar. With it, a unit
//! test composing `tracing_subscriber::registry().with(spans.layer())`
//! suffices.
//!
//! # Example
//!
//! ```no_run
//! use klieo_otel::test_utils::InMemorySpans;
//! use tracing_subscriber::layer::SubscriberExt;
//! use tracing_subscriber::util::SubscriberInitExt;
//!
//! let spans = InMemorySpans::new();
//! let _guard = tracing_subscriber::registry()
//!     .with(spans.layer())
//!     .set_default();
//!
//! tracing::info_span!("llm.call", gen_ai.system = "ollama").in_scope(|| {
//!     // ... agent code ...
//! });
//!
//! let captured = spans.snapshot();
//! assert!(captured.iter().any(|r| r.name == "llm.call"
//!     && r.fields.get("gen_ai.system").map(String::as_str) == Some("ollama")));
//! ```

use std::collections::HashMap;
use std::fmt;
use std::sync::{Arc, Mutex};

use tracing::field::{Field, Visit};
use tracing::span::Attributes;
use tracing::{Id, Subscriber};
use tracing_subscriber::layer::Context;
use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::Layer;

/// A single captured span observation.
#[derive(Debug, Clone)]
pub struct SpanRecord {
    /// Span name (the first argument to `tracing::*_span!`).
    pub name: String,
    /// `module_path!()` target — usually the crate::module of the span site.
    pub target: String,
    /// Static + dynamic fields the span carried at `new_span` time,
    /// stringified via `Debug`.
    pub fields: HashMap<String, String>,
}

/// In-memory span sink. Cheap to clone — internal storage is a single
/// `Arc<Mutex<Vec<SpanRecord>>>`.
#[derive(Debug, Clone, Default)]
pub struct InMemorySpans {
    inner: Arc<Mutex<Vec<SpanRecord>>>,
}

impl InMemorySpans {
    /// Construct a fresh sink with an empty buffer.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Build the `tracing::Layer` that pushes spans into this sink.
    /// Compose with `tracing_subscriber::registry()` like any other
    /// layer.
    #[must_use]
    pub fn layer(&self) -> InMemoryLayer {
        InMemoryLayer {
            inner: Arc::clone(&self.inner),
        }
    }

    /// Snapshot the buffer. Returns a clone of every record captured so
    /// far; subsequent emissions are not reflected in the returned
    /// vector. Use [`drain`](Self::drain) if you want to consume the
    /// buffer.
    pub fn snapshot(&self) -> Vec<SpanRecord> {
        self.inner.lock().unwrap_or_else(|p| p.into_inner()).clone()
    }

    /// Take every captured record and reset the buffer to empty.
    pub fn drain(&self) -> Vec<SpanRecord> {
        std::mem::take(&mut *self.inner.lock().unwrap_or_else(|p| p.into_inner()))
    }
}

/// The `tracing::Layer` impl returned by [`InMemorySpans::layer`].
#[derive(Debug, Clone)]
pub struct InMemoryLayer {
    inner: Arc<Mutex<Vec<SpanRecord>>>,
}

impl<S> Layer<S> for InMemoryLayer
where
    S: Subscriber + for<'a> LookupSpan<'a>,
{
    fn on_new_span(&self, attrs: &Attributes<'_>, _id: &Id, _ctx: Context<'_, S>) {
        let metadata = attrs.metadata();
        let mut visitor = FieldVisitor::default();
        attrs.record(&mut visitor);
        let record = SpanRecord {
            name: metadata.name().to_string(),
            target: metadata.target().to_string(),
            fields: visitor.0,
        };
        self.inner
            .lock()
            .unwrap_or_else(|p| p.into_inner())
            .push(record);
    }
}

#[derive(Default)]
struct FieldVisitor(HashMap<String, String>);

impl Visit for FieldVisitor {
    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
        self.0
            .insert(field.name().to_string(), format!("{value:?}"));
    }

    fn record_str(&mut self, field: &Field, value: &str) {
        self.0.insert(field.name().to_string(), value.to_string());
    }

    fn record_i64(&mut self, field: &Field, value: i64) {
        self.0.insert(field.name().to_string(), value.to_string());
    }

    fn record_u64(&mut self, field: &Field, value: u64) {
        self.0.insert(field.name().to_string(), value.to_string());
    }

    fn record_bool(&mut self, field: &Field, value: bool) {
        self.0.insert(field.name().to_string(), value.to_string());
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tracing_subscriber::layer::SubscriberExt;

    #[test]
    fn captures_span_name_and_string_attribute() {
        let spans = InMemorySpans::new();
        let dispatch =
            tracing::subscriber::set_default(tracing_subscriber::registry().with(spans.layer()));

        tracing::info_span!(target: "klieo::llm", "llm.call", gen_ai.system = "ollama")
            .in_scope(|| {});

        drop(dispatch);

        let captured = spans.snapshot();
        let llm = captured
            .iter()
            .find(|r| r.name == "llm.call")
            .expect("llm.call span captured");
        assert_eq!(llm.target, "klieo::llm");
        assert_eq!(
            llm.fields.get("gen_ai.system").map(String::as_str),
            Some("ollama")
        );
    }

    #[test]
    fn captures_sibling_spans_in_order() {
        let spans = InMemorySpans::new();
        let dispatch =
            tracing::subscriber::set_default(tracing_subscriber::registry().with(spans.layer()));

        tracing::info_span!("first").in_scope(|| {});
        tracing::info_span!("second").in_scope(|| {});
        tracing::info_span!("third").in_scope(|| {});

        drop(dispatch);

        let names: Vec<String> = spans
            .snapshot()
            .into_iter()
            .map(|r| r.name)
            .filter(|n| n == "first" || n == "second" || n == "third")
            .collect();
        assert_eq!(names, vec!["first", "second", "third"]);
    }

    #[test]
    fn drain_empties_the_buffer() {
        let spans = InMemorySpans::new();
        let dispatch =
            tracing::subscriber::set_default(tracing_subscriber::registry().with(spans.layer()));
        tracing::info_span!("once").in_scope(|| {});
        drop(dispatch);

        let first = spans.drain();
        assert_eq!(first.len(), 1);
        let second = spans.snapshot();
        assert!(second.is_empty(), "drain must reset the buffer");
    }

    #[test]
    fn captures_numeric_and_boolean_fields() {
        let spans = InMemorySpans::new();
        let dispatch =
            tracing::subscriber::set_default(tracing_subscriber::registry().with(spans.layer()));
        tracing::info_span!("metrics", count = 7_i64, hit = true, tokens = 42_u64).in_scope(|| {});
        drop(dispatch);

        let r = spans.snapshot();
        let m = r.iter().find(|x| x.name == "metrics").unwrap();
        assert_eq!(m.fields.get("count").map(String::as_str), Some("7"));
        assert_eq!(m.fields.get("hit").map(String::as_str), Some("true"));
        assert_eq!(m.fields.get("tokens").map(String::as_str), Some("42"));
    }
}