feat(tests): Improve test_utils to warn about mutli-runtime tests

iroh-net/Cargo.toml

@@ -89,6 +89,7 @@ ntest = "0.9"
 pretty_assertions = "1.4"
 rand_chacha = "0.3.1"
 tokio = { version = "1", features = ["io-util", "sync", "rt", "net", "fs", "macros", "time", "test-util"] }
+tracing-futures = "0.2.5"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 
 [build-dependencies]

iroh-net/src/magicsock.rs

@@ -3085,48 +3085,48 @@ pub(crate) mod tests {
 
     #[tokio::test(flavor = "multi_thread")]
     async fn test_two_devices_setup_teardown() -> Result<()> {
-        setup_logging();
-
-        let devices = Devices {
-            stun_ip: "127.0.0.1".parse()?,
-        };
-
-        for _ in 0..10 {
-            let (derp_map, _, cleanup) = run_derp_and_stun(devices.stun_ip).await?;
-            println!("setting up magic stack");
-            let m1 = MagicStack::new(derp_map.clone()).await?;
-            let m2 = MagicStack::new(derp_map.clone()).await?;
-
-            let cleanup_mesh = mesh_stacks(vec![m1.clone(), m2.clone()]).await?;
+        crate::test_utils::with_logging(async move {
+            let devices = Devices {
+                stun_ip: "127.0.0.1".parse()?,
+            };
 
-            // Wait for magicsock to be told about peers from mesh_stacks.
-            println!("waiting for connection");
-            let m1t = m1.clone();
-            let m2t = m2.clone();
-            time::timeout(Duration::from_secs(10), async move {
-                loop {
-                    let ab = m1t.tracked_endpoints().await.contains(&m2t.public());
-                    let ba = m2t.tracked_endpoints().await.contains(&m1t.public());
-                    if ab && ba {
-                        break;
+            for _ in 0..10 {
+                let (derp_map, _, cleanup) = run_derp_and_stun(devices.stun_ip).await?;
+                println!("setting up magic stack");
+                let m1 = MagicStack::new(derp_map.clone()).await?;
+                let m2 = MagicStack::new(derp_map.clone()).await?;
+
+                let cleanup_mesh = mesh_stacks(vec![m1.clone(), m2.clone()]).await?;
+
+                // Wait for magicsock to be told about peers from mesh_stacks.
+                println!("waiting for connection");
+                let m1t = m1.clone();
+                let m2t = m2.clone();
+                time::timeout(Duration::from_secs(10), async move {
+                    loop {
+                        let ab = m1t.tracked_endpoints().await.contains(&m2t.public());
+                        let ba = m2t.tracked_endpoints().await.contains(&m1t.public());
+                        if ab && ba {
+                            break;
+                        }
                     }
-                }
-            })
-            .await
-            .context("failed to connect peers")?;
+                })
+                .await
+                .context("failed to connect peers")?;
 
-            println!("closing endpoints");
-            m1.endpoint.close(0u32.into(), b"done").await?;
-            m2.endpoint.close(0u32.into(), b"done").await?;
+                println!("closing endpoints");
+                m1.endpoint.close(0u32.into(), b"done").await?;
+                m2.endpoint.close(0u32.into(), b"done").await?;
 
-            assert!(m1.endpoint.conn().inner.is_closed());
-            assert!(m2.endpoint.conn().inner.is_closed());
+                assert!(m1.endpoint.conn().inner.is_closed());
+                assert!(m2.endpoint.conn().inner.is_closed());
 
-            println!("cleaning up");
-            cleanup().await;
-            cleanup_mesh();
-        }
-        Ok(())
+                println!("cleaning up");
+                cleanup().await;
+                cleanup_mesh();
+            }
+            Ok(())
+        }).await
     }
 
     #[tokio::test]

iroh-net/src/test_utils.rs

@@ -1,38 +1,101 @@
 //! Internal utilities to support testing.
 
+use std::future::Future;
+
+use tokio::runtime::RuntimeFlavor;
 use tracing::level_filters::LevelFilter;
+use tracing_futures::WithSubscriber;
+use tracing_subscriber::fmt::format::FmtSpan;
 use tracing_subscriber::layer::{Layer, SubscriberExt};
 use tracing_subscriber::util::SubscriberInitExt;
 use tracing_subscriber::EnvFilter;
 
-/// Configures logging for the current test.
+/// 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).  For multi-threaded runtimes use [`with_logging`].
 ///
 /// 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
+///
+/// ```no_run
+/// #[tokio::test]
+/// async fn test_something() {
+///     let _guard = crate::test_utils::setup_logging();
+///     assert!(true);
+/// }
 #[must_use = "The tracing guard must only be dropped at the end of the test"]
 #[allow(dead_code)]
 pub(crate) fn setup_logging() -> 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, \
+                        use with_logging() instead"
+                );
+            }
+            _ => panic!("unknown runtime flavour"),
+        }
+    }
+    testing_subscriber().set_default()
+}
+
+/// 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:
+///
+/// ```no_run
+/// use tracing_future::WithSubscriber;
+/// use crate::test_utils::testing_subscriber;
+///
+/// #[tokio::test(flavor = "multi_thread")]
+/// async fn test_something() -> Result<()> {
+///    async move {
+///        Ok(())
+///    }.with_subscriber(testing_subscriber()).await
+/// }
+/// ```
+pub(crate) 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()
+                .with_span_events(FmtSpan::CLOSE)
                 .with_writer(|| TestWriter)
                 .with_filter(LevelFilter::TRACE),
         ),
     };
     let env_log_layer = var.map(|_| {
         tracing_subscriber::fmt::layer()
+            .with_span_events(FmtSpan::CLOSE)
             .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.