From dfeb0880eeb833ca6653fef197318fa985009cfb Mon Sep 17 00:00:00 2001
From: mat <git@matdoes.dev>
Date: Sat, 27 Sep 2025 21:20:55 -0700
Subject: add test for closing and opening containers

---
 azalea-client/src/test_utils/simulation.rs  | 12 +++++-
 azalea-client/tests/close_open_container.rs | 66 +++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 1 deletion(-)
 create mode 100644 azalea-client/tests/close_open_container.rs

diff --git a/azalea-client/src/test_utils/simulation.rs b/azalea-client/src/test_utils/simulation.rs
index aa3cef1b..4eff4454 100644
--- a/azalea-client/src/test_utils/simulation.rs
+++ b/azalea-client/src/test_utils/simulation.rs
@@ -25,7 +25,7 @@ use azalea_protocol::{
         },
     },
 };
-use azalea_registry::{Biome, DimensionType, EntityKind};
+use azalea_registry::{Biome, DataRegistry, DimensionType, EntityKind};
 use azalea_world::{Chunk, Instance, MinecraftEntityId, Section, palette::PalettedContainer};
 use bevy_app::App;
 use bevy_ecs::{component::Mutable, prelude::*, schedule::ExecutorKind};
@@ -127,6 +127,9 @@ impl Simulation {
     pub fn has_component<T: Component>(&self) -> bool {
         self.app.world().get::<T>(self.entity).is_some()
     }
+    pub fn with_component<T: Component>(&self, f: impl FnOnce(&T)) {
+        f(&mut self.app.world().entity(self.entity).get::<T>().unwrap());
+    }
     pub fn with_component_mut<T: Component<Mutability = Mutable>>(
         &mut self,
         f: impl FnOnce(&mut T),
@@ -293,6 +296,13 @@ fn tick_app(app: &mut App) {
     app.world_mut().run_schedule(GameTick);
 }
 
+pub fn default_login_packet() -> ClientboundLogin {
+    make_basic_login_packet(
+        DimensionType::new_raw(0), // overworld
+        ResourceLocation::new("minecraft:overworld"),
+    )
+}
+
 pub fn make_basic_login_packet(
     dimension_type: DimensionType,
     dimension: ResourceLocation,
diff --git a/azalea-client/tests/close_open_container.rs b/azalea-client/tests/close_open_container.rs
new file mode 100644
index 00000000..f0457b5b
--- /dev/null
+++ b/azalea-client/tests/close_open_container.rs
@@ -0,0 +1,66 @@
+use azalea_chat::FormattedText;
+use azalea_client::{inventory::Inventory, test_utils::prelude::*};
+use azalea_core::position::ChunkPos;
+use azalea_protocol::packets::{
+    ConnectionProtocol,
+    game::{ClientboundContainerClose, ClientboundOpenScreen, ClientboundSetChunkCacheCenter},
+};
+use azalea_registry::MenuKind;
+
+#[test]
+fn test_close_open_container() {
+    init_tracing();
+
+    let mut simulation = Simulation::new(ConnectionProtocol::Game);
+
+    simulation.receive_packet(default_login_packet());
+    simulation.tick();
+    // receive a chunk so the player is "loaded" now
+    simulation.receive_packet(ClientboundSetChunkCacheCenter { x: 1, z: 23 });
+    simulation.receive_packet(make_basic_empty_chunk(
+        ChunkPos::new(1, 23),
+        (384 + 64) / 16,
+    ));
+    simulation.tick();
+
+    // ensure no container is open
+    simulation.with_component(|inventory: &Inventory| {
+        assert!(inventory.container_menu.is_none());
+        assert_eq!(inventory.id, 0);
+    });
+
+    // open a container
+    simulation.receive_packet(ClientboundOpenScreen {
+        container_id: 1,
+        menu_type: MenuKind::Generic9x3,
+        title: FormattedText::default(),
+    });
+    simulation.tick();
+
+    simulation.with_component(|inventory: &Inventory| {
+        assert!(inventory.container_menu.is_some());
+        assert_eq!(inventory.id, 1);
+    });
+
+    // close and open
+    simulation.receive_packet(ClientboundContainerClose { container_id: 1 });
+    simulation.receive_packet(ClientboundOpenScreen {
+        container_id: 2,
+        menu_type: MenuKind::Generic9x3,
+        title: FormattedText::default(),
+    });
+    simulation.tick();
+    simulation.with_component(|inventory: &Inventory| {
+        // ensure that the new container was opened
+        assert!(inventory.container_menu.is_some());
+        assert_eq!(inventory.id, 2);
+    });
+
+    // close with the wrong container id should still close
+    simulation.receive_packet(ClientboundContainerClose { container_id: 123 });
+    simulation.tick();
+    simulation.with_component(|inventory: &Inventory| {
+        assert!(inventory.container_menu.is_none());
+        assert_eq!(inventory.id, 0);
+    });
+}
-- 
cgit v1.2.3