From dbbae9205e3a8ebce63650bf816b0d095ea0de6c Mon Sep 17 00:00:00 2001
From: Iman Malik <49374113+imanvt@users.noreply.github.com>
Date: Thu, 9 Jun 2022 13:55:41 -0400
Subject: [PATCH] CosmosClientBuilder: Adds Warmup Option to Create/Initialize
 Containers (#3257)

* Add Warmup option for CosmosClientBuilder

* Add test method for new builder warmup option

* Public API Change

* rename method

* remove async/await

* update API

* API changes to remove async
---
 .../src/Fluent/CosmosClientBuilder.cs         | 17 +++++++++++
 .../ClientCreateAndInitializeTest.cs          | 29 +++++++++++++++++++
 .../Contracts/DotNetSDKAPI.json               |  5 ++++
 3 files changed, 51 insertions(+)

diff --git a/Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs b/Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs
index 59147c2d76..f3bdb65e5c 100644
--- a/Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs
+++ b/Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs
@@ -8,6 +8,8 @@ namespace Microsoft.Azure.Cosmos.Fluent
     using System.Collections.Generic;
     using System.Net;
     using System.Net.Http;
+    using System.Threading;
+    using System.Threading.Tasks;
     using global::Azure.Core;
     using Microsoft.Azure.Cosmos.Core.Trace;
     using Microsoft.Azure.Documents;
@@ -130,6 +132,21 @@ public CosmosClient Build()
                 new CosmosClient(this.accountEndpoint, this.tokenCredential, this.clientOptions);
         }
 
+        /// <summary>
+        /// A method to create the cosmos client and initialize the provided containers.
+        /// </summary>
+        /// <param name="containers">Containers to be initialized identified by it's database name and container name.</param>
+        /// <param name="cancellationToken">(Optional) Cancellation Token</param>
+        /// <returns>
+        /// A CosmosClient object.
+        /// </returns>
+        public Task<CosmosClient> BuildAndInitializeAsync(IReadOnlyList<(string databaseId, string containerId)> containers, CancellationToken cancellationToken = default)
+        {
+            return this.tokenCredential == null ?
+                CosmosClient.CreateAndInitializeAsync(this.accountEndpoint, this.accountKey, containers, this.clientOptions, cancellationToken) :
+                CosmosClient.CreateAndInitializeAsync(this.accountEndpoint, this.tokenCredential, containers, this.clientOptions, cancellationToken);
+        }
+
         /// <summary>
         /// A method to create the cosmos client
         /// </summary>
diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/ClientCreateAndInitializeTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/ClientCreateAndInitializeTest.cs
index d1d45b472f..03cad902ac 100644
--- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/ClientCreateAndInitializeTest.cs
+++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/ClientCreateAndInitializeTest.cs
@@ -6,6 +6,7 @@
     using System.Net.Http;
     using System.Threading;
     using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Fluent;
     using Microsoft.Azure.Cosmos.Routing;
     using Microsoft.Azure.Documents;
     using Microsoft.VisualStudio.TestTools.UnitTesting;
@@ -80,6 +81,34 @@ public async Task CreateAndInitializeTest()
             cosmosClient.Dispose();
         }
 
+        [TestMethod]
+        public async Task CreateAndInitializeWithCosmosClientBuilderTest()
+        {
+            int httpCallsMade = 0;
+            HttpClientHandlerHelper httpClientHandlerHelper = new HttpClientHandlerHelper
+            {
+                RequestCallBack = (request, cancellationToken) =>
+                {
+                    httpCallsMade++;
+                    return null;
+                }
+            };
+
+            (string endpoint, string authKey) = TestCommon.GetAccountInfo();
+            List<(string, string)> containers = new List<(string, string)>
+            { ("ClientCreateAndInitializeDatabase", "ClientCreateAndInitializeContainer")};
+
+            CosmosClientBuilder builder = new CosmosClientBuilder(endpoint, authKey).WithHttpClientFactory(() => new HttpClient(httpClientHandlerHelper));
+            CosmosClient cosmosClient = await builder.BuildAndInitializeAsync(containers);
+            Assert.IsNotNull(cosmosClient);
+            int httpCallsMadeAfterCreation = httpCallsMade;
+
+            ContainerInternal container = (ContainerInternal)cosmosClient.GetContainer("ClientCreateAndInitializeDatabase", "ClientCreateAndInitializeContainer");
+            ItemResponse<ToDoActivity> readResponse = await container.ReadItemAsync<ToDoActivity>("1", new Cosmos.PartitionKey("Status1"));
+            Assert.AreEqual(httpCallsMade, httpCallsMadeAfterCreation);
+            cosmosClient.Dispose();
+        }
+
         [TestMethod]
         [ExpectedException(typeof(HttpRequestException))]
         public async Task AuthIncorrectTest()
diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Contracts/DotNetSDKAPI.json b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Contracts/DotNetSDKAPI.json
index 56ea78e37e..aa8403d330 100644
--- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Contracts/DotNetSDKAPI.json
+++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Contracts/DotNetSDKAPI.json
@@ -4356,6 +4356,11 @@
           "Attributes": [],
           "MethodInfo": "Microsoft.Azure.Cosmos.Fluent.CosmosClientBuilder WithThrottlingRetryOptions(System.TimeSpan, Int32);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
         },
+        "System.Threading.Tasks.Task`1[Microsoft.Azure.Cosmos.CosmosClient] BuildAndInitializeAsync(System.Collections.Generic.IReadOnlyList`1[System.ValueTuple`2[System.String,System.String]], System.Threading.CancellationToken)": {
+          "Type": "Method",
+          "Attributes": [],
+          "MethodInfo": "System.Threading.Tasks.Task`1[Microsoft.Azure.Cosmos.CosmosClient] BuildAndInitializeAsync(System.Collections.Generic.IReadOnlyList`1[System.ValueTuple`2[System.String,System.String]], System.Threading.CancellationToken);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
         "Void .ctor(System.String, Azure.Core.TokenCredential)": {
           "Type": "Constructor",
           "Attributes": [],