iroh_test/
logging.rs

1//! Logging during tests.
2
3use tokio::runtime::RuntimeFlavor;
4use tracing::level_filters::LevelFilter;
5use tracing_subscriber::{
6    layer::{Layer, SubscriberExt},
7    util::SubscriberInitExt,
8    EnvFilter,
9};
10
11/// Configures logging for the current test, **single-threaded runtime only**.
12///
13/// This setup can be used for any sync test or async test using a single-threaded tokio
14/// runtime (the default).
15///
16/// This configures logging that will interact well with tests: logs will be captured by the
17/// test framework and only printed on failure.
18///
19/// The logging is unfiltered, it logs all crates and modules on TRACE level.  If that's too
20/// much consider if your test is too large (or write a version that allows filtering...).
21///
22/// # Example
23///
24/// ```
25/// #[tokio::test]
26/// async fn test_something() {
27///     let _guard = iroh_test::logging::setup();
28///     assert!(true);
29/// }
30/// ```
31#[must_use = "The tracing guard must only be dropped at the end of the test"]
32pub fn setup() -> tracing::subscriber::DefaultGuard {
33    if let Ok(handle) = tokio::runtime::Handle::try_current() {
34        match handle.runtime_flavor() {
35            RuntimeFlavor::CurrentThread => (),
36            RuntimeFlavor::MultiThread => {
37                panic!("setup_logging() does not work in a multi-threaded tokio runtime");
38            }
39            _ => panic!("unknown runtime flavour"),
40        }
41    }
42    testing_subscriber().set_default()
43}
44
45/// The first call to this function will install a global logger.
46///
47/// The logger uses the `RUST_LOG` environment variable to decide on what level to log
48/// anything, which is set by our CI.  When running the tests with nextest the log
49/// output will be captured for just the executing test.
50///
51/// Logs to stdout since the assertion messages are logged on stderr by default.
52pub fn setup_multithreaded() {
53    tracing_subscriber::registry()
54        .with(
55            tracing_subscriber::fmt::layer()
56                .event_format(tracing_subscriber::fmt::format().with_line_number(true))
57                .with_writer(std::io::stdout),
58        )
59        .with(EnvFilter::from_default_env())
60        .try_init()
61        .ok();
62}
63
64// /// Invoke the future with test logging configured.
65// ///
66// /// This can be used to execute any future which uses tracing for logging, it sets up the
67// /// logging as [`setup_logging`] does but in a way which will work for both single and
68// /// multi-threaded tokio runtimes.
69// pub(crate) async fn with_logging<F: Future>(f: F) -> F::Output {
70//     f.with_subscriber(testing_subscriber()).await
71// }
72
73/// Returns the a [`tracing::Subscriber`] configured for our tests.
74///
75/// This subscriber will ensure that log output is captured by the test's default output
76/// capturing and thus is only shown with the test on failure.  By default it uses
77/// `RUST_LOG=trace` as configuration but you can specify the `RUST_LOG` environment
78/// variable explicitly to override this.
79///
80/// To use this in a tokio multi-threaded runtime use:
81///
82/// ```ignore
83/// use tracing_future::WithSubscriber;
84/// use iroh_test::logging::testing_subscriber;
85///
86/// #[tokio::test(flavor = "multi_thread")]
87/// async fn test_something() -> Result<()> {
88///    async move {
89///        Ok(())
90///    }.with_subscriber(testing_subscriber()).await
91/// }
92/// ```
93pub fn testing_subscriber() -> impl tracing::Subscriber {
94    let var = std::env::var_os("RUST_LOG");
95    let trace_log_layer = match var {
96        Some(_) => None,
97        None => Some(
98            tracing_subscriber::fmt::layer()
99                .event_format(tracing_subscriber::fmt::format().with_line_number(true))
100                .with_writer(|| TestWriter)
101                .with_filter(LevelFilter::TRACE),
102        ),
103    };
104    let env_log_layer = var.map(|_| {
105        tracing_subscriber::fmt::layer()
106            .event_format(tracing_subscriber::fmt::format().with_line_number(true))
107            .with_writer(|| TestWriter)
108            .with_filter(EnvFilter::from_default_env())
109    });
110    tracing_subscriber::registry()
111        .with(trace_log_layer)
112        .with(env_log_layer)
113}
114
115/// A tracing writer that interacts well with test output capture.
116///
117/// Using this writer will make sure that the output is captured normally and only printed
118/// when the test fails.  See [`setup`] to actually use this.
119#[derive(Debug)]
120struct TestWriter;
121
122impl std::io::Write for TestWriter {
123    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
124        print!(
125            "{}",
126            std::str::from_utf8(buf).expect("tried to log invalid UTF-8")
127        );
128        Ok(buf.len())
129    }
130    fn flush(&mut self) -> std::io::Result<()> {
131        std::io::stdout().flush()
132    }
133}