Fixing AutoFocus not working if it's spawn before Startup (#19618)

# Objective

- `AutoFocus` component does seem to work with the `set_initial_focus`
that was running in `Startup`.
- If an element is spawn during `PreStartup` or during
`OnEnter(SomeState)` that happens before `Startup`, the focus is
overridden by `set_initial_focus` which sets the focus to the primary
window.


## Solution

- `set_initial_focus` only sets the focus to the `PrimaryWindow` if no
other focus is set.

- *Note*: `cargo test --package bevy_input_focus` was not working, so
some changes are related to that.

## Testing

- `cargo test --package bevy_input_focus`: OK
- `cargo run --package ci`: OK

crates/bevy_input_focus/Cargo.toml

@@ -73,9 +73,6 @@ bevy_reflect = { path = "../bevy_reflect", version = "0.16.0-dev", features = [
 thiserror = { version = "2", default-features = false }
 log = { version = "0.4", default-features = false }
 
-[dev-dependencies]
-smol_str = "0.2"
-
 [lints]
 workspace = true
 

crates/bevy_input_focus/src/lib.rs

@@ -30,7 +30,7 @@ pub mod tab_navigation;
 mod autofocus;
 pub use autofocus::*;
 
-use bevy_app::{App, Plugin, PreUpdate, Startup};
+use bevy_app::{App, Plugin, PostStartup, PreUpdate};
 use bevy_ecs::{prelude::*, query::QueryData, system::SystemParam, traversal::Traversal};
 use bevy_input::{gamepad::GamepadButtonChangedEvent, keyboard::KeyboardInput, mouse::MouseWheel};
 use bevy_window::{PrimaryWindow, Window};
@@ -185,7 +185,7 @@ pub struct InputDispatchPlugin;
 
 impl Plugin for InputDispatchPlugin {
     fn build(&self, app: &mut App) {
-        app.add_systems(Startup, set_initial_focus)
+        app.add_systems(PostStartup, set_initial_focus)
             .init_resource::<InputFocus>()
             .init_resource::<InputFocusVisible>()
             .add_systems(
@@ -218,12 +218,14 @@ pub enum InputFocusSystems {
 #[deprecated(since = "0.17.0", note = "Renamed to `InputFocusSystems`.")]
 pub type InputFocusSet = InputFocusSystems;
 
-/// Sets the initial focus to the primary window, if any.
+/// If no entity is focused, sets the focus to the primary window, if any.
 pub fn set_initial_focus(
     mut input_focus: ResMut<InputFocus>,
     window: Single<Entity, With<PrimaryWindow>>,
 ) {
-    input_focus.0 = Some(*window);
+    if input_focus.0.is_none() {
+        input_focus.0 = Some(*window);
+    }
 }
 
 /// System which dispatches bubbled input events to the focused entity, or to the primary window
@@ -368,24 +370,12 @@ mod tests {
     use super::*;
 
     use alloc::string::String;
-    use bevy_ecs::{
-        lifecycle::HookContext, observer::On, system::RunSystemOnce, world::DeferredWorld,
-    };
+    use bevy_app::Startup;
+    use bevy_ecs::{observer::On, system::RunSystemOnce, world::DeferredWorld};
     use bevy_input::{
         keyboard::{Key, KeyCode},
         ButtonState, InputPlugin,
     };
-    use bevy_window::WindowResolution;
-    use smol_str::SmolStr;
-
-    #[derive(Component)]
-    #[component(on_add = set_focus_on_add)]
-    struct SetFocusOnAdd;
-
-    fn set_focus_on_add(mut world: DeferredWorld, HookContext { entity, .. }: HookContext) {
-        let mut input_focus = world.resource_mut::<InputFocus>();
-        input_focus.set(entity);
-    }
 
     #[derive(Component, Default)]
     struct GatherKeyboardEvents(String);
@@ -401,14 +391,16 @@ mod tests {
         }
     }
 
-    const KEY_A_EVENT: KeyboardInput = KeyboardInput {
-        key_code: KeyCode::KeyA,
-        logical_key: Key::Character(SmolStr::new_static("A")),
-        state: ButtonState::Pressed,
-        text: Some(SmolStr::new_static("A")),
-        repeat: false,
-        window: Entity::PLACEHOLDER,
-    };
+    fn key_a_event() -> KeyboardInput {
+        KeyboardInput {
+            key_code: KeyCode::KeyA,
+            logical_key: Key::Character("A".into()),
+            state: ButtonState::Pressed,
+            text: Some("A".into()),
+            repeat: false,
+            window: Entity::PLACEHOLDER,
+        }
+    }
 
     #[test]
     fn test_no_panics_if_resource_missing() {
@@ -438,6 +430,55 @@ mod tests {
             .unwrap();
     }
 
+    #[test]
+    fn initial_focus_unset_if_no_primary_window() {
+        let mut app = App::new();
+        app.add_plugins((InputPlugin, InputDispatchPlugin));
+
+        app.update();
+
+        assert_eq!(app.world().resource::<InputFocus>().0, None);
+    }
+
+    #[test]
+    fn initial_focus_set_to_primary_window() {
+        let mut app = App::new();
+        app.add_plugins((InputPlugin, InputDispatchPlugin));
+
+        let entity_window = app
+            .world_mut()
+            .spawn((Window::default(), PrimaryWindow))
+            .id();
+        app.update();
+
+        assert_eq!(app.world().resource::<InputFocus>().0, Some(entity_window));
+    }
+
+    #[test]
+    fn initial_focus_not_overridden() {
+        let mut app = App::new();
+        app.add_plugins((InputPlugin, InputDispatchPlugin));
+
+        app.world_mut().spawn((Window::default(), PrimaryWindow));
+
+        app.add_systems(Startup, |mut commands: Commands| {
+            commands.spawn(AutoFocus);
+        });
+
+        app.update();
+
+        let autofocus_entity = app
+            .world_mut()
+            .query_filtered::<Entity, With<AutoFocus>>()
+            .single(app.world())
+            .unwrap();
+
+        assert_eq!(
+            app.world().resource::<InputFocus>().0,
+            Some(autofocus_entity)
+        );
+    }
+
     #[test]
     fn test_keyboard_events() {
         fn get_gathered(app: &App, entity: Entity) -> &str {
@@ -454,18 +495,14 @@ mod tests {
         app.add_plugins((InputPlugin, InputDispatchPlugin))
             .add_observer(gather_keyboard_events);
 
-        let window = Window {
-            resolution: WindowResolution::new(800., 600.),
-            ..Default::default()
-        };
-        app.world_mut().spawn((window, PrimaryWindow));
+        app.world_mut().spawn((Window::default(), PrimaryWindow));
 
         // Run the world for a single frame to set up the initial focus
         app.update();
 
         let entity_a = app
             .world_mut()
-            .spawn((GatherKeyboardEvents::default(), SetFocusOnAdd))
+            .spawn((GatherKeyboardEvents::default(), AutoFocus))
             .id();
 
         let child_of_b = app
@@ -487,7 +524,7 @@ mod tests {
         assert!(!app.world().is_focus_visible(child_of_b));
 
         // entity_a should receive this event
-        app.world_mut().send_event(KEY_A_EVENT);
+        app.world_mut().send_event(key_a_event());
         app.update();
 
         assert_eq!(get_gathered(&app, entity_a), "A");
@@ -500,7 +537,7 @@ mod tests {
         assert!(!app.world().is_focus_visible(entity_a));
 
         // This event should be lost
-        app.world_mut().send_event(KEY_A_EVENT);
+        app.world_mut().send_event(key_a_event());
         app.update();
 
         assert_eq!(get_gathered(&app, entity_a), "A");
@@ -520,7 +557,8 @@ mod tests {
         assert!(app.world().is_focus_within(entity_b));
 
         // These events should be received by entity_b and child_of_b
-        app.world_mut().send_event_batch([KEY_A_EVENT; 4]);
+        app.world_mut()
+            .send_event_batch(core::iter::repeat_n(key_a_event(), 4));
         app.update();
 
         assert_eq!(get_gathered(&app, entity_a), "A");