bevy_assert 0.1.3

//! test suite for bevy
//!
//! See [`AppTest`].
use std::{
    collections::HashSet,
    fmt::Debug,
    hash::Hash,
    ops::{Deref, DerefMut},
    thread::sleep,
    time::Duration,
};

use bevy::{ecs::query::QueryFilter, prelude::*};

/// Test suite for bevy.
///
/// This suite makes available useful common assertions about the state of the
/// bevy app.
///
/// Create a [`Self::new`]
///
/// ```rust
/// use bevy_assert::AppTest;
///
/// let suite = AppTest::new();
/// ```
///
/// or from [`App`].
///
/// ```rust
/// use bevy::prelude::*;
/// use bevy_assert::AppTest;
///
/// let suite: AppTest = App::new().into();
/// ```
///
/// [`AppTest`] dereferences to [`App`] so you can construct the bevy state
/// normally by calling all the [`App`] methods directly.
///
/// ```rust
/// use bevy::{prelude::*, state::app::StatesPlugin};
/// use bevy_assert::AppTest;
///
/// let mut t = AppTest::new();
/// t.add_plugins(StatesPlugin);
/// ```
///
/// The method variants with `_r` suffix will return a Result instead of
/// panicking.
///
#[derive(Default)]
pub struct AppTest {
    app: App,
    retry_count: usize,
    retry_wait: Duration,
}

impl AppTest {
    /// Creates a new test suite from an empty app.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Asserts a condition
    ///
    /// Use this variant for a Result instead of panicking.
    ///
    /// # Errors
    ///
    /// Errors if assertion fails.
    pub fn assert_r(condition: bool, message: String) -> Result<(), String> {
        if !condition {
            return Err(message);
        }
        Ok(())
    }

    /// Asserts that a state is present in the app and is the specified one.
    ///
    /// # Panics
    ///
    /// Panics if assertion fails.
    pub fn assert_state<S: States>(&self, s: &S) {
        self.assert_state_r(s).unwrap();
    }

    /// Asserts that a state is present in the app and is the specified one.
    ///
    /// Use this variant for a Result instead of panicking.
    ///
    /// # Errors
    ///
    /// - If the state is not in the app.
    /// - If the state is not the specified one.
    pub fn assert_state_r<S: States + Debug>(&self, s: &S) -> Result<(), String> {
        let app_state = self
            .app
            .world()
            .get_resource::<State<S>>()
            .ok_or("State not found".to_string())?
            .get();
        if app_state != s {
            return Err(format!("State was not equal: {app_state:?} != {s:?}"));
        }
        Ok(())
    }

    /// Count all entities
    pub fn count_all(&mut self) -> usize {
        self.count::<()>()
    }

    /// Counts the amount of entities matching the selector type.
    ///
    /// Selection is formed by a [`QueryFilter`].
    pub fn count<F: QueryFilter>(&mut self) -> usize {
        let w = self.app.world_mut();
        w.query_filtered::<(), F>().iter(w).count()
    }

    /// Return a set of all entities
    pub fn all_entities(&mut self) -> HashSet<Entity> {
        self.entities::<()>()
    }

    /// Return a set of all entities matching the filter.
    ///
    /// Selection is formed by a [`QueryFilter`].
    pub fn entities<F: QueryFilter>(&mut self) -> HashSet<Entity> {
        let w = self.app.world_mut();
        w.query_filtered::<Entity, F>()
            .iter(w)
            .collect::<HashSet<_>>()
    }

