From 5b75060b94999e3a50a665a4e40cee3994162fa2 Mon Sep 17 00:00:00 2001
From: Dean Herbert <pe@ppy.sh>
Date: Mon, 10 Jun 2019 16:45:45 +0900
Subject: [PATCH] Add test for rollback logic correctly dereferencing files

---
 osu.Game.Tests/Beatmaps/IO/ImportBeatmapTest.cs | 13 +++++++++++++
 osu.Game/IO/FileStore.cs                        | 10 ++++++++++
 2 files changed, 23 insertions(+)

diff --git a/osu.Game.Tests/Beatmaps/IO/ImportBeatmapTest.cs b/osu.Game.Tests/Beatmaps/IO/ImportBeatmapTest.cs
index 4c9260f640..3f4f40781c 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 osu.Game.IO;
 using osu.Game.Tests.Resources;
 using SharpCompress.Archives.Zip;
 
@@ -97,6 +98,7 @@ namespace osu.Game.Tests.Beatmaps.IO
                 {
                     var osu = loadOsu(host);
                     var manager = osu.Dependencies.Get<BeatmapManager>();
+                    var files = osu.Dependencies.Get<FileStore>();
 
                     int fireCount = 0;
 
@@ -113,6 +115,12 @@ namespace osu.Game.Tests.Beatmaps.IO
 
                     Assert.AreEqual(0, fireCount -= 2);
 
+                    Assert.AreEqual(1, manager.GetAllUsableBeatmapSets().Count);
+                    Assert.AreEqual(1, manager.QueryBeatmapSets(_ => true).ToList().Count);
+                    Assert.AreEqual(12, manager.QueryBeatmaps(_ => true).ToList().Count);
+
+                    Assert.AreEqual(18, files.QueryFiles(_ => true).Count());
+
                     var breakTemp = TestResources.GetTestBeatmapForImport();
 
                     MemoryStream brokenOsu = new MemoryStream(new byte[] { 1, 3, 3, 7 });
@@ -131,6 +139,8 @@ namespace osu.Game.Tests.Beatmaps.IO
                     Assert.AreEqual(1, manager.QueryBeatmapSets(_ => true).ToList().Count);
                     Assert.AreEqual(12, manager.QueryBeatmaps(_ => true).ToList().Count);
 
+                    Assert.AreEqual(18, files.QueryFiles(_ => true).Count());
+
                     // this will trigger purging of the existing beatmap (online set id match) but should rollback due to broken osu.
                     await manager.Import(breakTemp);
 
@@ -140,6 +150,9 @@ namespace osu.Game.Tests.Beatmaps.IO
                     Assert.AreEqual(1, manager.GetAllUsableBeatmapSets().Count);
                     Assert.AreEqual(1, manager.QueryBeatmapSets(_ => true).ToList().Count);
                     Assert.AreEqual(12, manager.QueryBeatmaps(_ => true).ToList().Count);
+
+                    Assert.AreEqual(18, files.QueryFiles(_ => true).Count());
+                    Assert.AreEqual(18, files.QueryFiles(f => f.ReferenceCount == 1).Count());
                 }
                 finally
                 {
diff --git a/osu.Game/IO/FileStore.cs b/osu.Game/IO/FileStore.cs
index 458f8964f9..370d6786f5 100644
--- a/osu.Game/IO/FileStore.cs
+++ b/osu.Game/IO/FileStore.cs
@@ -2,8 +2,11 @@
 // See the LICENCE file in the repository root for full licence text.
 
 using System;
+using System.Collections.Generic;
 using System.IO;
 using System.Linq;
+using System.Linq.Expressions;
+using Microsoft.EntityFrameworkCore;
 using osu.Framework.Extensions;
 using osu.Framework.IO.Stores;
 using osu.Framework.Logging;
@@ -27,6 +30,13 @@ namespace osu.Game.IO
             Store = new StorageBackedResourceStore(Storage);
         }
 
+        /// <summary>
+        /// Perform a lookup query on available <see cref="FileInfo"/>s.
+        /// </summary>
+        /// <param name="query">The query.</param>
+        /// <returns>Results from the provided query.</returns>
+        public IEnumerable<FileInfo> QueryFiles(Expression<Func<FileInfo, bool>> query) => ContextFactory.Get().Set<FileInfo>().AsNoTracking().Where(f => f.ReferenceCount > 0).Where(query);
+
         public FileInfo Add(Stream data, bool reference = true)
         {
             using (var usage = ContextFactory.GetForWrite())