From 203691b1c776728fb48564c51bef60a20bc0d499 Mon Sep 17 00:00:00 2001
From: Dean Herbert <pe@ppy.sh>
Date: Tue, 29 May 2018 13:48:27 +0900
Subject: [PATCH] Add import rollback test

---
 .../Beatmaps/IO/ImportBeatmapTest.cs          | 75 +++++++++++++++++--
 1 file changed, 68 insertions(+), 7 deletions(-)

diff --git a/osu.Game.Tests/Beatmaps/IO/ImportBeatmapTest.cs b/osu.Game.Tests/Beatmaps/IO/ImportBeatmapTest.cs
index 38fd6e649b..586217a05f 100644
--- a/osu.Game.Tests/Beatmaps/IO/ImportBeatmapTest.cs
+++ b/osu.Game.Tests/Beatmaps/IO/ImportBeatmapTest.cs
@@ -12,6 +12,7 @@ using osu.Framework.Platform;
 using osu.Game.IPC;
 using osu.Framework.Allocation;
 using osu.Game.Beatmaps;
+using SharpCompress.Archives.Zip;
 
 namespace osu.Game.Tests.Beatmaps.IO
 {
@@ -77,8 +78,69 @@ namespace osu.Game.Tests.Beatmaps.IO
 
                     var manager = osu.Dependencies.Get<BeatmapManager>();
 
-                    Assert.IsTrue(manager.GetAllUsableBeatmapSets().Count == 1);
-                    Assert.IsTrue(manager.QueryBeatmapSets(_ => true).ToList().Count == 1);
+                    Assert.AreEqual(1, manager.GetAllUsableBeatmapSets().Count);
+                    Assert.AreEqual(1, manager.QueryBeatmapSets(_ => true).ToList().Count);
+                }
+                finally
+                {
+                    host.Exit();
+                }
+            }
+        }
+
+        [Test]
+        public void TestRollbackOnFailure()
+        {
+            //unfortunately for the time being we need to reference osu.Framework.Desktop for a game host here.
+            using (HeadlessGameHost host = new CleanRunHeadlessGameHost("TestRollbackOnFailure"))
+            {
+                try
+                {
+                    var osu = loadOsu(host);
+                    var manager = osu.Dependencies.Get<BeatmapManager>();
+
+                    int fireCount = 0;
+
+                    // ReSharper disable once AccessToModifiedClosure
+                    manager.ItemAdded += _ => fireCount++;
+                    manager.ItemRemoved += _ => fireCount++;
+
+                    var imported = loadOszIntoOsu(osu);
+
+                    Assert.AreEqual(0, fireCount -= 1);
+
+                    imported.Hash += "-changed";
+                    manager.Update(imported);
+
+                    Assert.AreEqual(0, fireCount -= 2);
+
+                    var breakTemp = createTemporaryBeatmap();
+
+                    MemoryStream brokenOsu = new MemoryStream(new byte[] { 1, 3, 3, 7 });
+                    MemoryStream brokenOsz = new MemoryStream(File.ReadAllBytes(breakTemp));
+
+                    File.Delete(breakTemp);
+
+                    using (var outStream = File.Open(breakTemp, FileMode.CreateNew))
+                    using (var zip = ZipArchive.Open(brokenOsz))
+                    {
+                        zip.AddEntry("broken.osu", brokenOsu, false);
+                        zip.SaveTo(outStream, SharpCompress.Common.CompressionType.Deflate);
+                    }
+
+                    Assert.AreEqual(1, manager.GetAllUsableBeatmapSets().Count);
+                    Assert.AreEqual(1, manager.QueryBeatmapSets(_ => true).ToList().Count);
+                    Assert.AreEqual(12, manager.QueryBeatmaps(_ => true).ToList().Count);
+
+                    // this will trigger purging of the existing beatmap (online set id match) but should rollback due to broken osu.
+                    manager.Import(breakTemp);
+
+                    // no events should be fired in the case of a rollback.
+                    Assert.AreEqual(0, fireCount);
+
+                    Assert.AreEqual(1, manager.GetAllUsableBeatmapSets().Count);
+                    Assert.AreEqual(1, manager.QueryBeatmapSets(_ => true).ToList().Count);
+                    Assert.AreEqual(12, manager.QueryBeatmaps(_ => true).ToList().Count);
                 }
                 finally
                 {
@@ -100,18 +162,17 @@ namespace osu.Game.Tests.Beatmaps.IO
 
                     var imported = loadOszIntoOsu(osu);
 
-                    //var change = manager.QueryBeatmapSets(_ => true).First();
                     imported.Hash += "-changed";
                     manager.Update(imported);
 
                     var importedSecondTime = loadOszIntoOsu(osu);
 
-                    // check the newly "imported" beatmap is actually just the restored previous import. since it matches hash.
                     Assert.IsTrue(imported.ID != importedSecondTime.ID);
                     Assert.IsTrue(imported.Beatmaps.First().ID < importedSecondTime.Beatmaps.First().ID);
 
-                    Assert.IsTrue(manager.GetAllUsableBeatmapSets().Count == 1);
-                    Assert.IsTrue(manager.QueryBeatmapSets(_ => true).ToList().Count == 1);
+                    // only one beatmap will exist as the online set ID matched, causing purging of the first import.
+                    Assert.AreEqual(1, manager.GetAllUsableBeatmapSets().Count);
+                    Assert.AreEqual(1, manager.QueryBeatmapSets(_ => true).ToList().Count);
                 }
                 finally
                 {
@@ -231,7 +292,7 @@ namespace osu.Game.Tests.Beatmaps.IO
             manager.Delete(imported);
 
             Assert.IsTrue(manager.GetAllUsableBeatmapSets().Count == 0);
-            Assert.IsTrue(manager.QueryBeatmapSets(_ => true).ToList().Count == 1);
+            Assert.AreEqual(1, manager.QueryBeatmapSets(_ => true).ToList().Count);
             Assert.IsTrue(manager.QueryBeatmapSets(_ => true).First().DeletePending);
         }