[Storage] Live tests - address more flaky tests.

sdk/storage/Azure.Storage.Blobs/tests/BlobQuickQueryTests.cs

@@ -12,6 +12,7 @@
 using Azure.Storage.Blobs.Specialized;
 using Azure.Storage.Test;
 using Azure.Storage.Test.Shared;
+using Azure.Storage.Tests.Shared;
 using NUnit.Framework;
 
 namespace Azure.Storage.Blobs.Test
@@ -25,6 +26,7 @@ public BlobQuickQueryTests(bool async, BlobClientOptions.ServiceVersion serviceV
 
         [RecordedTest]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2019_12_12)]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_Min()
         {
             // Arrange
@@ -45,8 +47,8 @@ public async Task QueryAsync_Min()
         }
 
         [RecordedTest]
-
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2019_12_12)]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_Snapshot()
         {
             // Arrange
@@ -70,6 +72,7 @@ public async Task QueryAsync_Snapshot()
 
         [RecordedTest]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2019_12_12)]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_Error()
         {
             // Arrange
@@ -87,6 +90,7 @@ await TestHelper.AssertExpectedExceptionAsync<RequestFailedException>(
         [RecordedTest]
         [Ignore("Don't want to record 16 MB of data.")]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2019_12_12)]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_MultipleDataRecords()
         {
             // Arrange
@@ -129,6 +133,7 @@ public async Task QueryAsync_MultipleDataRecords()
         [RecordedTest]
         [Ignore("Don't want to record 250 MB of data.")]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2019_12_12)]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_Large()
         {
             // Arrange
@@ -162,6 +167,7 @@ public async Task QueryAsync_Large()
 
         [RecordedTest]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2019_12_12)]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_Progress()
         {
             // Arrange
@@ -193,6 +199,7 @@ public async Task QueryAsync_Progress()
         [RecordedTest]
         [Ignore("https://github.com/Azure/azure-sdk-for-net/issues/12063")]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2019_12_12)]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_QueryTextConfigurations()
         {
             await using DisposingContainer test = await GetTestContainerAsync();
@@ -237,6 +244,7 @@ public async Task QueryAsync_QueryTextConfigurations()
 
         [RecordedTest]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2019_12_12)]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_NonFatalError()
         {
             // Arrange
@@ -277,6 +285,7 @@ public async Task QueryAsync_NonFatalError()
         [RecordedTest]
         [Ignore("https://github.com/Azure/azure-sdk-for-net/issues/12063")]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2019_12_12)]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_FatalError()
         {
             // Arrange
@@ -327,6 +336,7 @@ public async Task QueryAsync_FatalError()
 
         [RecordedTest]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2019_12_12)]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_AccessConditions()
         {
             var garbageLeaseId = GetGarbageLeaseId();
@@ -362,6 +372,7 @@ public async Task QueryAsync_AccessConditions()
 
         [RecordedTest]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2019_12_12)]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_AccessConditionsFail()
         {
             var garbageLeaseId = GetGarbageLeaseId();
@@ -393,6 +404,7 @@ await TestHelper.AssertExpectedExceptionAsync<RequestFailedException>(
 
         [RecordedTest]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2019_12_12)]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_IfTags()
         {
             // Arrange
@@ -430,6 +442,7 @@ public async Task QueryAsync_IfTags()
 
         [RecordedTest]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2019_12_12)]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_IfTags_Failed()
         {
             // Arrange
@@ -458,6 +471,7 @@ await TestHelper.AssertExpectedExceptionAsync<RequestFailedException>(
 
         [RecordedTest]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2020_02_10)]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_ArrowConfiguration()
         {
             // Arrange
@@ -498,6 +512,7 @@ public async Task QueryAsync_ArrowConfiguration()
         [Test]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2020_08_04)]
         [PlaybackOnly("https://github.com/Azure/azure-sdk-for-net/issues/19575")]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_ParquetConfiguration()
         {
             // Arrange
@@ -528,6 +543,7 @@ public async Task QueryAsync_ParquetConfiguration()
         [Test]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2020_08_04)]
         [PlaybackOnly("https://github.com/Azure/azure-sdk-for-net/issues/19575")]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_ParquetOutputError()
         {
             // Arrange
@@ -549,6 +565,7 @@ await TestHelper.AssertExpectedExceptionAsync<ArgumentException>(
 
         [RecordedTest]
         [ServiceVersion(Min = BlobClientOptions.ServiceVersion.V2020_02_10)]
+        [RetryOnException(TestConstants.QuickQueryRetryCount, typeof(IOException))]
         public async Task QueryAsync_ArrowConfigurationInput()
         {
             // Arrange

sdk/storage/Azure.Storage.Blobs/tests/BlobTestEnvironment.cs

@@ -19,12 +19,13 @@ protected override async ValueTask<bool> IsEnvironmentReadyAsync()
 
         private async Task<bool> DoesOAuthWorkAsync()
         {
-            TestContext.Out.WriteLine("Blob Probing OAuth");
-            BlobServiceClient serviceClient = new BlobServiceClient(
-                new Uri(TestConfigurations.DefaultTargetOAuthTenant.BlobServiceEndpoint),
-                GetOAuthCredential(TestConfigurations.DefaultTargetOAuthTenant));
+            TestContext.Error.WriteLine("Blob Probing OAuth");
+
             try
             {
+                BlobServiceClient serviceClient = new BlobServiceClient(
+                    new Uri(TestConfigurations.DefaultTargetOAuthTenant.BlobServiceEndpoint),
+                    GetOAuthCredential(TestConfigurations.DefaultTargetOAuthTenant));
                 await serviceClient.GetPropertiesAsync();
                 var containerName = Guid.NewGuid().ToString();
                 var containerClient = serviceClient.GetBlobContainerClient(containerName);
@@ -43,10 +44,10 @@ private async Task<bool> DoesOAuthWorkAsync()
                 }
             } catch (RequestFailedException e) when (e.Status == 403 && e.ErrorCode == "AuthorizationPermissionMismatch")
             {
-                TestContext.Out.WriteLine("Blob Probing OAuth - not ready");
+                TestContext.Error.WriteLine("Blob Probing OAuth - not ready");
                 return false;
             }
-            TestContext.Out.WriteLine("Blob Probing OAuth - ready");
+            TestContext.Error.WriteLine("Blob Probing OAuth - ready");
             return true;
         }
     }

sdk/storage/Azure.Storage.Common/tests/RetryOnExceptionAttributeTests.cs

@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using Azure.Storage.Tests.Shared;
+using NUnit.Framework;
+
+namespace Azure.Storage.Tests
+{
+    public class RetryOnExceptionAttributeTests
+    {
+        private int counter = 0;
+
+        [Test]
+        [RetryOnException(3, typeof(InvalidOperationException))]
+        public void ShouldRetryOnException()
+        {
+            if (counter++ == 0)
+            {
+                throw new InvalidOperationException();
+            }
+
+            Assert.AreEqual(2, counter);
+        }
+    }
+}

sdk/storage/Azure.Storage.Common/tests/Shared/RetryOnExceptionAttribute.cs

@@ -0,0 +1,99 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using NUnit.Framework;
+using NUnit.Framework.Interfaces;
+using NUnit.Framework.Internal;
+using NUnit.Framework.Internal.Commands;
+
+namespace Azure.Storage.Tests.Shared
+{
+    /// <summary>
+    /// This is clone of Retry from NUnit. Except it retries on exception. NUnit version retries on assertion error only.
+    /// Use only to deal with flaky live tests.
+    /// </summary>
+    [AttributeUsage(AttributeTargets.Method, AllowMultiple = false, Inherited = false)]
+    public class RetryOnExceptionAttribute : NUnitAttribute, IRepeatTest
+    {
+        private readonly int _tryCount;
+        private readonly Type _exceptionType;
+
+        public RetryOnExceptionAttribute(int tryCount, Type exceptionType)
+        {
+            _tryCount = tryCount;
+            _exceptionType = exceptionType;
+        }
+
+        #region IRepeatTest Members
+        public TestCommand Wrap(TestCommand command)
+        {
+            return new RetryCommand(command, _tryCount, _exceptionType);
+        }
+
+        #endregion
+
+        #region Nested RetryCommand Class
+
+        public class RetryCommand : DelegatingTestCommand
+        {
+            private readonly int _tryCount;
+            private readonly Type _exceptionType;
+
+            public RetryCommand(TestCommand innerCommand, int tryCount, Type exceptionType)
+                : base(innerCommand)
+            {
+                _tryCount = tryCount;
+                _exceptionType = exceptionType;
+            }
+
+            public override TestResult Execute(TestExecutionContext context)
+            {
+                int count = _tryCount;
+
+                while (count-- > 0)
+                {
+                    try
+                    {
+                        context.CurrentResult = innerCommand.Execute(context);
+                    }
+                    // Commands are supposed to catch exceptions, but some don't
+                    // and we want to look at restructuring the API in the future.
+                    catch (Exception ex)
+                    {
+                        if (context.CurrentResult == null)
+                            context.CurrentResult = context.CurrentTest.MakeTestResult();
+                        context.CurrentResult.RecordException(ex);
+                    }
+
+                    // Clear result for retry
+                    if (count > 0 && IsTestFailedWithExpectedException(context))
+                    {
+                        context.CurrentResult = context.CurrentTest.MakeTestResult();
+                        context.CurrentRepeatCount++; // increment Retry count for next iteration. will only happen if we are guaranteed another iteration
+                    }
+                    else
+                    {
+                        break;
+                    }
+                }
+
+                return context.CurrentResult;
+            }
+
+            private bool IsTestFailedWithExpectedException(TestExecutionContext context)
+            {
+                var failed = context.CurrentResult.ResultState.Status switch
+                {
+                    TestStatus.Passed => false,
+                    TestStatus.Skipped => false,
+                    _ => true
+                };
+
+                return failed && context.CurrentResult.Message.Contains(_exceptionType.FullName);
+            }
+        }
+
+        #endregion
+    }
+}

sdk/storage/Azure.Storage.Common/tests/Shared/TestConstants.cs

@@ -23,6 +23,9 @@ public class TestConstants
         public const int DataLakeRetryDelay = 70000;
         public const int RetryDelay = 10000;
 
+        // quick query fails with with connection reset, retry until it's solved properly https://github.com/Azure/azure-sdk-for-net/issues/17403
+        public const int QuickQueryRetryCount = 5;
+
         public string CacheControl { get; private set; }
         public string ContentDisposition { get; private set; }
         public string ContentEncoding { get; private set; }