Test staying on same difficulty due to unsaved changes

osu.Game.Tests/Visual/Editing/TestSceneDifficultySwitching.cs

@@ -8,6 +8,7 @@ using osu.Framework.Allocation;
 using osu.Framework.Testing;
 using osu.Game.Beatmaps;
 using osu.Game.Graphics.UserInterface;
+using osu.Game.Overlays.Dialog;
 using osu.Game.Screens.Edit;
 using osu.Game.Screens.Edit.Components.Menus;
 using osu.Game.Tests.Beatmaps.IO;
@@ -35,8 +36,9 @@ namespace osu.Game.Tests.Visual.Editing
             AddStep("set current beatmap", () => Beatmap.Value = beatmaps.GetWorkingBeatmap(importedBeatmapSet.Beatmaps.First()));
             AddStep("push loader", () => Stack.Push(new EditorLoader()));
 
-            AddUntilStep("wait for editor to load", () => Stack.CurrentScreen is Editor);
+            AddUntilStep("wait for editor push", () => Stack.CurrentScreen is Editor);
             AddStep("store editor", () => editor = (Editor)Stack.CurrentScreen);
+            AddUntilStep("wait for editor to load", () => editor.IsLoaded);
         }
 
         [Test]
@@ -53,6 +55,39 @@ namespace osu.Game.Tests.Visual.Editing
             AddAssert("stack empty", () => Stack.CurrentScreen == null);
         }
 
+        [Test]
+        public void TestPreventSwitchDueToUnsavedChanges()
+        {
+            BeatmapInfo targetDifficulty = null;
+            PromptForSaveDialog saveDialog = null;
+
+            AddStep("remove first hitobject", () =>
+            {
+                var editorBeatmap = editor.ChildrenOfType<EditorBeatmap>().Single();
+                editorBeatmap.RemoveAt(0);
+            });
+
+            AddStep("set target difficulty", () => targetDifficulty = importedBeatmapSet.Beatmaps.Last(beatmap => !beatmap.Equals(Beatmap.Value.BeatmapInfo)));
+            switchToDifficulty(() => targetDifficulty);
+
+            AddUntilStep("prompt for save dialog shown", () =>
+            {
+                saveDialog = this.ChildrenOfType<PromptForSaveDialog>().Single();
+                return saveDialog != null;
+            });
+            AddStep("continue editing", () =>
+            {
+                var continueButton = saveDialog.ChildrenOfType<PopupDialogCancelButton>().Last();
+                continueButton.TriggerClick();
+            });
+
+            confirmEditingBeatmap(() => importedBeatmapSet.Beatmaps.First());
+
+            AddRepeatStep("exit editor forcefully", () => Stack.Exit(), 2);
+            // ensure editor loader didn't resume.
+            AddAssert("stack empty", () => Stack.CurrentScreen == null);
+        }
+
         private void switchToDifficulty(Func<BeatmapInfo> difficulty)
         {
             AddUntilStep("wait for menubar to load", () => editor.ChildrenOfType<EditorMenuBar>().Any());

osu.Game/Screens/Edit/Editor.cs

@@ -498,7 +498,7 @@ namespace osu.Game.Screens.Edit
 
                 if (isNewBeatmap || HasUnsavedChanges)
                 {
-                    dialogOverlay?.Push(new PromptForSaveDialog(confirmExit, confirmExitWithSave));
+                    dialogOverlay?.Push(new PromptForSaveDialog(confirmExit, confirmExitWithSave, cancelPendingDifficultySwitch));
                     return true;
                 }
             }
@@ -737,7 +737,16 @@ namespace osu.Game.Screens.Edit
 
             loader.ValidForResume = true;
             this.Exit();
-            loader.PushEditor(beatmapInfo);
+            loader.ScheduleDifficultySwitch(beatmapInfo);
+        }
+
+        private void cancelPendingDifficultySwitch()
+        {
+            if (loader == null)
+                return;
+
+            loader.ValidForResume = false;
+            loader.CancelDifficultySwitch();
         }
 
         public double SnapTime(double time, double? referenceTime) => editorBeatmap.SnapTime(time, referenceTime);

osu.Game/Screens/Edit/EditorLoader.cs

@@ -4,6 +4,7 @@
 using JetBrains.Annotations;
 using osu.Framework.Allocation;
 using osu.Framework.Screens;
+using osu.Framework.Threading;
 using osu.Game.Beatmaps;
 using osu.Game.Screens.Play;
 
@@ -18,19 +19,31 @@ namespace osu.Game.Screens.Edit
         [Resolved]
         private BeatmapManager beatmapManager { get; set; }
 
+        [CanBeNull]
+        private ScheduledDelegate scheduledDifficultySwitch;
+
         public override void OnEntering(IScreen last)
         {
             base.OnEntering(last);
-            PushEditor();
+            pushEditor();
         }
 
-        public void PushEditor([CanBeNull] BeatmapInfo beatmapInfo = null) => Schedule(() =>
+        private void pushEditor()
         {
-            if (beatmapInfo != null)
-                Beatmap.Value = beatmapManager.GetWorkingBeatmap(beatmapInfo);
-
             this.Push(new Editor(this));
             ValidForResume = false;
-        });
+        }
+
+        public void ScheduleDifficultySwitch(BeatmapInfo beatmapInfo)
+        {
+            CancelDifficultySwitch();
+            scheduledDifficultySwitch = Schedule(() =>
+            {
+                Beatmap.Value = beatmapManager.GetWorkingBeatmap(beatmapInfo);
+                pushEditor();
+            });
+        }
+
+        public void CancelDifficultySwitch() => scheduledDifficultySwitch?.Cancel();
     }
 }

osu.Game/Screens/Edit/PromptForSaveDialog.cs

@@ -9,7 +9,7 @@ namespace osu.Game.Screens.Edit
 {
     public class PromptForSaveDialog : PopupDialog
     {
-        public PromptForSaveDialog(Action exit, Action saveAndExit)
+        public PromptForSaveDialog(Action exit, Action saveAndExit, Action cancel)
         {
             HeaderText = "Did you want to save your changes?";
 
@@ -30,6 +30,7 @@ namespace osu.Game.Screens.Edit
                 new PopupDialogCancelButton
                 {
                     Text = @"Oops, continue editing",
+                    Action = cancel
                 },
             };
         }