iroh_test/
logging.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
//! Logging during tests.

use tokio::runtime::RuntimeFlavor;
use tracing::level_filters::LevelFilter;
use tracing_subscriber::{
    layer::{Layer, SubscriberExt},
    util::SubscriberInitExt,
    EnvFilter,
};

/// Configures logging for the current test, **single-threaded runtime only**.
///
/// This setup can be used for any sync test or async test using a single-threaded tokio
/// runtime (the default).
///
/// This configures logging that will interact well with tests: logs will be captured by the
/// test framework and only printed on failure.
///
/// The logging is unfiltered, it logs all crates and modules on TRACE level.  If that's too
/// much consider if your test is too large (or write a version that allows filtering...).
///
/// # Example
///
/// ```
/// #[tokio::test]
/// async fn test_something() {
///     let _guard = iroh_test::logging::setup();
///     assert!(true);
/// }
/// ```
#[must_use = "The tracing guard must only be dropped at the end of the test"]
pub fn setup() -> tracing::subscriber::DefaultGuard {
    if let Ok(handle) = tokio::runtime::Handle::try_current() {
        match handle.runtime_flavor() {
            RuntimeFlavor::CurrentThread => (),
            RuntimeFlavor::MultiThread => {
                panic!("setup_logging() does not work in a multi-threaded tokio runtime");
            }
            _ => panic!("unknown runtime flavour"),
        }
    }
    testing_subscriber().set_default()
}

/// The first call to this function will install a global logger.
///
/// The logger uses the `RUST_LOG` environment variable to decide on what level to log
/// anything, which is set by our CI.  When running the tests with nextest the log
/// output will be captured for just the executing test.
///
/// Logs to stdout since the assertion messages are logged on stderr by default.
pub fn setup_multithreaded() {
    tracing_subscriber::registry()
        .with(
            tracing_subscriber::fmt::layer()
                .event_format(tracing_subscriber::fmt::format().with_line_number(true))
                .with_writer(std::io::stdout),
        )
        .with(EnvFilter::from_default_env())
        .try_init()
        .ok();
}

// /// Invoke the future with test logging configured.
// ///
// /// This can be used to execute any future which uses tracing for logging, it sets up the
// /// logging as [`setup_logging`] does but in a way which will work for both single and
// /// multi-threaded tokio runtimes.
// pub(crate) async fn with_logging<F: Future>(f: F) -> F::Output {
//     f.with_subscriber(testing_subscriber()).await
// }

/// Returns the a [`tracing::Subscriber`] configured for our tests.
///
/// This subscriber will ensure that log output is captured by the test's default output
/// capturing and thus is only shown with the test on failure.  By default it uses
/// `RUST_LOG=trace` as configuration but you can specify the `RUST_LOG` environment
/// variable explicitly to override this.
///
/// To use this in a tokio multi-threaded runtime use:
///
/// ```ignore
/// use tracing_future::WithSubscriber;
/// use iroh_test::logging::testing_subscriber;
///
/// #[tokio::test(flavor = "multi_thread")]
/// async fn test_something() -> Result<()> {
///    async move {
///        Ok(())
///    }.with_subscriber(testing_subscriber()).await
/// }
/// ```
pub fn testing_subscriber() -> impl tracing::Subscriber {
    let var = std::env::var_os("RUST_LOG");
    let trace_log_layer = match var {
        Some(_) => None,
        None => Some(
            tracing_subscriber::fmt::layer()
                .event_format(tracing_subscriber::fmt::format().with_line_number(true))
                .with_writer(|| TestWriter)
                .with_filter(LevelFilter::TRACE),
        ),
    };
    let env_log_layer = var.map(|_| {
        tracing_subscriber::fmt::layer()
            .event_format(tracing_subscriber::fmt::format().with_line_number(true))
            .with_writer(|| TestWriter)
            .with_filter(EnvFilter::from_default_env())
    });
    tracing_subscriber::registry()
        .with(trace_log_layer)
        .with(env_log_layer)
}

/// A tracing writer that interacts well with test output capture.
///
/// Using this writer will make sure that the output is captured normally and only printed
/// when the test fails.  See [`setup`] to actually use this.
#[derive(Debug)]
struct TestWriter;

impl std::io::Write for TestWriter {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        print!(
            "{}",
            std::str::from_utf8(buf).expect("tried to log invalid UTF-8")
        );
        Ok(buf.len())
    }
    fn flush(&mut self) -> std::io::Result<()> {
        std::io::stdout().flush()
    }
}