rmv_bevy_testing_tools/
app.rs

1use std::time::Duration;
2
3use bevy_app::{App, AppExit};
4use bevy_derive::{Deref, DerefMut};
5use bevy_ecs::{
6    event::{Event, SendBatchIds},
7    query::{QueryData, QueryFilter, QuerySingleError, ReadOnlyQueryData, WorldQuery},
8};
9use bevy_internal::time::{Time, Virtual};
10use bevy_state::state::{FreelyMutableState, NextState, State, States};
11
12use crate::{
13    events::{CollectedEvents, EventCollectorPlugin, EventFilterPlugin},
14    traits::{AdvanceTime, BasicQuery, CollectEvents, ImmediateQuery, ManageState, SendEvents},
15};
16
17#[derive(Debug, Deref, DerefMut)]
18pub struct TestApp(pub App);
19
20impl TestApp {
21    pub fn get_state<S: States>(&self) -> Option<&S> {
22        self.world().get_resource::<State<S>>().map(|s| s.get())
23    }
24    pub fn get_next_state<S: FreelyMutableState>(&self) -> Option<&NextState<S>> {
25        self.world().get_resource::<NextState<S>>()
26    }
27    pub fn set_next_state<S: FreelyMutableState>(&mut self, next: S) -> Option<()> {
28        self.world_mut()
29            .get_resource_mut::<NextState<S>>()
30            .map(|mut s| s.set(next))
31    }
32}
33
34impl Drop for TestApp {
35    fn drop(&mut self) {
36        self.world_mut().send_event(AppExit::Success);
37    }
38}
39
40impl ManageState for TestApp {
41    fn get_state<S: States>(&self) -> Option<&S> {
42        self.world().get_resource::<State<S>>().map(|s| s.get())
43    }
44    fn get_next_state<S: FreelyMutableState>(&self) -> Option<&NextState<S>> {
45        self.world().get_resource::<NextState<S>>()
46    }
47    fn set_next_state<S: FreelyMutableState>(&mut self, next: S) -> Option<()> {
48        self.world_mut()
49            .get_resource_mut::<NextState<S>>()
50            .map(|mut s| s.set(next))
51    }
52}
53
54impl BasicQuery for TestApp {
55    fn query_any<'a, Q, C>(&mut self) -> bool
56    where
57        Q: QueryData<Item<'a> = C>,
58    {
59        let mut q = self.world_mut().query::<Q>();
60        q.iter(self.world()).next().is_some()
61    }
62}
63
64impl CollectEvents for TestApp {
65    fn collect_events<E: Event + Clone>(&mut self) -> &mut TestApp {
66        self.add_plugins(EventCollectorPlugin::<E>::default());
67        self
68    }
69
70    fn collect_events_only<E: Event + Clone + PartialEq>(&mut self, event: E) -> &mut TestApp {
71        self.add_plugins(EventFilterPlugin::<E>::Only(event.clone()));
72        self
73    }
74
75    fn collect_events_any_of<E: Event + Clone + PartialEq>(&mut self, events: &[E]) -> &mut Self {
76        self.add_plugins(EventFilterPlugin::<E>::AnyOf(events.into()));
77        self
78    }
79
80    fn get_collected_events<E: Event + Clone>(&self) -> Option<Vec<E>> {
81        self.world()
82            .get_resource::<CollectedEvents<E>>()
83            .map(|e| e.get().clone())
84    }
85}
86
87impl SendEvents for TestApp {
88    fn send_event_default<E: Event + Default>(&mut self) {
89        self.world_mut().send_event_default::<E>();
90    }
91    fn send_event<E: Event>(&mut self, event: E) {
92        self.world_mut().send_event::<E>(event);
93    }
94    fn send_event_batch<E: Event>(
95        &mut self,
96        events: impl IntoIterator<Item = E>,
97    ) -> Option<SendBatchIds<E>> {
98        self.world_mut().send_event_batch(events)
99    }
100}
101
102impl ImmediateQuery for TestApp {
103    fn query_single<D>(&mut self) -> Result<<D as WorldQuery>::Item<'_>, QuerySingleError>
104    where
105        D: ReadOnlyQueryData,
106    {
107        let mut query = self.world_mut().query::<D>();
108        query.get_single(self.world_mut())
109    }
110    fn query_single_filtered<D, F>(
111        &mut self,
112    ) -> Result<<D as WorldQuery>::Item<'_>, QuerySingleError>
113    where
114        D: ReadOnlyQueryData,
115        F: QueryFilter,
116    {
117        let mut query = self.world_mut().query_filtered::<D, F>();
118        query.get_single(self.world_mut())
119    }
120    #[cfg(feature = "iter_tools")]
121    fn query_vec<D>(&mut self) -> Vec<<D as WorldQuery>::Item<'_>>
122    where
123        D: ReadOnlyQueryData,
124    {
125        use iter_tools::Itertools;
126
127        let mut query = self.world_mut().query::<D>();
128        query.iter(self.world_mut()).collect_vec()
129    }
130
131    fn query_collect<D, C>(&mut self) -> C
132    where
133        D: ReadOnlyQueryData,
134        for<'a> C: std::iter::FromIterator<<D as bevy_ecs::query::WorldQuery>::Item<'a>>,
135    {
136        let mut query = self.world_mut().query::<D>();
137        let result = query.iter(self.world_mut()).collect::<C>();
138        result
139    }
140}
141
142impl AdvanceTime for TestApp {
143    fn advance_time_by(app: &mut TestApp, duration: Duration) {
144        app.world_mut()
145            .get_resource_mut::<Time<Virtual>>()
146            .unwrap()
147            .advance_by(duration);
148    }
149    fn advance_time_to(app: &mut TestApp, duration: Duration) {
150        app.world_mut()
151            .get_resource_mut::<Time<Virtual>>()
152            .unwrap()
153            .advance_to(duration);
154    }
155}