[Storage][Core] Wait for environment to become eventually consistent. (…

…#20574)

* draft.

* tweaks + tests.

* pr feedback.

* PR feedback.

* rename.

* rename.

* valuetask.

* readme.

* pr feedback.

sdk/core/Azure.Core.TestFramework/README.md

@@ -124,6 +124,31 @@ public partial class ConfigurationSamples: SamplesBase<AppConfigurationTestEnvir
 }
 ```
 
+If resources require some time to become eventually consistent and there's a scenario that can be used to detect if asynchronous process completed
+then you can consider implementing `TestEnvironment.IsEnvironmentReadyAsync`. Test framework will probe the scenario couple of times before starting tests or
+fail test run if resources don't become available:
+
+``` C#
+public class AppConfigurationTestEnvironment : TestEnvironment
+{
+    // in addition to other members
+    protected override async ValueTask<bool> IsEnvironmentReadyAsync()
+    {
+        var connectionString = TestEnvironment.ConnectionString;
+        var client = new ConfigurationClient(connectionString);
+        try
+        {
+            await service.GetConfigurationSettingAsync("Setting");
+        } 
+        catch (RequestFailedException e) when (e.Status == 403)
+        {
+            return false;
+        }
+        return true;
+    }
+}
+```
+
 ## TokenCredential
 
 If a test or sample uses `TokenCredential` to construct the client use `TestEnvironment.Credential` to retrieve it.
@@ -431,4 +456,4 @@ For this to work with tests, your test class must have an `InternalsVisisbleTo`
 [assembly: InternalsVisibleTo("DynamicProxyGenAssembly2, PublicKey=0024000004800000940000000602000000240000525341310004000001000100c547cac37abd99c8db225ef2f6c8a3602f3b3606cc9891605d02baa56104f4cfc0734aa39b93bf7852f7d9266654753cc297e7d2edfe0bac1cdcf9f717241550e0a7b191195b7667bb4f64bcb8e2121380fd1d9d46ad2d92d2d15605093924cceaf74c4861eff62abf69b9291ed0a340e113be11e6a7d3113e92484cf7045cc7")]
 ```
 
-If this is neglected, _clientDiagnostics will be null at test runtime.
+If this is neglected, _clientDiagnostics will be null at test runtime.

sdk/core/Azure.Core.TestFramework/src/LiveTestBase.cs

@@ -1,6 +1,9 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+using System.Threading.Tasks;
+using NUnit.Framework;
+
 namespace Azure.Core.TestFramework
 {
     /// <summary>
@@ -16,5 +19,11 @@ protected LiveTestBase()
         }
 
         protected TEnvironment TestEnvironment { get; }
+
+        [OneTimeSetUp]
+        public async ValueTask WaitForEnvironment()
+        {
+            await TestEnvironment.WaitForEnvironmentAsync();
+        }
     }
-}
+}

sdk/core/Azure.Core.TestFramework/src/RecordedTestBaseOfTEnvironment.cs

@@ -1,6 +1,10 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+using System;
+using System.Threading.Tasks;
+using NUnit.Framework;
+
 namespace Azure.Core.TestFramework
 {
 #pragma warning disable SA1649 // File name should match first type name
@@ -23,5 +27,11 @@ public override void StartTestRecording()
         }
 
         public TEnvironment TestEnvironment { get; }
+
+        [OneTimeSetUp]
+        public async ValueTask WaitForEnvironment()
+        {
+            await TestEnvironment.WaitForEnvironmentAsync();
+        }
     }
 }

sdk/core/Azure.Core.TestFramework/src/TestEnvironment.cs

@@ -13,6 +13,7 @@
 using System.ComponentModel;
 using System.Linq;
 using NUnit.Framework;
+using System.Collections.Concurrent;
 
 namespace Azure.Core.TestFramework
 {
@@ -25,6 +26,8 @@ public abstract class TestEnvironment
         [EditorBrowsableAttribute(EditorBrowsableState.Never)]
         public static string RepositoryRoot { get; }
 
+        private static readonly ConcurrentDictionary<Type, Task> s_environmentStateCache = new ConcurrentDictionary<Type, Task>();
+
         private readonly string _prefix;
 
         private TokenCredential _credential;
@@ -175,6 +178,50 @@ public TokenCredential Credential
             }
         }
 
