From 7d8df1d8d4dcb91aa1adbfe3c49b9e9c35eee0fb Mon Sep 17 00:00:00 2001
From: Dean Herbert <pe@ppy.sh>
Date: Wed, 4 Jun 2025 17:30:57 +0900
Subject: [PATCH] Add back test covering immediate exit song select

---
 .../Navigation/TestSceneScreenNavigation.cs      |  9 +--------
 .../Navigation/TestSceneSongSelectNavigation.cs  | 16 ++++++++++++++++
 osu.Game/OsuGame.cs                              |  3 +++
 3 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/osu.Game.Tests/Visual/Navigation/TestSceneScreenNavigation.cs b/osu.Game.Tests/Visual/Navigation/TestSceneScreenNavigation.cs
index e2b2e41456..8ba914c05f 100644
--- a/osu.Game.Tests/Visual/Navigation/TestSceneScreenNavigation.cs
+++ b/osu.Game.Tests/Visual/Navigation/TestSceneScreenNavigation.cs
@@ -793,15 +793,8 @@ namespace osu.Game.Tests.Visual.Navigation
         {
             AddStep("push song select", () => Game.ScreenStack.Push(new TestPlaySongSelect()));
             AddStep("press back button", () => Game.ChildrenOfType<BackButton>().First().Action!.Invoke());
-            AddWaitStep("wait two frames", 2);
-        }
 
-        [Test]
-        public void TestPushSongSelectAndPressBackButtonImmediatelyV2()
-        {
-            AddStep("push song select", () => Game.ScreenStack.Push(new TestPlaySongSelect()));
-            AddStep("press back button", () => Game.ChildrenOfType<BackButton>().First().Action!.Invoke());
-            AddWaitStep("wait two frames", 2);
+            ConfirmAtMainMenu();
         }
 
         [Test]
diff --git a/osu.Game.Tests/Visual/Navigation/TestSceneSongSelectNavigation.cs b/osu.Game.Tests/Visual/Navigation/TestSceneSongSelectNavigation.cs
index f369a52ae7..264f09f6b9 100644
--- a/osu.Game.Tests/Visual/Navigation/TestSceneSongSelectNavigation.cs
+++ b/osu.Game.Tests/Visual/Navigation/TestSceneSongSelectNavigation.cs
@@ -1,11 +1,14 @@
 // Copyright (c) ppy Pty Ltd <contact@ppy.sh>. Licensed under the MIT Licence.
 // See the LICENCE file in the repository root for full licence text.
 
+using System.Linq;
 using NUnit.Framework;
 using osu.Framework.Allocation;
 using osu.Framework.Extensions;
+using osu.Framework.Testing;
 using osu.Game.Beatmaps;
 using osu.Game.Overlays;
+using osu.Game.Screens.Footer;
 using osu.Game.Screens.Play;
 using osu.Game.Screens.SelectV2;
 using osu.Game.Tests.Beatmaps.IO;
@@ -19,6 +22,19 @@ namespace osu.Game.Tests.Visual.Navigation
     /// </summary>
     public class TestSceneSongSelectNavigation : OsuGameTestScene
     {
+        [Test]
+        public void TestPushSongSelectAndPressBackButtonImmediately()
+        {
+            AddStep("push song select", () => Game.ScreenStack.Push(new SoloSongSelect()));
+
+            // TODO: without this step, a critical bug will be hit, see inline comment in `OsuGame.handleBackButton`.
+            AddUntilStep("Wait for song select", () => Game.ScreenStack.CurrentScreen is SoloSongSelect select && select.IsLoaded);
+
+            AddStep("press back button", () => Game.ChildrenOfType<ScreenBackButton>().First().Action!.Invoke());
+
+            ConfirmAtMainMenu();
+        }
+
         [TestCase(true)]
         [TestCase(false)]
         public void TestSongContinuesAfterExitPlayer(bool withUserPause)
diff --git a/osu.Game/OsuGame.cs b/osu.Game/OsuGame.cs
index 32ffc52be1..628d9d990c 100644
--- a/osu.Game/OsuGame.cs
+++ b/osu.Game/OsuGame.cs
@@ -1306,6 +1306,9 @@ namespace osu.Game
 
         private void handleBackButton()
         {
+            // TODO: this is SUPER SUPER bad.
+            // It can potentially exit the wrong screen if screens are not loaded yet.
+            // ScreenFooter / ScreenBackButton should be aware of which screen it is currently being handled by.
             if (!(ScreenStack.CurrentScreen is IOsuScreen currentScreen)) return;
 
             if (!((Drawable)currentScreen).IsLoaded || (currentScreen.AllowUserExit && !currentScreen.OnBackButton())) ScreenStack.Exit();