From 76844850801db52705ef1ad5e207f340bd0638d4 Mon Sep 17 00:00:00 2001
From: Caio Saldanha <camaiaor@microsoft.com>
Date: Mon, 26 Apr 2021 05:07:50 -0700
Subject: [PATCH] Added RetryOnTooManyRequestsPolicy to tests

---
 .../tests/MetricsAdvisorLiveTestBase.cs       | 14 +++-
 .../tests/RetryOnTooManyRequestsPolicy.cs     | 67 +++++++++++++++++++
 2 files changed, 79 insertions(+), 2 deletions(-)
 create mode 100644 sdk/metricsadvisor/Azure.AI.MetricsAdvisor/tests/RetryOnTooManyRequestsPolicy.cs

diff --git a/sdk/metricsadvisor/Azure.AI.MetricsAdvisor/tests/MetricsAdvisorLiveTestBase.cs b/sdk/metricsadvisor/Azure.AI.MetricsAdvisor/tests/MetricsAdvisorLiveTestBase.cs
index 5c791adcc421a..71c5b863af441 100644
--- a/sdk/metricsadvisor/Azure.AI.MetricsAdvisor/tests/MetricsAdvisorLiveTestBase.cs
+++ b/sdk/metricsadvisor/Azure.AI.MetricsAdvisor/tests/MetricsAdvisorLiveTestBase.cs
@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using Azure.AI.MetricsAdvisor.Administration;
 using Azure.AI.MetricsAdvisor.Models;
+using Azure.Core;
 using Azure.Core.TestFramework;
 using NUnit.Framework;
 
@@ -38,7 +39,7 @@ public MetricsAdvisorLiveTestBase(bool isAsync) : base(isAsync)
         public MetricsAdvisorAdministrationClient GetMetricsAdvisorAdministrationClient(bool useTokenCredential = false)
         {
             var endpoint = new Uri(TestEnvironment.MetricsAdvisorUri);
-            var instrumentedOptions = InstrumentClientOptions(new MetricsAdvisorClientsOptions());
+            var instrumentedOptions = GetInstrumentedOptions();
 
             MetricsAdvisorAdministrationClient client = useTokenCredential
                 ? new(endpoint, TestEnvironment.Credential, instrumentedOptions)
@@ -50,7 +51,7 @@ public MetricsAdvisorAdministrationClient GetMetricsAdvisorAdministrationClient(
         public MetricsAdvisorClient GetMetricsAdvisorClient(bool useTokenCredential = false)
         {
             var endpoint = new Uri(TestEnvironment.MetricsAdvisorUri);
-            var instrumentedOptions = InstrumentClientOptions(new MetricsAdvisorClientsOptions());
+            var instrumentedOptions = GetInstrumentedOptions();
 
             MetricsAdvisorClient client = useTokenCredential
                 ? new(endpoint, TestEnvironment.Credential, instrumentedOptions)
@@ -88,5 +89,14 @@ protected void ValidateGroupKey(DimensionKey groupKey)
                 Assert.That(column.Value, Is.Not.Null.And.Not.Empty);
             }
         }
+
+        private MetricsAdvisorClientsOptions GetInstrumentedOptions()
+        {
+            var options = new MetricsAdvisorClientsOptions();
+
+            options.AddPolicy(new RetryOnTooManyRequestsPolicy(), HttpPipelinePosition.PerRetry);
+
+            return InstrumentClientOptions(options);
+        }
     }
 }
diff --git a/sdk/metricsadvisor/Azure.AI.MetricsAdvisor/tests/RetryOnTooManyRequestsPolicy.cs b/sdk/metricsadvisor/Azure.AI.MetricsAdvisor/tests/RetryOnTooManyRequestsPolicy.cs
new file mode 100644
index 0000000000000..5612a0b3b9488
--- /dev/null
+++ b/sdk/metricsadvisor/Azure.AI.MetricsAdvisor/tests/RetryOnTooManyRequestsPolicy.cs
@@ -0,0 +1,67 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Azure.Core;
+using Azure.Core.Pipeline;
+
+namespace Azure.AI.MetricsAdvisor.Tests
+{
+    public class RetryOnTooManyRequestsPolicy : HttpPipelinePolicy
+    {
+        private const int TooManyRequestsStatusCode = 429;
+
+        private const string RetryAfterHeaderName = "Retry-After";
+
+        public override void Process(HttpMessage message, ReadOnlyMemory<HttpPipelinePolicy> pipeline) =>
+            ProcessAsync(message, pipeline, false).EnsureCompleted();
+
+        public override ValueTask ProcessAsync(HttpMessage message, ReadOnlyMemory<HttpPipelinePolicy> pipeline) =>
+            ProcessAsync(message, pipeline, true);
+
+        private async ValueTask ProcessAsync(HttpMessage message, ReadOnlyMemory<HttpPipelinePolicy> pipeline, bool async)
+        {
+            bool gotResponse = false;
+
+            while (!gotResponse)
+            {
+                if (async)
+                {
+                    await ProcessNextAsync(message, pipeline);
+                }
+                else
+                {
+                    ProcessNext(message, pipeline);
+                }
+
+                if (message.Response.Status == TooManyRequestsStatusCode)
+                {
+                    TimeSpan delay = message.Response.Headers.TryGetValue(RetryAfterHeaderName, out string retryAfterValue)
+                        && int.TryParse(retryAfterValue, out int delayInSeconds)
+                        ? TimeSpan.FromSeconds(delayInSeconds)
+                        : TimeSpan.FromSeconds(1);
+
+                    await WaitAsync(delay, async, message.CancellationToken);
+                }
+                else
+                {
+                    gotResponse = true;
+                }
+            }
+        }
+
+        private async ValueTask WaitAsync(TimeSpan delay, bool async, CancellationToken cancellationToken)
+        {
+            if (async)
+            {
+                await Task.Delay(delay, cancellationToken);
+            }
+            else
+            {
+                cancellationToken.WaitHandle.WaitOne(delay);
+            }
+        }
+    }
+}