From 5c07a8ace76868a5294b1a74c3a6844836282fce Mon Sep 17 00:00:00 2001
From: Andriy Svyryd <Andriy.Svyryd@microsoft.com>
Date: Thu, 5 Sep 2024 10:55:26 -0700
Subject: [PATCH] Set shared database throughput when using RBAC

---
 .../TestUtilities/CosmosTestStore.cs          | 44 ++++++++++++++++---
 1 file changed, 37 insertions(+), 7 deletions(-)

diff --git a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs
index 4b2c755da80..805053b7f7f 100644
--- a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs
+++ b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs
@@ -172,7 +172,15 @@ private async Task CreateFromFile(DbContext context)
     {
         if (await EnsureCreatedAsync(context).ConfigureAwait(false))
         {
-            await context.Database.EnsureCreatedAsync().ConfigureAwait(false);
+            if (!TestEnvironment.UseTokenCredential)
+            {
+                await context.Database.EnsureCreatedAsync().ConfigureAwait(false);
+            }
+            else
+            {
+                await CreateContainersAsync(context).ConfigureAwait(false);
+            }
+
             var cosmosClient = context.GetService<ICosmosClientWrapper>();
             var serializer = CosmosClientWrapper.Serializer;
             using var fs = new FileStream(_dataFilePath!, FileMode.Open, FileAccess.Read);
@@ -253,7 +261,7 @@ public async Task<bool> EnsureCreatedAsync(DbContext context, CancellationToken
 
         var databaseAccount = await GetDBAccount(cancellationToken).ConfigureAwait(false);
         var collection = databaseAccount.Value.GetCosmosDBSqlDatabases();
-        var sqlDatabaseCreateUpdateOptions = new CosmosDBSqlDatabaseCreateOrUpdateContent(
+        var sqlDatabaseCreateUpdateContent = new CosmosDBSqlDatabaseCreateOrUpdateContent(
             TestEnvironment.AzureLocation,
             new CosmosDBSqlDatabaseResourceInfo(Name));
         if (await collection.ExistsAsync(Name, cancellationToken))
@@ -261,9 +269,28 @@ public async Task<bool> EnsureCreatedAsync(DbContext context, CancellationToken
             return false;
         }
 
-        var databaseResponse = (await collection.CreateOrUpdateAsync(
-            WaitUntil.Completed, Name, sqlDatabaseCreateUpdateOptions, cancellationToken).ConfigureAwait(false)).GetRawResponse();
-        return databaseResponse.Status == (int)HttpStatusCode.OK;
+        var model = context.GetService<IDesignTimeModel>().Model;
+
+        var modelThrouput = model.GetThroughput();
+        if (modelThrouput == null
+            && GetContainersToCreate(model).All(c => c.Throughput == null))
+        {
+            modelThrouput = ThroughputProperties.CreateManualThroughput(400);
+        }
+
+        if (modelThrouput != null)
+        {
+            sqlDatabaseCreateUpdateContent.Options = new CosmosDBCreateUpdateConfig
+            {
+                Throughput = modelThrouput.Throughput,
+                AutoscaleMaxThroughput = modelThrouput.AutoscaleMaxThroughput
+            };
+        }
+
+        var databaseResponse = await collection.CreateOrUpdateAsync(
+            WaitUntil.Completed, Name, sqlDatabaseCreateUpdateContent, cancellationToken).ConfigureAwait(false);
+
+        return databaseResponse.GetRawResponse().Status == (int)HttpStatusCode.OK;
     }
 
     private async Task<bool> EnsureDeletedAsync(DbContext context, CancellationToken cancellationToken = default)
@@ -362,8 +389,11 @@ private async Task CreateContainersAsync(DbContext context)
             var content = new CosmosDBSqlContainerCreateOrUpdateContent(TestEnvironment.AzureLocation, resource);
             if (container.Throughput != null)
             {
-                content.Options.AutoscaleMaxThroughput = container.Throughput.AutoscaleMaxThroughput;
-                content.Options.Throughput = container.Throughput.Throughput;
+                content.Options = new CosmosDBCreateUpdateConfig
+                {
+                    AutoscaleMaxThroughput = container.Throughput.AutoscaleMaxThroughput,
+                    Throughput = container.Throughput.Throughput
+                };
             }
 
             await database.Value.GetCosmosDBSqlContainers().CreateOrUpdateAsync(