tests: Add test utilities to configure logging (#1060)

This is useful in many tests, not just netcheck.  Make it more
globally available.

It also makes it better to enable by default: it just always logs
all modules on the TRACE level.  But it is properly captured by the
test runner and only shown if the test fails.

To change the logging the RUST_LOG environment variable can
be used like normal to override the TRACE logging.

src/hp/netcheck.rs

@@ -1753,21 +1753,12 @@ impl<T: Future + Unpin> Future for MaybeFuture<T> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::test_utils::setup_logging;
     use bytes::BytesMut;
-    use tracing::level_filters::LevelFilter;
-    use tracing_subscriber::{prelude::*, EnvFilter};
-
-    fn setup_logging() {
-        tracing_subscriber::registry()
-            .with(tracing_subscriber::fmt::layer().with_writer(std::io::stderr))
-            .with(EnvFilter::from_default_env())
-            .try_init()
-            .ok();
-    }
 
     #[tokio::test]
     async fn test_basic() -> Result<()> {
-        setup_logging();
+        let _guard = setup_logging();
 
         let (stun_addr, stun_stats, done) = stun::test::serve("0.0.0.0".parse().unwrap()).await?;
 
@@ -1814,10 +1805,11 @@ mod tests {
         // TODO: remove this once this test is no longer flaky.
         // Setup logging, use .set_default() to overrule possible earlier global init from
         // other tests.
-        let log_layer = tracing_subscriber::fmt::layer()
-            .with_writer(std::io::stderr)
-            .with_filter(LevelFilter::TRACE);
-        let _guard = tracing_subscriber::registry().with(log_layer).set_default();
+        // let log_layer = tracing_subscriber::fmt::layer()
+        //     .with_writer(std::io::stderr)
+        //     .with_filter(LevelFilter::TRACE);
+        // let _guard = tracing_subscriber::registry().with(log_layer).set_default();
+        let _guard = setup_logging();
 
         let mut client = Client::new(None)
             .await
@@ -1864,7 +1856,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_udp_tokio() -> Result<()> {
-        setup_logging();
+        let _guard = setup_logging();
         let local_addr = "127.0.0.1";
         let bind_addr = "0.0.0.0";
 
@@ -1897,7 +1889,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_udp_blocked() -> Result<()> {
-        setup_logging();
+        let _guard = setup_logging();
 
         let blackhole = tokio::net::UdpSocket::bind("127.0.0.1:0").await?;
         let stun_addr = blackhole.local_addr()?;
@@ -1931,6 +1923,8 @@ mod tests {
 
     #[tokio::test(flavor = "current_thread", start_paused = true)]
     async fn test_add_report_history_set_preferred_derp() -> Result<()> {
+        let _guard = setup_logging();
+
         // report returns a *Report from (DERP host, Duration)+ pairs.
         fn report(a: impl IntoIterator<Item = (&'static str, u64)>) -> Option<Arc<Report>> {
             let mut report = Report::default();
@@ -2106,7 +2100,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_hairpin() -> Result<()> {
-        setup_logging();
+        let _guard = setup_logging();
 
         // Hairpinning is initiated after we discover our own IPv4 socket address (IP +
         // port) via STUN, so the test needs to have a STUN server and perform STUN over

src/lib.rs

@@ -24,6 +24,9 @@ mod util;
 #[allow(missing_docs)]
 pub mod hp;
 
+#[cfg(test)]
+pub(crate) mod test_utils;
+
 pub use tls::{Keypair, PeerId, PeerIdError, PublicKey, SecretKey, Signature};
 pub use util::{pathbuf_from_name, Hash};
 

src/test_utils.rs

@@ -0,0 +1,55 @@
+//! Internal utilities to support testing.
+
+use tracing::level_filters::LevelFilter;
+use tracing_subscriber::layer::{Layer, SubscriberExt};
+use tracing_subscriber::util::SubscriberInitExt;
+use tracing_subscriber::EnvFilter;
+
+/// Configures logging for the current test.
+///
+/// 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...).
+#[must_use = "The tracing guard must only be dropped at the end of the test"]
+pub(crate) fn setup_logging() -> tracing::subscriber::DefaultGuard {
+    let var = std::env::var_os("RUST_LOG");
+    let trace_log_layer = match var {
+        Some(_) => None,
+        None => Some(
+            tracing_subscriber::fmt::layer()
+                .with_writer(|| TestWriter)
+                .with_filter(LevelFilter::TRACE),
+        ),
+    };
+    let env_log_layer = var.map(|_| {
+        tracing_subscriber::fmt::layer()
+            .with_writer(|| TestWriter)
+            .with_filter(EnvFilter::from_default_env())
+    });
+    tracing_subscriber::registry()
+        .with(trace_log_layer)
+        .with(env_log_layer)
+        .set_default()
+}
+
+/// 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_logging`] 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()
+    }
+}