RavenDB-21646 : address PR comments

aviv86 · Nov 9, 2023 · ad40650 · ad40650
1 parent 469816e
commit ad40650
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 52 deletions.
diff --git a/test/SlowTests/Server/Replication/ReplicationCleanTombstones.cs b/test/SlowTests/Server/Replication/ReplicationCleanTombstones.cs
@@ -15,7 +15,6 @@
 using Raven.Client.Json.Serialization;
 using Raven.Client.ServerWide;
 using Raven.Server;
-using Raven.Server.Documents;
 using Raven.Server.Documents.ETL;
 using Raven.Server.Documents.Replication;
 using Raven.Server.Documents.Replication.ReplicationItems;
@@ -815,7 +814,7 @@ public async Task CanBackupAndRestoreTombstonesWithSameId(Options options)
                     session.SaveChanges();
                 }
 
-                var backupTaskId = await CreateAndRunBackup(store, options.DatabaseMode, backupPath);
+                var backupTaskId = await Backup.CreateAndRunBackupAsync(store, options.DatabaseMode, backupPath);
 
                 using (var session = store.OpenSession())
                 {
@@ -841,15 +840,19 @@ public async Task CanBackupAndRestoreTombstonesWithSameId(Options options)
                     session.SaveChanges();
                 }
 
-                var configuration = await RunBackupAndCreateRestoreConfiguration(store, options.DatabaseMode, backupTaskId, backupPath);
+                var configuration = await Backup.RunBackupAndCreateRestoreConfigurationAsync(store, options.DatabaseMode, backupTaskId, backupPath);
 
                 using (Backup.RestoreDatabase(store, configuration))
                 {
                     var stats = await GetDatabaseStatisticsAsync(store, configuration.DatabaseName);
                     Assert.Equal(1, stats.CountOfDocuments); // the marker
                     Assert.Equal(2, stats.CountOfTombstones);
 
-                    var storage = await GetDocumentStorageFor(store, options.DatabaseMode, configuration.DatabaseName, id);
+                    var dbName = options.DatabaseMode == RavenDatabaseMode.Single
+                        ? configuration.DatabaseName
+                        : await Sharding.GetShardDatabaseNameForDocAsync(store, id, configuration.DatabaseName);
+                    var storage = (await GetDocumentDatabaseInstanceForAsync(dbName)).DocumentsStorage;
+
                     using (storage.ContextPool.AllocateOperationContext(out DocumentsOperationContext ctx))
                     using (ctx.OpenReadTransaction())
                     {
@@ -862,53 +865,5 @@ public async Task CanBackupAndRestoreTombstonesWithSameId(Options options)
                 }
             }
         }
-
-        private async Task<long> CreateAndRunBackup(IDocumentStore store, RavenDatabaseMode mode, string backupPath)
-        {
-            var config = Backup.CreateBackupConfiguration(backupPath);
-
-            if (mode == RavenDatabaseMode.Single)
-                return await Backup.UpdateConfigAndRunBackupAsync(Server, config, store);
-
-            var backupCompleted = await Sharding.Backup.WaitForBackupToComplete(store);
-            var backupTaskId = await Sharding.Backup.UpdateConfigurationAndRunBackupAsync(Server, store, config);
-            Assert.True(WaitHandle.WaitAll(backupCompleted, TimeSpan.FromSeconds(10)));
-
-            return backupTaskId;
-        }
-
-        private async Task<RestoreBackupConfiguration> RunBackupAndCreateRestoreConfiguration(DocumentStore store, RavenDatabaseMode mode, long backupTaskId, string backupPath)
-        {
-            var databaseName = GetDatabaseName();
-            var configuration = new RestoreBackupConfiguration { DatabaseName = databaseName };
-
-            if (mode == RavenDatabaseMode.Single)
-            {
-                await Backup.RunBackupAndReturnStatusAsync(Server, backupTaskId, store, isFullBackup: false);
-                configuration.BackupLocation = Directory.GetDirectories(backupPath).First();
-            }
-            else
-            {
-                var backupCompleted = await Sharding.Backup.WaitForBackupToComplete(store);
-                await Sharding.Backup.RunBackupAsync(store.Database, backupTaskId, isFullBackup: false, servers: new List<RavenServer> { Server });
-                Assert.True(WaitHandle.WaitAll(backupCompleted, TimeSpan.FromSeconds(10)));
-
-                var dirs = Directory.GetDirectories(backupPath);
-                var sharding = await Sharding.GetShardingConfigurationAsync(store);
-                var settings = Sharding.Backup.GenerateShardRestoreSettings(dirs, sharding);
-                configuration.ShardRestoreSettings = settings;
-            }
-
-            return configuration;
-        }
-
-        private async Task<DocumentsStorage> GetDocumentStorageFor(IDocumentStore store, RavenDatabaseMode mode, string databaseName, string id)
-        {
-            var db = await Server.ServerStore.DatabasesLandlord.TryGetOrCreateResourceStore(
-                mode == RavenDatabaseMode.Single
-                    ? databaseName
-                    : await Sharding.GetShardDatabaseNameForDocAsync(store, id, databaseName));
-            return db.DocumentsStorage;
-        }
     }
 }
