Add retries to DatabaseWriter on partial success, and add unit tests

KustoSchemaTools.Tests/DefaultDatabaseWriterTests.cs

@@ -0,0 +1,230 @@
+using Kusto.Data;
+using KustoSchemaTools.Changes;
+using KustoSchemaTools.Model;
+using KustoSchemaTools.Parser;
+using KustoSchemaTools.Parser.KustoWriter;
+using Microsoft.Extensions.Logging;
+using Moq;
+
+namespace KustoSchemaTools.Tests
+{
+    public class DefaultDatabaseWriterTests
+    {
+        [Fact]
+        public async Task UpdatePrimary_ShouldRetryUntilChangesSucceed()
+        {
+            // Arrange
+            var sourceDb = new Database { Name = "SourceDb" };
+            var targetDb = new Database { Name = "TargetDb" };
+
+            var logger = new ConsoleLogger();
+            var mockClient = new Mock<KustoClient>(MockBehavior.Loose, "test");
+
+            var writer = new TestableDefaultDatabaseWriter();
+
+            // Create mock changes for the test
+            writer.Setup(true, false, false);
+            writer.Setup(true, true);
+
+            // Act
+            await writer.WriteAsync(sourceDb, targetDb, mockClient.Object, logger);
+            Assert.Equal(2, writer.GenerateChangesCallCount);
+            Assert.Equal(2, writer.ApplyChangesToDatabaseCallCount);
+            Assert.Equal(0, writer.ExceptionCount);
+        }
+
+        [Fact]
+        public async Task UpdatePrimary_NoRetryIfAllSucces()
+        {
+            // Arrange
+            var sourceDb = new Database { Name = "SourceDb" };
+            var targetDb = new Database { Name = "TargetDb" };
+            var logger = new ConsoleLogger();
+            var mockClient = new Mock<KustoClient>(MockBehavior.Loose, "test");
+            var writer = new TestableDefaultDatabaseWriter();
+
+            // Create mock changesets for the test
+            writer.Setup(true, true, true);
+            writer.Setup(true, true, true);
+
+            // all success
+            await writer.WriteAsync(sourceDb, targetDb, mockClient.Object, logger);
+            Assert.Equal(1, writer.GenerateChangesCallCount);
+            Assert.Equal(1, writer.ApplyChangesToDatabaseCallCount);
+            Assert.Equal(0, writer.ExceptionCount);
+        }
+
+        [Fact]
+        public async Task UpdatePrimary_NoRetryIfAllFail()
+        {
+            // Arrange
+            var sourceDb = new Database { Name = "SourceDb" };
+            var targetDb = new Database { Name = "TargetDb" };
+            var logger = new ConsoleLogger();
+            var mockClient = new Mock<KustoClient>(MockBehavior.Loose, "test");
+            var writer = new TestableDefaultDatabaseWriter();
+
+            // Create mock changesets for the test
+            writer.Setup(false, false, false);
+            writer.Setup(false, false, false);
+
+            await writer.WriteAsync(sourceDb, targetDb, mockClient.Object, logger);
+            Assert.Equal(1, writer.GenerateChangesCallCount);
+            Assert.Equal(1, writer.ApplyChangesToDatabaseCallCount);
+            Assert.Equal(1, writer.ExceptionCount);
+        }
+
+
+        [Fact]
+        public async Task UpdatePrimary_ShouldRetryUntilNoSuccessfulChanges()
+        {
+            // Arrange
+            var sourceDb = new Database { Name = "SourceDb" };
+            var targetDb = new Database { Name = "TargetDb" };
+            var logger = new ConsoleLogger();
+            var mockClient = new Mock<KustoClient>(MockBehavior.Loose, "test");
+            var writer = new TestableDefaultDatabaseWriter();
+
+            // Create mock changes for the test
+            writer.Setup(true, false, false);
+            writer.Setup(true, false);
+            writer.Setup(false);
+            writer.Setup(false); // one extra
+
+            // Act
+            await writer.WriteAsync(sourceDb, targetDb, mockClient.Object, logger);
+            Assert.Equal(3, writer.GenerateChangesCallCount);
+            Assert.Equal(3, writer.ApplyChangesToDatabaseCallCount);
+            Assert.Equal(1, writer.ExceptionCount);
+        }
+
+        #region Helper Methods
+
+
+        static IChange CreateMockChange(string name)
+        {
+            var mockChange = new Mock<IChange>();
+
+            var scripts = new List<DatabaseScriptContainer> {
+                new DatabaseScriptContainer(
+                    new DatabaseScript(name, 0),
+                    "TestKind"
+                )
+            };
+
+            mockChange.Setup(c => c.Scripts).Returns(scripts);
+            return mockChange.Object;
+        }
+
+        static ScriptExecuteCommandResult CreateMockResult(IChange change, bool isSuccess)
+        {
+            return new ScriptExecuteCommandResult {
+                OperationId = Guid.NewGuid(),
+                CommandType = "Script",
+                Result = isSuccess ? "Completed" : "Failed",
+                CommandText = change.Scripts.First().Script.Text,
+                Reason = isSuccess ? null : "Test failure"
+            };
+        }
+
+        #endregion
+
+        #region Test Support Classes
+
+        /// <summary>
+        /// A testable version of DefaultDatabaseWriter that we can use to track calls and control behavior
+        /// </summary>
+        private class TestableDefaultDatabaseWriter : DefaultDatabaseWriter
+        {
+            public Dictionary<IChange, ScriptExecuteCommandResult> ResultsCache { get; } = new Dictionary<IChange, ScriptExecuteCommandResult>();
+            private readonly IList<List<IChange>> _generateChangesResults = new List<List<IChange>>();
+            public int GenerateChangesCallCount { get; private set; } = 0;
+            public int ApplyChangesToDatabaseCallCount { get; private set; } = 0;
+            public int ExceptionCount { get; private set; } = 0;
+
+            public void ResetCounts()
+            {
+                GenerateChangesCallCount = ApplyChangesToDatabaseCallCount = ExceptionCount = 0;
+            }
+
+            public void Setup(params bool[] results)
+            {
+                var changeSet = new List<IChange> { };
+                foreach (var result in results)
+                {
+                    var change = CreateMockChange($"change");
+                    var scriptExecuteResult = CreateMockResult(change, result);
+                    ResultsCache.Add(change, scriptExecuteResult);
+                    changeSet.Add(change);
+                }
+                _generateChangesResults.Add(changeSet);
+            }
+
+            // Override the base method to use our predetermined results
+            internal override List<IChange> GenerateChanges(Database targetDb, Database sourceDb, ILogger logger)
+            {
+                if (GenerateChangesCallCount >= _generateChangesResults.Count)
+                {
+                    throw new InvalidOperationException("No more predefined change sets available.");
+                }
+
+                var changes = _generateChangesResults.ElementAt(GenerateChangesCallCount);
+                GenerateChangesCallCount++;
+                return changes;
+            }
+
+            // Override the protected method to return our predetermined results
+            internal override Task<List<ScriptExecuteCommandResult>> ApplyChangesToDatabase(string databaseName, List<IChange> changes, KustoClient client, ILogger logger)
+            {
+                ApplyChangesToDatabaseCallCount++;
+                var results = new List<ScriptExecuteCommandResult>();
+                foreach (var c in changes) {
+                    var result = ResultsCache[c];
+                    results.Add(result);
+                }
+                return Task.FromResult(results);
+            }
+
+            // Expose the internal UpdatePrimary method for testing
+            public Task<List<ScriptExecuteCommandResult>> TestUpdatePrimary(Database sourceDb, Database targetDb, KustoClient client, ILogger logger)
+            {
+                return UpdatePrimary(sourceDb, targetDb, client, logger);
+            }
+
+            // Override WriteAsync to control exception throwing
+            public override async Task WriteAsync(Database sourceDb, Database targetDb, KustoClient client, ILogger logger)
+            {
+                try
+                {
+                    await base.WriteAsync(sourceDb, targetDb, client, logger);
+                }
+                catch (Exception)
+                {
+                    ExceptionCount++;
+                }
+            }
+        }
+
+        // Custom console logger implementation
+        private class ConsoleLogger : ILogger
+        {
+            public IDisposable? BeginScope<TState>(TState state) where TState : notnull => new DummyDisposable();
+            public bool IsEnabled(LogLevel logLevel) => true;
+
+            public void Log<TState>(LogLevel logLevel, EventId eventId, TState state, Exception? exception, Func<TState, Exception?, string> formatter)
+            {
+                Console.WriteLine($"[{logLevel}] {formatter(state, exception)}");
+                if (exception != null)
+                {
+                    Console.WriteLine($"Exception: {exception.Message}");
+                }
+            }
+
+            private class DummyDisposable : IDisposable
+            {
+                public void Dispose() { }
+            }
+        }
+        #endregion
+    }
+}