+        /// <summary>
+        /// Returns whether environment is ready to use. Should be overriden to provide service specific sampling scenario.
+        /// The test framework will wait until this returns true before starting tests.
+        /// Use this place to hook up logic that polls if eventual consistency has happened.
+        ///
+        /// Return true if environment is ready to use.
+        /// Return false if environment is not ready to use and framework should wait.
+        /// Throw if you want to fail the run fast.
+        /// </summary>
+        /// <returns>Whether environment is ready to use.</returns>
+        protected virtual ValueTask<bool> IsEnvironmentReadyAsync()
+        {
+            return new ValueTask<bool>(true);
+        }
+
+        /// <summary>
+        /// Waits until environment becomes ready to use. See <see cref="IsEnvironmentReadyAsync"/> to define sampling scenario.
+        /// </summary>
+        /// <returns>A task.</returns>
+        public async ValueTask WaitForEnvironmentAsync()
+        {
+            if (Mode == RecordedTestMode.Live)
+            {
+                await s_environmentStateCache.GetOrAdd(GetType(), t => WaitForEnvironmentInternalAsync());
+            }
+        }
+
+        private async Task WaitForEnvironmentInternalAsync()
+        {
+            int numberOfTries = 60;
+            TimeSpan delay = TimeSpan.FromSeconds(10);
+            for (int i = 0; i < numberOfTries; i++)
+            {
+                var isReady = await IsEnvironmentReadyAsync();
+                if (isReady)
+                {
+                    return;
+                }
+                await Task.Delay(delay);
+            }
+
+            throw new InvalidOperationException("The environment has not become ready, check your TestEnvironment.IsEnvironmentReady scenario.");
+        }
+
         /// <summary>
         /// Returns and records an environment variable value when running live or recorded value during playback.
         /// </summary>

sdk/core/Azure.Core/tests/TestEnvironmentTests.cs

@@ -8,8 +8,10 @@
 using System.Security.Cryptography;
 using System.Text.Json;
 using System.Threading;
+using System.Threading.Tasks;
 using Azure.Core.Serialization;
 using Azure.Core.TestFramework;
+using Moq;
 using NUnit.Framework;
 
 namespace Azure.Core.Tests
@@ -169,6 +171,34 @@ public void RecordedOptionalVariablePrefersPrefix()
             Assert.AreEqual("5", env.AzureEnvironment);
         }
 
+        [Test]
+        public async Task ShouldCacheExceptionIfWaitingForEnvironmentFailed()
+        {
+            var env = new WaitForEnvironmentTestEnvironmentFailureMode();
+
+            try
+            {
+                await env.WaitForEnvironmentAsync();
+                Assert.Fail();
+            }
+            catch (InvalidOperationException e)
+            {
+                StringAssert.Contains("kaboom", e.Message);
+            }
+
+            try
+            {
+                await env.WaitForEnvironmentAsync();
+                Assert.Fail();
+            }
+            catch (InvalidOperationException e)
+            {
+                StringAssert.Contains("kaboom", e.Message);
+            }
+
+            Assert.AreEqual(1, WaitForEnvironmentTestEnvironmentFailureMode.InvocationCount);
+        }
+
         private class RecordedVariableMisuse : RecordedTestBase<MockTestEnvironment>
         {
             // To make NUnit happy
@@ -223,5 +253,81 @@ private class MockTestEnvironment : TestEnvironment
             public string MissingOptionalSecret => GetRecordedOptionalVariable("MissingOptionalSecret", option => option.IsSecret("INVALID"));
             public string ConnectionStringWithSecret => GetRecordedVariable("ConnectionStringWithSecret", option => option.HasSecretConnectionStringParameter("key"));
         }