diff --git a/test/Tests.Infrastructure/RavenTestBase.Backup.cs b/test/Tests.Infrastructure/RavenTestBase.Backup.cs
@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.IO;
 using System.Linq;
 using System.Text;
 using System.Threading;
@@ -17,6 +18,7 @@
 using Raven.Server.ServerWide.Context;
 using Raven.Server.Utils;
 using Raven.Tests.Core.Utils.Entities;
+using Tests.Infrastructure;
 using Xunit;
 
 namespace FastTests
@@ -270,6 +272,45 @@ public async Task<long> GetBackupOperationIdAsync(IDocumentStore store, long tas
                 return result.Status.LastOperationId.Value;
             }
 
+            public async Task<long> CreateAndRunBackupAsync(IDocumentStore store, RavenDatabaseMode mode, string backupPath)
+            {
+                var config = CreateBackupConfiguration(backupPath);
+
+                if (mode == RavenDatabaseMode.Single)
+                    return await UpdateConfigAndRunBackupAsync(_parent.Server, config, store);
+
+                var backupCompleted = await _parent.Sharding.Backup.WaitForBackupToComplete(store);
+                var backupTaskId = await _parent.Sharding.Backup.UpdateConfigurationAndRunBackupAsync(_parent.Server, store, config);
+                Assert.True(WaitHandle.WaitAll(backupCompleted, TimeSpan.FromSeconds(10)));
+
+                return backupTaskId;
+            }
+
+            public async Task<RestoreBackupConfiguration> RunBackupAndCreateRestoreConfigurationAsync(DocumentStore store, RavenDatabaseMode mode, long backupTaskId, string backupPath)
+            {
+                var databaseName = _parent.GetDatabaseName();
+                var configuration = new RestoreBackupConfiguration { DatabaseName = databaseName };
+
+                if (mode == RavenDatabaseMode.Single)
+                {
+                    await RunBackupAndReturnStatusAsync(_parent.Server, backupTaskId, store, isFullBackup: false);
+                    configuration.BackupLocation = Directory.GetDirectories(backupPath).First();
+                }
+                else
+                {
+                    var backupCompleted = await _parent.Sharding.Backup.WaitForBackupToComplete(store);
+                    await _parent.Sharding.Backup.RunBackupAsync(store.Database, backupTaskId, isFullBackup: false, servers: new List<RavenServer> { _parent.Server });
+                    Assert.True(WaitHandle.WaitAll(backupCompleted, TimeSpan.FromSeconds(10)));
+
+                    var dirs = Directory.GetDirectories(backupPath);
+                    var sharding = await _parent.Sharding.GetShardingConfigurationAsync(store);
+                    var settings = _parent.Sharding.Backup.GenerateShardRestoreSettings(dirs, sharding);
+                    configuration.ShardRestoreSettings = settings;
+                }
+
+                return configuration;
+            }
+
             internal static string PrintBackupStatus(PeriodicBackupStatus status)
             {
                 var sb = new StringBuilder();