Fix Cosmos HTTP client lifecycle across sources, sinks and Operations

Extensions/Cosmos/Cosmos.DataTransfer.CosmosExtension.UnitTests/CertificateConfigurationTests.cs

@@ -51,7 +51,7 @@ public void CreateClient_WithDisableSslValidation_LogsWarningAndSetsCallback()
                 DisableSslValidation = true
             };
 
-            var client = CosmosExtensionServices.CreateClient(settings, "TestDisplay", loggerMock.Object);
+            using var client = CosmosExtensionServices.CreateClient(settings, "TestDisplay", loggerMock.Object);
 
             // Verify warning was logged when DisableSslValidation is true
             loggerMock.Verify(
@@ -79,7 +79,7 @@ public void CreateClient_WithoutDisableSslValidation_DoesNotLogWarning()
                 DisableSslValidation = false
             };
 
-            var client = CosmosExtensionServices.CreateClient(settings, "TestDisplay", loggerMock.Object);
+            using var client = CosmosExtensionServices.CreateClient(settings, "TestDisplay", loggerMock.Object);
 
             // Verify warning was NOT logged when DisableSslValidation is false
             loggerMock.Verify(
@@ -116,7 +116,8 @@ public void CreateClient_WithEnableNetHttpLogging_LogsWarningAndPreservesClientO
                 AllowBulkExecution = true
             };
 
-            using var client = CosmosExtensionServices.CreateClient(settings, "TestDisplay", loggerMock.Object);
+            using var client = CosmosExtensionServices.CreateClient(settings, "TestDisplay", loggerMock.Object,
+                null, (_, _) => new OfflineTokenCredential());
 
             loggerMock.Verify(
                 x => x.Log(

Extensions/Cosmos/Cosmos.DataTransfer.CosmosExtension.UnitTests/CosmosDataSourceExtensionTests.cs

@@ -3,7 +3,13 @@
 using Microsoft.Extensions.Logging.Abstractions;
 using Microsoft.Azure.Cosmos;
 using System.Net;
+using System.CommandLine;
+using System.CommandLine.Invocation;
+using System.ComponentModel.Composition.Hosting;
+using Cosmos.DataTransfer.Core;
 using Cosmos.DataTransfer.Interfaces;
+using Microsoft.Extensions.Configuration;
+using Moq;
 using Newtonsoft.Json;
 using Newtonsoft.Json.Linq;
 
@@ -272,6 +278,7 @@ public async Task Test_Json_RoundTrip_WithLongValues() {
         var (_, database, connectionString) = GetSetup();
         var config = TestHelpers.CreateConfig(new Dictionary<string,string>() {
             { "UseRbacAuth", "false" },
+            { "UseDefaultCredentials", "true" },
             { "ConnectionString", connectionString },
             { "ConnectionMode", "Gateway" },
             { "AllowBulkExecution", "false" },
@@ -305,6 +312,58 @@ await sinkExtension.WriteAsync(items.ToAsyncEnumerable(), config,
         await jsonSink.FormatDataAsync(fetched.ToAsyncEnumerable(), writer, config, NullLogger.Instance);
         var res = Encoding.UTF8.GetString(writer.ToArray());
         Assert.AreEqual("[{\"partitionKey\":\"pk\",\"id\":\"1\",\"long\":638676324052177500}]", res);
+
+        foreach (var preAuthenticate in new[] { false, true })
+        {
+            await RunSequentialOperationsAsync(connectionString, database.Id,
+                config["Container"]!, preAuthenticate);
+        }
+    }
+
+    private static async Task RunSequentialOperationsAsync(
+        string connectionString, string database, string sourceContainer, bool preAuthenticate)
+    {
+        var source = new CosmosDataSourceExtension();
+        var sink = new CosmosDataSinkExtension();
+        using var catalog = new CompositionContainer();
+        var loader = new Mock<IExtensionLoader>();
+        loader.Setup(l => l.BuildExtensionCatalog(It.IsAny<string>())).Returns(catalog);
+        loader.Setup(l => l.LoadExtensions<IDataSourceExtension>(catalog)).Returns([source]);
+        loader.Setup(l => l.LoadExtensions<IDataSinkExtension>(catalog)).Returns([sink]);
+        var values = new Dictionary<string, string?>();
+        foreach (var side in new[] { "SourceSettings", "SinkSettings" })
+        {
+            values[$"{side}:ConnectionString"] = connectionString;
+            values[$"{side}:Database"] = database;
+            values[$"{side}:ConnectionMode"] = "Gateway";
+            values[$"{side}:LimitToEndpoint"] = "true";
+            values[$"{side}:UseDefaultCredentials"] = "true";
+            values[$"{side}:PreAuthenticate"] = preAuthenticate.ToString();
+            values[$"{side}:AllowBulkExecution"] = "false";
+        }
+        values["SourceSettings:Container"] = sourceContainer;
+        values["SinkSettings:PartitionKeyPath"] = "/partitionKey";
+        for (var i = 0; i < 2; i++)
+        {
+            values[$"Operations:{i}:SinkSettings:Container"] = $"operation-{preAuthenticate}-{i}";
+        }
+        var configuration = new ConfigurationBuilder().AddInMemoryCollection(values).Build();
+        var command = new RunCommand.CommandHandler(loader.Object, configuration, NullLoggerFactory.Instance)
+        {
+            Source = source.DisplayName,
+            Sink = sink.DisplayName,
+        };
+        var context = new InvocationContext(new RootCommand().Parse(Array.Empty<string>()));
+        Assert.AreEqual(0, await command.InvokeAsync(context), "Sequential Cosmos Operations must succeed.");
+        for (var i = 0; i < 2; i++)
+        {
+            var outputConfig = new ConfigurationBuilder()
+                .AddConfiguration(configuration.GetSection("SinkSettings"))
+                .AddConfiguration(configuration.GetSection($"Operations:{i}:SinkSettings")).Build();
+            var output = await source.ReadAsync(outputConfig, NullLogger.Instance).ToArrayAsync();
+            Assert.AreEqual(1, output.Length);
+            Assert.AreEqual(638676324052177500L, output[0].GetValue("long"));
+        }
     }
 
     [TestMethod]