2 files changed, 65 insertions, 21 deletions

diff --git a/azalea-client/src/test_utils/simulation.rs b/azalea-client/src/test_utils/simulation.rs
index 5a1c9c52..a6bc22ab 100644
--- a/azalea-client/src/test_utils/simulation.rs
+++ b/azalea-client/src/test_utils/simulation.rs
@@ -57,8 +57,9 @@ pub struct Simulation {
 }
 
 impl Simulation {
-    pub fn new(initial_connection_protocol: ConnectionProtocol) -> Self {
+    pub fn new(conn_protocol: ConnectionProtocol) -> Self {
         let mut app = create_simulation_app();
+
         let mut entity = app.world_mut().spawn_empty();
         let (player, rt) =
             create_local_player_bundle(entity.id(), ConnectionProtocol::Configuration);
@@ -78,7 +79,7 @@ impl Simulation {
         let mut simulation = Self { app, entity, rt };
 
         #[allow(clippy::single_match)]
-        match initial_connection_protocol {
+        match conn_protocol {
             ConnectionProtocol::Configuration => {}
             ConnectionProtocol::Game => {
                 simulation.receive_packet(ClientboundRegistryData {
@@ -97,7 +98,7 @@ impl Simulation {
                 simulation.receive_packet(ClientboundFinishConfiguration);
                 simulation.tick();
             }
-            _ => unimplemented!("unsupported ConnectionProtocol {initial_connection_protocol:?}"),
+            _ => unimplemented!("unsupported ConnectionProtocol {conn_protocol:?}"),
         }
 
         simulation
@@ -300,10 +301,13 @@ fn create_local_player_bundle(
 fn create_simulation_app() -> App {
     let mut app = App::new();
 
+    let mut plugins = bevy_app::PluginGroup::build(crate::DefaultPlugins);
     #[cfg(feature = "log")]
-    app.add_plugins(
-        bevy_app::PluginGroup::build(crate::DefaultPlugins).disable::<bevy_log::LogPlugin>(),
-    );
+    {
+        plugins = plugins.disable::<bevy_log::LogPlugin>();
+    }
+
+    app.add_plugins(plugins);
 
     app.edit_schedule(bevy_app::Main, |schedule| {
         // makes test results more reproducible
diff --git a/azalea-client/src/test_utils/tracing.rs b/azalea-client/src/test_utils/tracing.rs
index 207a4625..3be92260 100644
--- a/azalea-client/src/test_utils/tracing.rs
+++ b/azalea-client/src/test_utils/tracing.rs
@@ -1,27 +1,67 @@
+use std::sync::LazyLock;
+
 use bevy_log::tracing_subscriber::{
-    self, EnvFilter, Layer,
+    self, EnvFilter, Layer, Registry,
+    filter::Filtered,
+    fmt,
     layer::{Context, SubscriberExt},
+    reload::{self, Handle},
     util::SubscriberInitExt,
 };
+use parking_lot::{Mutex, MutexGuard};
 use tracing::{Event, Level, Subscriber};
 
-pub fn init_tracing() {
-    init_tracing_with_level(Level::WARN);
+static LOCK: LazyLock<Mutex<()>> = LazyLock::new(|| Mutex::new(()));
+
+pub fn init<'a>() -> MutexGuard<'a, ()> {
+    init_with_level(Level::WARN)
+}
+
+// can't treat these as one layer due to https://github.com/tokio-rs/tracing/issues/1629
+struct LayerReloadHandles {
+    filter: Handle<Filtered<fmt::Layer<Registry>, EnvFilter, Registry>, Registry>,
+    test: Handle<TestTracingLayer, Registry>,
 }
 
-pub fn init_tracing_with_level(max_level: Level) {
+static RELOAD_HANDLES: LazyLock<LayerReloadHandles> = LazyLock::new(|| {
+    let (filter_layer, filter_reload_handle) = reload::Layer::new(
+        fmt::layer().with_filter(
+            EnvFilter::builder()
+                .with_default_directive(Level::WARN.into())
+                .from_env_lossy(),
+        ),
+    );
+    let (test_layer, test_reload_handle) = reload::Layer::new(TestTracingLayer {
+        panic_on_level: Level::WARN,
+    });
+
     tracing_subscriber::registry()
-        .with(
-            tracing_subscriber::fmt::layer().with_filter(
-                EnvFilter::builder()
-                    .with_default_directive(max_level.into())
-                    .from_env_lossy(),
-            ),
-        )
-        .with(TestTracingLayer {
-            panic_on_level: max_level,
-        })
+        .with(filter_layer.and_then(test_layer))
         .init();
+
+    LayerReloadHandles {
+        filter: filter_reload_handle,
+        test: test_reload_handle,
+    }
+});
+
+pub fn init_with_level<'a>(max_level: Level) -> MutexGuard<'a, ()> {
+    let lock = LOCK.lock();
+
+    RELOAD_HANDLES
+        .filter
+        .modify(|layer| {
+            *layer.filter_mut() = EnvFilter::builder()
+                .with_default_directive(max_level.into())
+                .from_env_lossy()
+        })
+        .unwrap();
+    RELOAD_HANDLES
+        .test
+        .modify(|layer| layer.panic_on_level = max_level)
+        .unwrap();
+
+    lock
 }
 
 struct TestTracingLayer {
@@ -31,7 +71,7 @@ impl<S: Subscriber> Layer<S> for TestTracingLayer {
     fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) {
         let level = *event.metadata().level();
         if level <= self.panic_on_level {
-            panic!("logged on level {level}");
+            panic!("Logged on level {level}.");
         }
     }
 }