apollo_router/logging/
mod.rs

1#[macro_export]
2/// This is a really simple macro to assert a snapshot of the logs.
3/// To use it call `.with_subscriber(assert_snapshot_subscriber!())` in your test just before calling `await`.
4/// This will assert a snapshot of the logs in pretty yaml format.
5/// You can also use subscriber::with_default(assert_snapshot_subscriber!(), || { ... }) to assert the logs in non async code.
6macro_rules! assert_snapshot_subscriber {
7    () => {
8        $crate::assert_snapshot_subscriber!(tracing_core::LevelFilter::INFO, {})
9    };
10
11    ($redactions:tt) => {
12        $crate::assert_snapshot_subscriber!(tracing_core::LevelFilter::INFO, $redactions)
13    };
14
15    ($level:expr) => {
16        $crate::assert_snapshot_subscriber!($level, {})
17    };
18
19    ($level:expr, $redactions:tt) => {
20        $crate::logging::test::SnapshotSubscriber::create_subscriber($level, |yaml| {
21            insta::with_settings!({sort_maps => true}, {
22                // the tests here will force maps to sort
23                let mut settings = insta::Settings::clone_current();
24                settings.set_snapshot_suffix("logs");
25                settings.set_sort_maps(true);
26                settings.bind(|| {
27                    insta::assert_yaml_snapshot!(yaml, $redactions);
28                });
29            });
30        })
31    };
32}
33
34#[cfg(test)]
35pub(crate) mod test {
36    use std::sync::Arc;
37    use std::sync::Mutex;
38
39    use serde_json::Value;
40    use tracing_core::LevelFilter;
41    use tracing_core::Subscriber;
42    use tracing_subscriber::layer::SubscriberExt;
43
44    use crate::plugins::telemetry::dynamic_attribute::DynAttributeLayer;
45
46    pub(crate) struct SnapshotSubscriber {
47        buffer: Arc<Mutex<Vec<u8>>>,
48        assertion: fn(serde_json::Value),
49    }
50
51    impl std::io::Write for SnapshotSubscriber {
52        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
53            let buf_len = buf.len();
54            self.buffer.lock().unwrap().append(&mut buf.to_vec());
55            Ok(buf_len)
56        }
57
58        fn flush(&mut self) -> std::io::Result<()> {
59            Ok(())
60        }
61    }
62
63    impl Drop for SnapshotSubscriber {
64        fn drop(&mut self) {
65            let log = String::from_utf8(self.buffer.lock().unwrap().to_vec()).unwrap();
66            let parsed: Value = if log.is_empty() {
67                serde_json::json!([])
68            } else {
69                let parsed_log: Vec<Value> = log
70                    .lines()
71                    .map(|line| {
72                        let mut line: serde_json::Value = serde_json::from_str(line).unwrap();
73                        // move the message field to the top level
74                        let fields = line
75                            .as_object_mut()
76                            .unwrap()
77                            .get_mut("fields")
78                            .unwrap()
79                            .as_object_mut()
80                            .unwrap();
81                        let message = fields.remove("message").unwrap_or_default();
82                        line.as_object_mut()
83                            .unwrap()
84                            .insert("message".to_string(), message);
85                        line
86                    })
87                    .collect();
88                serde_json::json!(parsed_log)
89            };
90
91            (self.assertion)(parsed)
92        }
93    }
94
95    impl SnapshotSubscriber {
96        pub(crate) fn create_subscriber(
97            level: LevelFilter,
98            assertion: fn(Value),
99        ) -> impl Subscriber {
100            let collector = Self {
101                buffer: Arc::new(Mutex::new(Vec::new())),
102                assertion,
103            };
104
105            tracing_subscriber::registry::Registry::default()
106                .with(level)
107                .with(DynAttributeLayer::new())
108                .with(
109                    tracing_subscriber::fmt::Layer::default()
110                        .json()
111                        .without_time()
112                        .with_target(false)
113                        .with_file(false)
114                        .with_line_number(false)
115                        .with_writer(Mutex::new(collector)),
116                )
117        }
118    }
119}