    /// Simulates a click of a button/key.
    ///
    /// A click is intended as a press followed by a release.
    ///
    /// Any clickable element can be passed, including [`KeyCode`] or
    /// [`MouseButton`].
    ///
    /// *this calls app update internally*
    pub fn click<K: Clone + Eq + Hash + Send + Sync + 'static>(&mut self, k: K) {
        let mut input = self.app.world_mut().resource_mut::<ButtonInput<K>>();
        input.press(k.clone());
        input.release(k);
        self.app.update();
    }

    /// Asserts that the button/key was clicked.
    ///
    /// A click is intended as a press followed by a release.
    ///
    /// # Panics
    ///
    /// Panics if assertion fails
    pub fn assert_clicked<K: Clone + Eq + Hash + Send + Sync + 'static>(&mut self, k: K) {
        self.assert_clicked_r(k).unwrap();
    }

    /// Asserts that the button/key was clicked.
    ///
    /// A click is intended as a press followed by a release.
    ///
    /// Use this variant for a Result instead of panicking.
    ///
    /// # Errors
    ///
    /// - If the button/key was not clicked.
    pub fn assert_clicked_r<K: Clone + Eq + Hash + Send + Sync + 'static>(
        &mut self,
        k: K,
    ) -> Result<(), String> {
        let input = self.app.world().resource::<ButtonInput<K>>();
        if !input.just_pressed(k) {
            return Err("No key was pressed".to_string());
        }
        Ok(())
    }

    /// Sends a message.
    ///
    /// Note: messages are cleared after every call to [`App::update`].
    pub fn send_message<M: Message>(&mut self, m: M) {
        self.world_mut().write_message(m);
    }

    /// Asserts that the message is present in the app.
    ///
    /// Note: messages are cleared after every call to [`App::update`].
    ///
    /// # Panics
    ///
    /// Panics if assertion fails
    pub fn assert_message<M: Message + PartialEq>(&mut self, m: &M) {
        self.assert_message_r(m).unwrap();
    }

    /// Asserts that the message is present in the app.
    ///
    /// Note: messages are cleared after every call to [`App::update`].
    ///
    /// Use this variant for a Result instead of panicking.
    ///
    /// # Errors
    ///
    /// - If the message is not present.
    pub fn assert_message_r<M: Message + PartialEq>(&mut self, m: &M) -> Result<(), String> {
        let msgs = self.world().resource::<Messages<M>>();
        let mut cursor = msgs.get_cursor();
        for msg in cursor.read(msgs) {
            if m == msg {
                return Ok(());
            }
        }
        Err("Message not found".to_string())
    }

    /// Asserts at least one message of this type is present in the app.
    ///
    /// Note: messages are cleared after every call to [`App::update`].
    ///
    /// # Panics
    ///
    /// Panics if assertion fails
    pub fn assert_any_message<M: Message>(&mut self) {
        self.assert_any_message_r::<M>().unwrap();
    }

    /// Asserts at least one message of this type is present in the app.
    ///
    /// Note: messages are cleared after every call to [`App::update`].
    ///
    /// Use this variant for a Result instead of panicking.
    ///
    /// # Errors
    ///
    /// - If no message is present.
    pub fn assert_any_message_r<M: Message>(&mut self) -> Result<(), String> {
        let msgs = self.world().resource::<Messages<M>>();
        if msgs.is_empty() {
            return Err("Message not found".to_string());
        }
        Ok(())
    }

    /// Gets the component from the specified entity.
    ///
    /// # Panics
    ///
    /// If no component or entity is found
    pub fn get<C: Component>(&mut self, id: Entity) -> &C {
        self.get_r::<C>(id).unwrap()
    }

    /// Gets the component from the specified entity.
    ///
    /// # Errors
    ///
    /// If no component or entity is found
    pub fn get_r<C: Component>(&mut self, id: Entity) -> Result<&C, String> {
        self.world()
            .entity(id)
            .get::<C>()
            .ok_or("Component or entity not found".to_string())
    }

    /// Retries the given function
    ///
    /// Note: this retries on Result, so use the _r variant of assertions
    /// within the lambda.
    ///
    /// # Panics
    ///
    /// If retries are exhausted.
    ///
    /// *this calls app update internally*
    pub fn retry<T, E: Debug, F: FnMut(&mut AppTest) -> Result<T, E>>(&mut self, f: F) -> T {
        self.retry_r(f).unwrap()
    }

    /// Retries the given function
    ///
    /// Note: this retries on Result, so use the _r variant of assertions
    /// within the lambda.
    ///
    /// # Errors
    ///
    /// If retries are exhausted.
    ///
    /// *this calls app update internally*
    pub fn retry_r<T, E, F: FnMut(&mut AppTest) -> Result<T, E>>(
        &mut self,
        mut f: F,
    ) -> Result<T, E> {
        let mut ct = 0;
        loop {
            let res = f(self);
            match res {
                Ok(t) => return Ok(t),
                Err(e) => {
                    if ct >= self.retry_count {
                        return Err(e);
                    }
                }
            }
            sleep(self.retry_wait);
            self.update();
            ct += 1;
        }
    }

    /// Sets the configutation for retry behavior.
    ///
    /// This is only used in [`Self::retry`] and [`Self::retry_r`].
    pub fn setup_retry(&mut self, count: usize, wait: Duration) {
        self.retry_count = count;
        self.retry_wait = wait;
    }
}

impl From<App> for AppTest {
    fn from(value: App) -> Self {
        Self {
            app: value,
            retry_count: 0,
            retry_wait: Duration::from_millis(100),
        }
    }
}

impl Deref for AppTest {
    type Target = App;

    fn deref(&self) -> &Self::Target {
        &self.app
    }
}

impl DerefMut for AppTest {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.app
    }
}