Use helper test functions

giusdp · Jan 16, 2024 · 9e035e0 · 9e035e0
1 parent 7968ba9
commit 9e035e0
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 69 deletions.
diff --git a/src/builder/build_command.rs b/src/builder/build_command.rs
@@ -315,6 +315,8 @@ mod tests {
     use bevy::{prelude::*, utils::HashMap};
     use rstest::rstest;
 
+    use crate::tests::{count, single};
+
     use super::*;
 
     #[test]
@@ -324,14 +326,8 @@ mod tests {
         let builder = TalkBuilder::default()
             .say("Hello")
             .choose(vec![
-                (
-                    "Choice 1".to_string(),
-                    TalkBuilder::default().say("Hi").to_owned(),
-                ),
-                (
-                    "Choice 2".to_string(),
-                    TalkBuilder::default().say("World!").to_owned(),
-                ),
+                ("Choice 1".to_string(), TalkBuilder::default().say("Hi")),
+                ("Choice 2".to_string(), TalkBuilder::default().say("World!")),
             ])
             .say("something");
 
@@ -452,24 +448,10 @@ mod tests {
         );
 
         // Assert that the connections are made correctly
-        assert_eq!(
-            world.query::<(Entity, Leaf<FollowedBy>)>().single(&world).0,
-            root_ent
-        );
-
-        assert_eq!(
-            world.query::<(Entity, Root<FollowedBy>)>().single(&world).0,
-            leaf_ent
-        );
-
+        assert_eq!(single::<(Entity, Leaf<FollowedBy>)>(&mut world).0, root_ent);
+        assert_eq!(single::<(Entity, Root<FollowedBy>)>(&mut world).0, leaf_ent);
         if previous_node_was_choice {
-            assert_eq!(
-                world
-                    .query::<(Entity, Branch<FollowedBy>)>()
-                    .iter(&world)
-                    .count(),
-                2
-            );
+            assert_eq!(count::<(Entity, Branch<FollowedBy>)>(&mut world), 2);
         }
     }
 
@@ -485,10 +467,7 @@ mod tests {
         // Assert that the relationships are built correctly
         assert_ne!(ent, root);
         assert_eq!(leaves.len(), 1);
-        assert_eq!(
-            world.query::<Relations<FollowedBy>>().iter(&world).count(),
-            2
-        );
+        assert_eq!(count::<Relations<FollowedBy>>(&mut world), 2);
     }
     #[test]
     fn test_add_relationships() {
@@ -507,31 +486,10 @@ mod tests {
         form_graph(root, &talk_builder, &mut build_node_entities, &mut world);
 
         // Assert that the relationships are built correctly
-        assert_eq!(
-            world.query::<Relations<FollowedBy>>().iter(&world).count(),
-            5
-        );
-        assert_eq!(
-            world
-                .query::<(Entity, Leaf<FollowedBy>)>()
-                .iter(&world)
-                .count(),
-            1
-        );
-        assert_eq!(
-            world
-                .query::<(Entity, Branch<FollowedBy>)>()
-                .iter(&world)
-                .count(),
-            3
-        );
-        assert_eq!(
-            world
-                .query::<(Entity, Root<FollowedBy>)>()
-                .iter(&world)
-                .count(),
-            1
-        );
+        assert_eq!(count::<Relations<FollowedBy>>(&mut world), 5);
+        assert_eq!(count::<(Entity, Leaf<FollowedBy>)>(&mut world), 1);
+        assert_eq!(count::<(Entity, Branch<FollowedBy>)>(&mut world), 3);
+        assert_eq!(count::<(Entity, Root<FollowedBy>)>(&mut world), 1);
     }
 }
 
@@ -685,7 +643,6 @@ mod integration_tests {
         let mut world = World::default();
         BuildTalkCommand::new(world.spawn_empty().id(), builder).apply(&mut world);
 
-        // TODO: I should assert on the actual structure of the graph instead of simple number of nodes, leaf and roots.
         assert_relationship_nodes(6, 6, 1, &mut world);
     }
 

diff --git a/src/lib.rs b/src/lib.rs
@@ -192,6 +192,13 @@ mod tests {
     }
 
     #[inline]
+    #[track_caller]
+    pub fn count<Q: WorldQuery>(world: &mut World) -> usize {
+        world.query::<Q>().iter(&world).count()
+    }
+
+    #[inline]
+    #[track_caller]
     pub fn single<Q: WorldQuery>(world: &mut World) -> ROQueryItem<Q> {
         world.query::<Q>().single(world)
     }

diff --git a/src/talk_asset.rs b/src/talk_asset.rs
@@ -160,7 +160,7 @@ fn prepare_builder(
 
 #[cfg(test)]
 mod tests {
-    use crate::{prelude::*, FollowedBy};
+    use crate::{prelude::*, tests::count, FollowedBy};
 
     use aery::{edges::Root, operations::utils::Relations, tuple_traits::RelationEntries};
     use bevy::{ecs::system::Command, prelude::*, utils::hashbrown::HashMap};
@@ -208,7 +208,7 @@ mod tests {
             map.insert(index + 2, (target, "Hello"));
         }
         let mut world = build(TalkData::new(script, vec![]));
-        assert_eq!(world.query::<&TextNode>().iter(&world).count(), nodes);
+        assert_eq!(count::<&TextNode>(&mut world), nodes);
         assert_on_text_nodes(world, map);
     }
 
@@ -221,7 +221,7 @@ mod tests {
         };
 
         let mut world = build(TalkData::new(script, vec![]));
-        assert_eq!(world.query::<&TextNode>().iter(&world).count(), 3);
+        assert_eq!(count::<&TextNode>(&mut world), 3);
         let mut map = HashMap::new();
         map.insert(2, (Some(3), "1"));
         map.insert(3, (Some(4), "10"));
@@ -247,9 +247,9 @@ mod tests {
 
         let mut world = build(TalkData::new(script, vec![]));
 
-        assert_eq!(world.query::<&TextNode>().iter(&world).count(), 2);
-        assert_eq!(world.query::<&ChoiceNode>().iter(&world).count(), 1);
-        assert_eq!(world.query::<Root<FollowedBy>>().iter(&world).count(), 1);
+        assert_eq!(count::<&TextNode>(&mut world), 2);
+        assert_eq!(count::<&ChoiceNode>(&mut world), 1);
+        assert_eq!(count::<Root<FollowedBy>>(&mut world), 1);
         let mut map: HashMap<usize, (Vec<u32>, Vec<&str>)> = HashMap::new();
         map.insert(2, (vec![3, 4], vec!["Choice 1", "Choice 2"]));
         assert_on_choice_nodes(&mut world, map);
@@ -275,9 +275,9 @@ mod tests {
         };
         let mut world = build(TalkData::new(script, vec![]));
 
-        assert_eq!(world.query::<&TextNode>().iter(&world).count(), 4);
-        assert_eq!(world.query::<&ChoiceNode>().iter(&world).count(), 1);
-        assert_eq!(world.query::<Root<FollowedBy>>().iter(&world).count(), 1);
+        assert_eq!(count::<&TextNode>(&mut world), 4);
+        assert_eq!(count::<&ChoiceNode>(&mut world), 1);
+        assert_eq!(count::<Root<FollowedBy>>(&mut world), 1);
 
         let mut choice_map = HashMap::new();
         choice_map.insert(4, (vec![5, 6], vec!["Choice 1", "Choice 2"]));
@@ -319,9 +319,9 @@ mod tests {
         };
         let mut world = build(TalkData::new(script, vec![]));
 
-        assert_eq!(world.query::<&TextNode>().iter(&world).count(), 3);
-        assert_eq!(world.query::<&ChoiceNode>().iter(&world).count(), 2);
-        assert_eq!(world.query::<Root<FollowedBy>>().iter(&world).count(), 1);
+        assert_eq!(count::<&TextNode>(&mut world), 3);
+        assert_eq!(count::<&ChoiceNode>(&mut world), 2);
+        assert_eq!(count::<Root<FollowedBy>>(&mut world), 1);
 
         let mut choice_map = HashMap::new();
         choice_map.insert(2, (vec![3, 5], vec!["First Choice 1", "First Choice 2"]));
@@ -373,8 +373,8 @@ mod tests {
         }
         let mut world = build(TalkData::new(script, actors));
 
-        assert_eq!(world.query::<&TextNode>().iter(&world).count(), nodes);
-        assert_eq!(world.query::<&Actor>().iter(&world).count(), 3);
+        assert_eq!(count::<&TextNode>(&mut world), nodes);
+        assert_eq!(count::<&Actor>(&mut world), 3);
 
         assert_on_text_nodes(world, map);
     }