rig_core/test_utils/
trace_capture.rs

1//! A `tracing` layer that records spans and events for test assertions.
2//!
3//! ```
4//! use rig_core::test_utils::TraceCapture;
5//!
6//! let capture = TraceCapture::default();
7//! tracing::subscriber::with_default(capture.subscriber(), || {
8//!     let _span = tracing::info_span!("work", answer = 42_u64).entered();
9//! });
10//! assert_eq!(capture.spans()[0].u64("answer"), Some(42));
11//! ```
12
13use std::sync::{Arc, Mutex, PoisonError};
14
15use serde_json::{Map, Value};
16use tracing::field::{Field, Visit};
17use tracing::span::{Attributes, Id, Record};
18use tracing::{Event, Level, Subscriber};
19use tracing_subscriber::layer::{Context, Layer, SubscriberExt};
20use tracing_subscriber::registry::LookupSpan;
21
22/// Records every span and event it sees. Clones share one record, so a test
23/// keeps a clone and installs another with [`TraceCapture::subscriber`].
24///
25/// Field values are JSON: strings and `Debug`/`Display` renderings are
26/// strings, integers are numbers and booleans are booleans.
27#[derive(Clone, Default)]
28pub struct TraceCapture(Arc<Mutex<Captured>>);
29
30#[derive(Default)]
31struct Captured {
32    spans: Vec<CapturedSpan>,
33    events: Vec<CapturedEvent>,
34}
35
36impl Captured {
37    fn span_mut(&mut self, id: &Id) -> Option<&mut CapturedSpan> {
38        self.spans
39            .iter_mut()
40            .rev()
41            .find(|span| span.id == id.into_u64())
42    }
43}
44
45/// One span as [`TraceCapture`] saw it.
46#[derive(Clone, Debug)]
47pub struct CapturedSpan {
48    /// The span's id, unique while the span is open.
49    pub id: u64,
50    /// The span's name.
51    pub name: &'static str,
52    /// The span's target.
53    pub target: &'static str,
54    /// The explicit parent, or the current span for a contextual span.
55    pub parent: Option<u64>,
56    /// The parent's name.
57    pub parent_name: Option<&'static str>,
58    /// Every field the span declares, in declaration order.
59    pub declared: Vec<&'static str>,
60    /// The values given when the span was created.
61    pub initial: Map<String, Value>,
62    /// Every value recorded later, in order, repeats included.
63    pub recorded: Vec<(String, Value)>,
64    /// The ids this span follows from.
65    pub follows_from: Vec<u64>,
66}
67
68impl CapturedSpan {
69    /// The latest value of `field`: its last recording, else its initial value.
70    pub fn value(&self, field: &str) -> Option<&Value> {
71        self.recorded
72            .iter()
73            .rev()
74            .find(|(name, _)| name == field)
75            .map(|(_, value)| value)
76            .or_else(|| self.initial.get(field))
77    }
78
79    /// [`Self::value`] as text: a string as it is, anything else as JSON.
80    pub fn text(&self, field: &str) -> Option<String> {
81        self.value(field).map(text)
82    }
83
84    /// The latest value of `field` when it is an unsigned integer.
85    pub fn u64(&self, field: &str) -> Option<u64> {
86        self.value(field).and_then(Value::as_u64)
87    }
88
89    /// How many times `field` was recorded after creation.
90    pub fn record_count(&self, field: &str) -> usize {
91        self.recorded
92            .iter()
93            .filter(|(name, _)| name == field)
94            .count()
95    }
96
97    /// The initial values with every later recording applied in order.
98    pub fn values(&self) -> Map<String, Value> {
99        let mut values = self.initial.clone();
100        values.extend(self.recorded.iter().cloned());
101        values
102    }
103
104    /// The span without its id, which differs between runs: name, target,
105    /// parent name, declared fields and [`Self::values`].
106    pub fn summary(&self) -> Value {
107        serde_json::json!({
108            "name": self.name,
109            "target": self.target,
110            "parent": self.parent_name,
111            "fields": self.declared,
112            "values": self.values(),
113        })
114    }
115}
116
117/// One event as [`TraceCapture`] saw it.
118#[derive(Clone, Debug)]
119pub struct CapturedEvent {
120    /// The event's level.
121    pub level: Level,
122    /// The event's target.
123    pub target: &'static str,
124    /// The event's fields, the message under `message`.
125    pub fields: Map<String, Value>,
126}
127
128impl CapturedEvent {
129    /// The event's message.
130    pub fn message(&self) -> String {
131        self.fields.get("message").map(text).unwrap_or_default()
132    }
133}
134
135fn text(value: &Value) -> String {
136    match value {
137        Value::String(text) => text.clone(),
138        other => other.to_string(),
139    }
140}
141
142impl TraceCapture {
143    /// A registry with this capture as its only layer.
144    pub fn subscriber(&self) -> impl Subscriber + Send + Sync + 'static {
145        tracing_subscriber::registry().with(self.clone())
146    }
147
148    /// Every span seen so far, in creation order.
149    pub fn spans(&self) -> Vec<CapturedSpan> {
150        self.lock().spans.clone()
151    }
152
153    /// Every event seen so far, in order.
154    pub fn events(&self) -> Vec<CapturedEvent> {
155        self.lock().events.clone()
156    }
157
158    /// The last span opened so far.
159    pub fn last_span(&self) -> Option<CapturedSpan> {
160        self.lock().spans.last().cloned()
161    }
162
163    /// Every value `field` took on any span, at creation or later, span by
164    /// span in creation order.
165    pub fn values_of(&self, field: &str) -> Vec<Value> {
166        let captured = self.lock();
167        let mut values = Vec::new();
168        for span in &captured.spans {
169            values.extend(span.initial.get(field).cloned());
170            values.extend(
171                span.recorded
172                    .iter()
173                    .filter(|(name, _)| name == field)
174                    .map(|(_, value)| value.clone()),
175            );
176        }
177        values
178    }
179
180    /// The events at WARN, in order, each as its message followed by
181    /// ` name=value` for every other field.
182    pub fn warnings(&self) -> Vec<String> {
183        self.lock()
184            .events
185            .iter()
186            .filter(|event| event.level == Level::WARN)
187            .map(|event| {
188                let mut rendered = event.message();
189                for (name, value) in event.fields.iter().filter(|(name, _)| *name != "message") {
190                    rendered.push_str(&format!(" {name}={}", text(value)));
191                }
192                rendered
193            })
194            .collect()
195    }
196
197    /// Forgets every span and event seen so far.
198    pub fn clear(&self) {
199        let mut captured = self.lock();
200        captured.spans.clear();
201        captured.events.clear();
202    }
203
204    fn lock(&self) -> std::sync::MutexGuard<'_, Captured> {
205        self.0.lock().unwrap_or_else(PoisonError::into_inner)
206    }
207}
208
209impl<S> Layer<S> for TraceCapture
210where
211    S: Subscriber + for<'a> LookupSpan<'a>,
212{
213    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
214        let parent = match attrs.parent() {
215            Some(parent) => ctx.span(parent),
216            None if attrs.is_contextual() => ctx.lookup_current(),
217            None => None,
218        };
219        let mut initial = Map::new();
220        attrs.record(&mut Values(&mut initial));
221        let metadata = attrs.metadata();
222        self.lock().spans.push(CapturedSpan {
223            id: id.into_u64(),
224            name: metadata.name(),
225            target: metadata.target(),
226            parent: parent.as_ref().map(|span| span.id().into_u64()),
227            parent_name: parent.as_ref().map(|span| span.name()),
228            declared: metadata.fields().iter().map(|field| field.name()).collect(),
229            initial,
230            recorded: Vec::new(),
231            follows_from: Vec::new(),
232        });
233    }
234
235    fn on_record(&self, id: &Id, values: &Record<'_>, _: Context<'_, S>) {
236        let mut fields = Map::new();
237        values.record(&mut Values(&mut fields));
238        if let Some(span) = self.lock().span_mut(id) {
239            span.recorded.extend(fields);
240        }
241    }
242
243    fn on_follows_from(&self, id: &Id, follows: &Id, _: Context<'_, S>) {
244        if let Some(span) = self.lock().span_mut(id) {
245            span.follows_from.push(follows.into_u64());
246        }
247    }
248
249    fn on_event(&self, event: &Event<'_>, _: Context<'_, S>) {
250        let mut fields = Map::new();
251        event.record(&mut Values(&mut fields));
252        self.lock().events.push(CapturedEvent {
253            level: *event.metadata().level(),
254            target: event.metadata().target(),
255            fields,
256        });
257    }
258}
259
260struct Values<'a>(&'a mut Map<String, Value>);
261
262impl Visit for Values<'_> {
263    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
264        self.0
265            .insert(field.name().into(), Value::String(format!("{value:?}")));
266    }
267
268    fn record_str(&mut self, field: &Field, value: &str) {
269        self.0.insert(field.name().into(), Value::from(value));
270    }
271
272    fn record_u64(&mut self, field: &Field, value: u64) {
273        self.0.insert(field.name().into(), Value::from(value));
274    }
275
276    fn record_i64(&mut self, field: &Field, value: i64) {
277        self.0.insert(field.name().into(), Value::from(value));
278    }
279
280    fn record_bool(&mut self, field: &Field, value: bool) {
281        self.0.insert(field.name().into(), Value::from(value));
282    }
283}
rig_core/test_utils/trace_capture.rs

rig_core/test_utils/
trace_capture.rs
