diff --git a/crates/bevy_state/src/state_scoped.rs b/crates/bevy_state/src/state_scoped.rs index ea2acdcf816db..b2694eb09ebff 100644 --- a/crates/bevy_state/src/state_scoped.rs +++ b/crates/bevy_state/src/state_scoped.rs @@ -160,3 +160,75 @@ pub fn despawn_entities_on_enter_state( } } } + +#[cfg(test)] +mod tests { + use super::*; + + use bevy_app::App; + + use crate::{ + app::{AppExtStates, StatesPlugin}, + prelude::CommandsStatesExt, + }; + + #[test] + fn despawn_on_exit_from_computed_state() { + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, States)] + enum State { + On, + Off, + } + + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] + struct ComputedState; + impl bevy_state::state::ComputedStates for ComputedState { + type SourceStates = State; + + fn compute(sources: Self::SourceStates) -> Option { + match sources { + State::On => Some(ComputedState), + State::Off => None, + } + } + } + + let mut app = App::new(); + app.add_plugins(StatesPlugin); + + app.insert_state(State::On); + app.add_computed_state::(); + app.update(); + + assert_eq!( + app.world() + .resource::>() + .get(), + &State::On + ); + assert_eq!( + app.world() + .resource::>() + .get(), + &ComputedState + ); + + let entity = app.world_mut().spawn(DespawnOnExit(ComputedState)).id(); + assert!(app.world().get_entity(entity).is_ok()); + + app.world_mut().commands().set_state(State::Off); + app.update(); + + assert_eq!( + app.world() + .resource::>() + .get(), + &State::Off + ); + assert!(app + .world() + .get_resource::>() + .is_none()); + assert!(app.world().get_entity(entity).is_err()); + } +}