+
+        public class WaitForEnvironmentTestClassOne : RecordedTestBase<WaitForEnvironmentTestEnvironmentOne>
+        {
+            public WaitForEnvironmentTestClassOne(bool isAsync) : base(isAsync, RecordedTestMode.Live)
+            {
+            }
+
+            [Test]
+            public void ShouldCacheStateCorrectly()
+            {
+                Assert.AreEqual(2, WaitForEnvironmentTestEnvironmentOne.InvocationCount);
+            }
+        }
+
+        public class WaitForEnvironmentTestClassTwo : RecordedTestBase<WaitForEnvironmentTestEnvironmentTwo>
+        {
+            public WaitForEnvironmentTestClassTwo(bool isAsync) : base(isAsync, RecordedTestMode.Live)
+            {
+            }
+
+            [Test]
+            public void ShouldCacheStateCorrectly()
+            {
+                Assert.AreEqual(2, WaitForEnvironmentTestEnvironmentTwo.InvocationCount);
+            }
+        }
+
+        // This one uses same env as WaitForEnvironmentTestClassTwo to prove value is cached.
+        public class WaitForEnvironmentTestClassThree : RecordedTestBase<WaitForEnvironmentTestEnvironmentTwo>
+        {
+            public WaitForEnvironmentTestClassThree(bool isAsync) : base(isAsync, RecordedTestMode.Live)
+            {
+            }
+
+            [Test]
+            public void ShouldCacheStateCorrectly()
+            {
+                Assert.AreEqual(2, WaitForEnvironmentTestEnvironmentTwo.InvocationCount);
+            }
+        }
+
+        public class WaitForEnvironmentTestEnvironmentOne : TestEnvironment
+        {
+            public static int InvocationCount { get; private set; }
+
+            protected override ValueTask<bool> IsEnvironmentReadyAsync()
+            {
+                return new ValueTask<bool>(InvocationCount++ < 1 ? false : true);
+            }
+        }
+
+        public class WaitForEnvironmentTestEnvironmentTwo : TestEnvironment
+        {
+            public static int InvocationCount { get; private set; }
+
+            protected override ValueTask<bool> IsEnvironmentReadyAsync()
+            {
+                return new ValueTask<bool>(InvocationCount++ < 1 ? false : true);
+            }
+        }
+
+        public class WaitForEnvironmentTestEnvironmentFailureMode : TestEnvironment
+        {
+            public WaitForEnvironmentTestEnvironmentFailureMode()
+            {
+                Mode = RecordedTestMode.Live;
+            }
+
+            public static int InvocationCount { get; private set; }
+
+            protected override ValueTask<bool> IsEnvironmentReadyAsync()
+            {
+                InvocationCount++;
+                throw new InvalidOperationException("kaboom");
+            }
+        }
     }
 }

sdk/storage/Azure.Storage.Blobs.Batch/tests/Azure.Storage.Blobs.Batch.Tests.csproj

@@ -23,6 +23,7 @@
   <ItemGroup>
     <Compile Include="$(AzureCoreSharedSources)ArrayBufferWriter.cs" Link="Shared\Core\%(RecursiveDir)\%(Filename)%(Extension)" />
     <Compile Include="$(MSBuildThisFileDirectory)..\..\Azure.Storage.Blobs\tests\BlobTestBase.cs" Link="Shared\%(RecursiveDir)\%(Filename)%(Extension)" />
+    <Compile Include="$(MSBuildThisFileDirectory)..\..\Azure.Storage.Blobs\tests\BlobTestEnvironment.cs" Link="Shared\%(RecursiveDir)\%(Filename)%(Extension)" />
     <Compile Include="$(AzureStorageSharedSources)StorageConnectionString.cs" Link="Shared\%(RecursiveDir)\%(Filename)%(Extension)" />
     <Compile Include="$(AzureStorageSharedSources)SharedAccessSignatureCredentials.cs" Link="Shared\%(RecursiveDir)\%(Filename)%(Extension)" />
     <Compile Include="$(AzureStorageSharedSources)UriExtensions.cs" Link="Shared\%(RecursiveDir)\%(Filename)%(Extension)" />

sdk/storage/Azure.Storage.Blobs.ChangeFeed/tests/ChangeFeedTestBase.cs

@@ -19,7 +19,7 @@ namespace Azure.Storage.Blobs.ChangeFeed.Tests
         StorageVersionExtensions.MaxVersion,
     RecordingServiceVersion = StorageVersionExtensions.MaxVersion,
     LiveServiceVersions = new object[] { StorageVersionExtensions.LatestVersion })]
-    public class ChangeFeedTestBase : StorageTestBase
+    public class ChangeFeedTestBase : StorageTestBase<StorageTestEnvironment>
     {
         protected readonly BlobClientOptions.ServiceVersion _serviceVersion;