From 59be93cbd81bab376c0aefed425b7c8de2d77496 Mon Sep 17 00:00:00 2001
From: bstriker <strikeforcezero@gmail.com>
Date: Sat, 30 Mar 2024 14:59:40 -0400
Subject: [PATCH] Add regression test for recursive despawn on ui nodes

---
 crates/bevy_ui/src/layout/mod.rs | 81 ++++++++++++++++++++++++++++++++
 1 file changed, 81 insertions(+)

diff --git a/crates/bevy_ui/src/layout/mod.rs b/crates/bevy_ui/src/layout/mod.rs
index f82c6190c1e7d0..cb6ab16f194112 100644
--- a/crates/bevy_ui/src/layout/mod.rs
+++ b/crates/bevy_ui/src/layout/mod.rs
@@ -1134,4 +1134,85 @@ mod tests {
             }
         }
     }
+
+    struct DespawnTestEntityReference {
+        parent_entity: Entity,
+        child1_entity: Entity,
+        child2_entity: Entity,
+    }
+    fn recursive_despawn_setup() -> (World, Schedule, DespawnTestEntityReference) {
+        let (mut world, mut ui_schedule) = setup_ui_test_world();
+
+
+        let mut child1_entity = None;
+        let mut child2_entity = None;
+        let parent_entity = world.spawn(NodeBundle::default()).with_children(|children| {
+            child1_entity = Some(children.spawn(NodeBundle::default()).with_children(|children| {
+                child2_entity = Some(children.spawn(NodeBundle::default()).id());
+            }).id());
+        }).id();
+
+        ui_schedule.run(&mut world);
+
+        let ui_surface = world.get_resource::<UiSurface>().unwrap();
+        // 1 for root node, 1 for implicit viewport node
+        // 2 children
+        assert_eq!(ui_surface.taffy.total_node_count(), 4);
+
+        (world, ui_schedule, DespawnTestEntityReference {
+            parent_entity,
+            child1_entity: child1_entity.expect("expected child 1"),
+            child2_entity: child2_entity.expect("expected child 2"),
+        })
+    }
+    
+    #[test]
+    fn test_recursive_despawn_on_parent() {
+        let (mut world, mut ui_schedule, DespawnTestEntityReference {
+            parent_entity,
+            child1_entity,
+            child2_entity
+        }) = recursive_despawn_setup();
+
+        let ui_surface = world.get_resource::<UiSurface>().unwrap();
+        
+        let parent_taffy = *ui_surface.entity_to_taffy.get(&parent_entity).unwrap();
+        let child1_taffy = *ui_surface.entity_to_taffy.get(&child1_entity).unwrap();
+        let child2_taffy = *ui_surface.entity_to_taffy.get(&child2_entity).unwrap();
+        assert_eq!(ui_surface.taffy.parent(child2_taffy), Some(child1_taffy));
+        assert_eq!(ui_surface.taffy.parent(child1_taffy), Some(parent_taffy));
+
+        world.commands().entity(parent_entity).despawn_recursive();
+
+        ui_schedule.run(&mut world);
+
+        let ui_surface = world.get_resource::<UiSurface>().unwrap();
+        // all nodes should be removed
+        assert_eq!(ui_surface.taffy.total_node_count(), 0);
+    }
+
+    #[test]
+    fn test_recursive_despawn_on_child() {
+        let (mut world, mut ui_schedule, DespawnTestEntityReference {
+            parent_entity,
+            child1_entity,
+            child2_entity
+        }) = recursive_despawn_setup();
+        
+        let ui_surface = world.get_resource::<UiSurface>().unwrap();
+        
+        let parent_taffy = *ui_surface.entity_to_taffy.get(&parent_entity).unwrap();
+        let child1_taffy = *ui_surface.entity_to_taffy.get(&child1_entity).unwrap();
+        let child2_taffy = *ui_surface.entity_to_taffy.get(&child2_entity).unwrap();
+        assert_eq!(ui_surface.taffy.parent(child2_taffy), Some(child1_taffy));
+        assert_eq!(ui_surface.taffy.parent(child1_taffy), Some(parent_taffy));
+
+        world.commands().entity(child1_entity).despawn_recursive();
+
+        ui_schedule.run(&mut world);
+
+        let ui_surface = world.get_resource::<UiSurface>().unwrap();
+        // only root node and implicit viewport left
+        assert_eq!(ui_surface.taffy.total_node_count(), 2);
+    }
 }