[release/8.0] Use TokenCredential for Cosmos tests (#33966)

dotnet · Jun 12, 2024 · 15d21bf · 15d21bf
1 parent 67fd137
commit 15d21bf
Show file tree

Hide file tree

Showing 15 changed files with 178 additions and 31 deletions.
diff --git a/All.sln b/All.sln
@@ -1,5 +1,4 @@
-
-Microsoft Visual Studio Solution File, Format Version 12.00
+Microsoft Visual Studio Solution File, Format Version 12.00
 # Visual Studio Version 17
 VisualStudioVersion = 17.0.31521.260
 MinimumVisualStudioVersion = 17.0.31521.260

diff --git a/azure-pipelines.yml b/azure-pipelines.yml
@@ -120,7 +120,7 @@ extends:
               - script: eng\common\cibuild.cmd -configuration $(_BuildConfig) -prepareMachine $(_InternalBuildArgs) $(_InternalRuntimeDownloadArgs) $(_AdditionalBuildArgs)
                 env:
                   Test__Cosmos__DefaultConnection: $(_CosmosConnectionUrl)
-                name: Build
+                displayName: Build
             templateContext:
               outputs:
               - output: pipelineArtifact
@@ -131,7 +131,7 @@ extends:
           - job: macOS
             pool:
               name: Azure Pipelines
-              image: macOS-11
+              image: macOS-12
               os: macOS
             variables:
               # Rely on task Arcade injects, not auto-injected build step.
@@ -149,7 +149,7 @@ extends:
                   Test__Cosmos__DefaultConnection: $(_CosmosConnectionUrl)
                   # Work-around for https://github.com/dotnet/runtime/issues/70758
                   COMPlus_EnableWriteXorExecute: 0
-                name: Build
+                displayName: Build
             templateContext:
               outputs:
               - output: pipelineArtifact
@@ -172,12 +172,10 @@ extends:
             steps:
               - bash: |
                   echo "##vso[task.setvariable variable=_CosmosConnectionUrl]https://ef-nightly-test.documents.azure.com:443/"
-                  echo "##vso[task.setvariable variable=_CosmosToken]$(ef-nightly-cosmos-key)"
                 displayName: Prepare to run Cosmos tests on ef-nightly-test
                 condition: and(eq(variables['_CosmosConnectionUrl'], 'true'), or(endsWith(variables['_runCounter'], '0'), endsWith(variables['_runCounter'], '2'), endsWith(variables['_runCounter'], '4'), endsWith(variables['_runCounter'], '6'), endsWith(variables['_runCounter'], '8')))
               - bash: |
                   echo "##vso[task.setvariable variable=_CosmosConnectionUrl]https://ef-pr-test.documents.azure.com:443/"
-                  echo "##vso[task.setvariable variable=_CosmosToken]$(ef-pr-cosmos-test)"
                 displayName: Prepare to run Cosmos tests on ef-pr-test
                 condition: and(eq(variables['_CosmosConnectionUrl'], 'true'), or(endsWith(variables['_runCounter'], '1'), endsWith(variables['_runCounter'], '3'), endsWith(variables['_runCounter'], '5'), endsWith(variables['_runCounter'], '7'), endsWith(variables['_runCounter'], '9')))
               - task: Bash@3
@@ -188,10 +186,20 @@ extends:
                 env:
                   Token: $(dn-bot-dnceng-artifact-feeds-rw)
               - script: eng/common/cibuild.sh --configuration $(_BuildConfig) --prepareMachine $(_InternalRuntimeDownloadArgs)
+                displayName: Build
+              - task: AzureCLI@2
+                displayName: Run Cosmos tests
+                inputs:
+                  azureSubscription: EFCosmosTesting
+                  scriptType: bash
+                  scriptLocation: 'inlineScript'
+                  inlineScript: |
+                    ./test.sh --ci --configuration $(_BuildConfig) --projects $(Build.SourcesDirectory)/test/EFCore.Cosmos.FunctionalTests/EFCore.Cosmos.FunctionalTests.csproj
                 env:
                   Test__Cosmos__DefaultConnection: $(_CosmosConnectionUrl)
-                  Test__Cosmos__AuthToken: $(_CosmosToken)
-                name: Build
+                  Test__Cosmos__UseTokenCredential: true
+                  Test__Cosmos__SubscriptionId: d709b837-4a74-4aec-addc-b6e4b9b23e7e
+                  Test__Cosmos__ResourceGroup: efcosmosci
             templateContext:
               outputs:
               - output: pipelineArtifact

diff --git a/eng/Versions.props b/eng/Versions.props
@@ -38,5 +38,7 @@
     <!-- NB: This version affects Visual Studio compatibility. See https://learn.microsoft.com/visualstudio/extensibility/roslyn-version-support -->
     <MicrosoftCodeAnalysisVersion>4.5.0</MicrosoftCodeAnalysisVersion>
     <MicrosoftCodeAnalysisTestingVersion>1.1.2-beta1.24121.1</MicrosoftCodeAnalysisTestingVersion>
+    <AzureIdentityVersion>1.11.3</AzureIdentityVersion>
+    <AzureResourceManagerCosmosDBVersion>1.3.2</AzureResourceManagerCosmosDBVersion>
   </PropertyGroup>
 </Project>
diff --git a/test/EFCore.Cosmos.FunctionalTests/ConfigPatternsCosmosTest.cs b/test/EFCore.Cosmos.FunctionalTests/ConfigPatternsCosmosTest.cs
@@ -6,6 +6,7 @@
 // ReSharper disable UnusedAutoPropertyAccessor.Local
 namespace Microsoft.EntityFrameworkCore.Cosmos;
 
+[CosmosCondition(CosmosCondition.DoesNotUseTokenCredential)]
 public class ConfigPatternsCosmosTest : IClassFixture<ConfigPatternsCosmosTest.CosmosFixture>
 {
     private const string DatabaseName = "ConfigPatternsCosmos";

diff --git a/test/EFCore.Cosmos.FunctionalTests/ConnectionSpecificationTest.cs b/test/EFCore.Cosmos.FunctionalTests/ConnectionSpecificationTest.cs
@@ -6,6 +6,7 @@
 
 namespace Microsoft.EntityFrameworkCore.Cosmos;
 
+[CosmosCondition(CosmosCondition.DoesNotUseTokenCredential)]
 public class ConnectionSpecificationTest
 {
     [ConditionalFact]

diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosConcurrencyTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosConcurrencyTest.cs
@@ -3,6 +3,7 @@
 
 namespace Microsoft.EntityFrameworkCore.Cosmos;
 
+[CosmosCondition(CosmosCondition.DoesNotUseTokenCredential)]
 public class CosmosConcurrencyTest : IClassFixture<CosmosConcurrencyTest.CosmosFixture>
 {
     private const string DatabaseName = "CosmosConcurrencyTest";

diff --git a/test/EFCore.Cosmos.FunctionalTests/EFCore.Cosmos.FunctionalTests.csproj b/test/EFCore.Cosmos.FunctionalTests/EFCore.Cosmos.FunctionalTests.csproj
@@ -56,7 +56,8 @@
   <ItemGroup>
     <PackageReference Include="Microsoft.Extensions.Configuration.EnvironmentVariables" Version="$(MicrosoftExtensionsConfigurationEnvironmentVariablesVersion)" />
     <PackageReference Include="Microsoft.Extensions.Configuration.Json" Version="$(MicrosoftExtensionsConfigurationJsonVersion)" />
-    <PackageReference Include="Azure.Identity" Version="1.10.2" />
+    <PackageReference Include="Azure.Identity" Version="$(AzureIdentityVersion)" />
+    <PackageReference Include="Azure.ResourceManager.CosmosDB" Version="$(AzureResourceManagerCosmosDBVersion)" />
   </ItemGroup>
 
 </Project>
diff --git a/test/EFCore.Cosmos.FunctionalTests/EmbeddedDocumentsTest.cs b/test/EFCore.Cosmos.FunctionalTests/EmbeddedDocumentsTest.cs
@@ -498,7 +498,6 @@ public virtual async Task Can_use_non_int_keys_for_embedded_entities()
         Guid addressGuid;
         await using (var context = new EmbeddedTransportationContext(options))
         {
-            await context.Database.EnsureCreatedAsync();
             var person = new Person { Id = 1 };
             address = new Address { Street = "Second", City = "Village" };
             person.Addresses.Add(address);

diff --git a/test/EFCore.Cosmos.FunctionalTests/Internal/CosmosDatabaseCreatorTest.cs b/test/EFCore.Cosmos.FunctionalTests/Internal/CosmosDatabaseCreatorTest.cs
@@ -3,6 +3,7 @@
 
 namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal;
 
+[CosmosCondition(CosmosCondition.DoesNotUseTokenCredential)]
 public class CosmosDatabaseCreatorTest
 {
     public static IEnumerable<object[]> IsAsyncData = new[]

diff --git a/test/EFCore.Cosmos.FunctionalTests/ReloadTest.cs b/test/EFCore.Cosmos.FunctionalTests/ReloadTest.cs
@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using Azure.Core;
 using Newtonsoft.Json.Linq;
 
 namespace Microsoft.EntityFrameworkCore.Cosmos;
@@ -43,21 +44,35 @@ public class ReloadTestContext : DbContext
         private readonly string _connectionUri;
         private readonly string _authToken;
         private readonly string _name;
+        private readonly TokenCredential _tokenCredential;
 
         public ReloadTestContext(CosmosTestStore testStore)
         {
             _connectionUri = testStore.ConnectionUri;
             _authToken = testStore.AuthToken;
             _name = testStore.Name;
+            _tokenCredential = testStore.TokenCredential;
         }
 
         protected override void OnConfiguring(DbContextOptionsBuilder optionsBuilder)
-            => optionsBuilder
-                .UseCosmos(
-                    _connectionUri,
-                    _authToken,
-                    _name,
-                    b => b.ApplyConfiguration());
+        {
+            if (TestEnvironment.UseTokenCredential)
+            {
+                optionsBuilder.UseCosmos(
+                            _connectionUri,
+                            _tokenCredential,
+                            _name,
+                            b => b.ApplyConfiguration());
+            }
+            else
+            {
+                optionsBuilder.UseCosmos(
+                            _connectionUri,
+                            _authToken,
+                            _name,
+                            b => b.ApplyConfiguration());
+            }
+        }
 
         protected override void OnModelCreating(ModelBuilder modelBuilder)
         {

diff --git a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosCondition.cs b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosCondition.cs
@@ -0,0 +1,11 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.EntityFrameworkCore.TestUtilities;
+
+[Flags]
+public enum CosmosCondition
+{
+    UsesTokenCredential = 1 << 0,
+    DoesNotUseTokenCredential = 1 << 1,
+}
diff --git a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosConditionAttribute.cs b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosConditionAttribute.cs
@@ -0,0 +1,38 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Microsoft.EntityFrameworkCore.TestUtilities.Xunit;
+
+namespace Microsoft.EntityFrameworkCore.TestUtilities;
+
+[AttributeUsage(AttributeTargets.Method | AttributeTargets.Class)]
+public sealed class CosmosConditionAttribute(CosmosCondition conditions) : Attribute, ITestCondition
+{
+    public CosmosCondition Conditions { get; set; } = conditions;
+
+    public ValueTask<bool> IsMetAsync()
+    {
+        var isMet = true;
+
+        if (Conditions.HasFlag(CosmosCondition.UsesTokenCredential))
+        {
+            isMet &= TestEnvironment.UseTokenCredential;
+        }
+
+        if (Conditions.HasFlag(CosmosCondition.DoesNotUseTokenCredential))
+        {
+            isMet &= !TestEnvironment.UseTokenCredential;
+        }
+
+        return ValueTask.FromResult(isMet);
+    }
+
+    public string SkipReason
+        => string.Format(
+                "The test Cosmos account does not meet these conditions: '{0}'",
+                string.Join(
+                    ", ", Enum.GetValues(typeof(CosmosCondition))
+                        .Cast<Enum>()
+                        .Where(Conditions.HasFlag)
+                        .Select(f => Enum.GetName(typeof(CosmosCondition), f))));
+}
diff --git a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestHelpers.cs b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestHelpers.cs
@@ -17,7 +17,12 @@ public override IServiceCollection AddProviderServices(IServiceCollection servic
         => services.AddEntityFrameworkCosmos();
 
     public override DbContextOptionsBuilder UseProviderOptions(DbContextOptionsBuilder optionsBuilder)
-        => optionsBuilder.UseCosmos(
+        => TestEnvironment.UseTokenCredential
+        ? optionsBuilder.UseCosmos(
+            TestEnvironment.DefaultConnection,
+            TestEnvironment.TokenCredential,
+            "UnitTests")
+        : optionsBuilder.UseCosmos(
             TestEnvironment.DefaultConnection,
             TestEnvironment.AuthToken,
             "UnitTests");

diff --git a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs
@@ -1,8 +1,13 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Net;
 using System.Net.Sockets;
+using Azure;
 using Azure.Core;
+using Azure.ResourceManager;
+using Azure.ResourceManager.CosmosDB;
+using Azure.ResourceManager.CosmosDB.Models;
 using Microsoft.Azure.Cosmos;
 using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal;
 using Newtonsoft.Json;
@@ -76,11 +81,9 @@ protected override DbContext CreateDefaultContext()
         => new TestStoreContext(this);
 
     public override DbContextOptionsBuilder AddProviderOptions(DbContextOptionsBuilder builder)
-        => builder.UseCosmos(
-            ConnectionUri,
-            AuthToken,
-            Name,
-            _configureCosmos);
+        => TestEnvironment.UseTokenCredential
+        ? builder.UseCosmos(ConnectionUri, TokenCredential, Name, _configureCosmos)
+        : builder.UseCosmos(ConnectionUri, AuthToken, Name, _configureCosmos);
 
     public static async ValueTask<bool> IsConnectionAvailableAsync()
     {
@@ -160,8 +163,9 @@ protected override void Initialize(Func<DbContext> createContext, Action<DbConte
 
     private async Task CreateFromFile(DbContext context)
     {
-        if (await context.Database.EnsureCreatedAsync())
+        if (await EnsureCreatedAsync(context))
         {
+            await context.Database.EnsureCreatedAsync();
             var cosmosClient = context.GetService<ICosmosClientWrapper>();
             var serializer = CosmosClientWrapper.Serializer;
             using var fs = new FileStream(_dataFilePath, FileMode.Open, FileAccess.Read);
@@ -217,13 +221,52 @@ await cosmosClient.CreateItemAsync(
         }
     }
 
+    private static readonly ArmClient _armClient = new(TestEnvironment.TokenCredential);
+
     public override void Clean(DbContext context)
         => CleanAsync(context).GetAwaiter().GetResult();
 
+    public async Task<bool> EnsureCreatedAsync(DbContext context, CancellationToken cancellationToken = default)
+    {
+        if (!TestEnvironment.UseTokenCredential)
+        {
+            var cosmosClientWrapper = context.GetService<ICosmosClientWrapper>();
+            return await cosmosClientWrapper.CreateDatabaseIfNotExistsAsync(null, cancellationToken);
+        }
+
+        var databaseAccount = await GetDBAccount(cancellationToken);
+        var collection = databaseAccount.Value.GetCosmosDBSqlDatabases();
+        var sqlDatabaseCreateUpdateOptions = new CosmosDBSqlDatabaseCreateOrUpdateContent(TestEnvironment.AzureLocation,
+            new CosmosDBSqlDatabaseResourceInfo(Name));
+        var databaseResponse = (await collection.CreateOrUpdateAsync(
+            WaitUntil.Completed, Name, sqlDatabaseCreateUpdateOptions, cancellationToken)).GetRawResponse();
+        return databaseResponse.Status == (int)HttpStatusCode.Created;
+    }
+
+    private async Task<bool> EnsureDeletedAsync(DbContext context, CancellationToken cancellationToken = default)
+    {
+        if (!TestEnvironment.UseTokenCredential)
+        {
+            return await context.Database.EnsureDeletedAsync(cancellationToken);
+        }
+
+        var databaseAccount = await GetDBAccount(cancellationToken);
+        var databaseResponse = (await databaseAccount.Value.GetCosmosDBSqlDatabase(Name, cancellationToken).Value.DeleteAsync(
+            WaitUntil.Completed, cancellationToken)).GetRawResponse();
+        return databaseResponse.Status == (int)HttpStatusCode.Created;
+    }
+
+    private async Task<global::Azure.Response<CosmosDBAccountResource>> GetDBAccount(CancellationToken cancellationToken)
+    {
+        var accountName = new Uri(ConnectionUri).Host.Split('.').First();
+        var databaseAccountIdentifier = CosmosDBAccountResource.CreateResourceIdentifier(
+            TestEnvironment.SubscriptionId, TestEnvironment.ResourceGroup, accountName);
+        return await _armClient.GetCosmosDBAccountResource(databaseAccountIdentifier).GetAsync(cancellationToken);
+    }
+
     public override async Task CleanAsync(DbContext context)
     {
-        var cosmosClientWrapper = context.GetService<ICosmosClientWrapper>();
-        var created = await cosmosClientWrapper.CreateDatabaseIfNotExistsAsync(null);
+        var created = await EnsureCreatedAsync(context);
         try
         {
             if (!created)
@@ -274,7 +317,7 @@ await container.DeleteItemAsync<object>(
         {
             try
             {
-                await context.Database.EnsureDeletedAsync();
+                await EnsureDeletedAsync(context);
             }
             catch (Exception)
             {
@@ -302,7 +345,7 @@ public override async Task DisposeAsync()
                 GetTestStoreIndex(ServiceProvider).RemoveShared(GetType().Name + Name);
             }
 
-            await _storeContext.Database.EnsureDeletedAsync();
+            await EnsureDeletedAsync(_storeContext);
         }
 
         _storeContext.Dispose();
@@ -318,7 +361,16 @@ public TestStoreContext(CosmosTestStore testStore)
         }
 
         protected override void OnConfiguring(DbContextOptionsBuilder optionsBuilder)
-            => optionsBuilder.UseCosmos(_testStore.ConnectionUri, _testStore.AuthToken, _testStore.Name, _testStore._configureCosmos);
+        {
+            if (TestEnvironment.UseTokenCredential)
+            {
+                optionsBuilder.UseCosmos(_testStore.ConnectionUri, _testStore.TokenCredential, _testStore.Name, _testStore._configureCosmos);
+            }
+            else
+            {
+                optionsBuilder.UseCosmos(_testStore.ConnectionUri, _testStore.AuthToken, _testStore.Name, _testStore._configureCosmos);
+            }
+        }
     }
 
     private class FakeUpdateEntry : IUpdateEntry