[LLC] Implement RequestOptions.AddPolicy

sdk/core/Azure.Core.Experimental/tests/LowLevelClient/PetStoreClient.cs

@@ -146,7 +146,7 @@ public virtual Response GetPet(string id, RequestContext context = null)
         /// <param name="options"> The request options. </param>
         private HttpMessage CreateGetPetRequest(string id, RequestContext context = null)
         {
-            var message = Pipeline.CreateMessage();
+            var message = Pipeline.CreateMessage(context);
             var request = message.Request;
             request.Method = RequestMethod.Get;
             var uri = new RawRequestUriBuilder();

sdk/core/Azure.Core.Experimental/tests/LowLevelClientTests.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Linq;
 using System.Text.Json;
 using System.Threading.Tasks;
 using Azure.Core.Experimental.Tests;
@@ -211,6 +212,145 @@ public void ThrowOnErrorThrowsOnError()
             });
         }
 
+        [Test]
+        public async Task CanAddPerCallPolicy()
+        {
+            var petResponse = new MockResponse(200);
+
+            Pet pet = new("snoopy", "beagle");
+            petResponse.SetContent(SerializationHelpers.Serialize(pet, SerializePet));
+
+            var mockTransport = new MockTransport(petResponse);
+            PetStoreClient client = CreateClient(mockTransport);
+
+            var context = new RequestContext();
+            context.AddPolicy(new AddHeaderPolicy("PerCallHeader", "Value"), HttpPipelinePosition.PerCall);
+
+            Response response = await client.GetPetAsync("snoopy", context);
+
+            Request request = mockTransport.Requests[0];
+            Assert.IsTrue(request.Headers.TryGetValues("PerCallHeader", out var values));
+            Assert.AreEqual(1, values.Count());
+            Assert.AreEqual("Value", values.ElementAt(0));
+        }
+
+        [Test]
+        public async Task CanAddPerRetryPolicy()
+        {
+            var retryResponse = new MockResponse(408); // Request Timeout
+            var petResponse = new MockResponse(200);
+
+            Pet pet = new("snoopy", "beagle");
+            petResponse.SetContent(SerializationHelpers.Serialize(pet, SerializePet));
+
+            // retry twice -- this will add the header three times.
+            var mockTransport = new MockTransport(retryResponse, retryResponse, petResponse);
+            PetStoreClient client = CreateClient(mockTransport);
+
+            var context = new RequestContext();
+            context.AddPolicy(new AddHeaderPolicy("PerRetryHeader", "Value"), HttpPipelinePosition.PerRetry);
+
+            Response response = await client.GetPetAsync("snoopy", context);
+
+            Request request = mockTransport.Requests[0];
+            Assert.IsTrue(request.Headers.TryGetValues("PerRetryHeader", out var values));
+            Assert.AreEqual(3, values.Count());
+            Assert.AreEqual("Value", values.ElementAt(0));
+            Assert.AreEqual("Value", values.ElementAt(1));
+            Assert.AreEqual("Value", values.ElementAt(2));
+        }
+
+        [Test]
+        public async Task CanAddBeforeTransportPolicy()
+        {
+            var retryResponse = new MockResponse(408); // Request Timeout
+            var petResponse = new MockResponse(200);
+
+            Pet pet = new("snoopy", "beagle");
+            petResponse.SetContent(SerializationHelpers.Serialize(pet, SerializePet));
+
+            // retry twice
+            var mockTransport = new MockTransport(retryResponse, retryResponse, petResponse);
+            PetStoreClient client = CreateClient(mockTransport);
+
+            var context = new RequestContext();
+            context.AddPolicy(new AddHeaderPolicy("BeforeTransportHeader", "Value"), HttpPipelinePosition.BeforeTransport);
+
+            Response response = await client.GetPetAsync("snoopy", context);
+
+            Request request = mockTransport.Requests[0];
+            Assert.IsTrue(request.Headers.TryGetValues("BeforeTransportHeader", out var values));
+            Assert.AreEqual(3, values.Count());
+            Assert.AreEqual("Value", values.ElementAt(0));
+            Assert.AreEqual("Value", values.ElementAt(1));
+            Assert.AreEqual("Value", values.ElementAt(2));
+        }
+
+        [Test]
+        public async Task CanAddPoliciesAllPositions()
+        {
+            var retryResponse = new MockResponse(408); // Request Timeout
+            var petResponse = new MockResponse(200);
+
+            Pet pet = new("snoopy", "beagle");
+            petResponse.SetContent(SerializationHelpers.Serialize(pet, SerializePet));
+
+            // retry twice -- this will add the header three times.
+            var mockTransport = new MockTransport(retryResponse, retryResponse, petResponse);
+            PetStoreClient client = CreateClient(mockTransport);
+
+            var context = new RequestContext();
+            context.AddPolicy(new AddHeaderPolicy("PerCallHeader1", "PerCall1"), HttpPipelinePosition.PerCall);
+            context.AddPolicy(new AddHeaderPolicy("PerCallHeader2", "PerCall2"), HttpPipelinePosition.PerCall);
+            context.AddPolicy(new AddHeaderPolicy("PerRetryHeader", "PerRetry"), HttpPipelinePosition.PerRetry);
+            context.AddPolicy(new AddHeaderPolicy("BeforeTransportHeader", "BeforeTransport"), HttpPipelinePosition.BeforeTransport);
+
+            Response response = await client.GetPetAsync("snoopy", context);
+
+            Request request = mockTransport.Requests[0];
+
+            Assert.IsTrue(request.Headers.TryGetValues("PerCallHeader1", out var perCall1Values));
+            Assert.AreEqual(1, perCall1Values.Count());
+            Assert.AreEqual("PerCall1", perCall1Values.ElementAt(0));
+
+            Assert.IsTrue(request.Headers.TryGetValues("PerCallHeader2", out var perCall2Values));
+            Assert.AreEqual(1, perCall2Values.Count());
+            Assert.AreEqual("PerCall2", perCall2Values.ElementAt(0));
+
+            Assert.IsTrue(request.Headers.TryGetValues("PerRetryHeader", out var perRetryValues));
+            Assert.AreEqual("PerRetry", perRetryValues.ElementAt(0));
+            Assert.AreEqual("PerRetry", perRetryValues.ElementAt(1));
+            Assert.AreEqual("PerRetry", perRetryValues.ElementAt(2));
+
+            Assert.IsTrue(request.Headers.TryGetValues("BeforeTransportHeader", out var beforeTransportValues));
+            Assert.AreEqual("BeforeTransport", beforeTransportValues.ElementAt(0));
+            Assert.AreEqual("BeforeTransport", beforeTransportValues.ElementAt(1));
+            Assert.AreEqual("BeforeTransport", beforeTransportValues.ElementAt(2));
+        }
+
+        [Test]
+        public async Task ThrowsIfUsePipelineConstructor()
+        {
+            HttpPipeline pipeline = new HttpPipeline(new MockTransport());
+
+            var context = new RequestContext();
+            context.AddPolicy(new AddHeaderPolicy("PerCallHeader", "Value"), HttpPipelinePosition.PerCall);
+
+            var message = pipeline.CreateMessage(context);
+
+            bool throws = false;
+            try
+            {
+                await pipeline.SendAsync(message, context.CancellationToken);
+            }
+            catch (InvalidOperationException)
+            {
+                throws = true;
+            }
+
+            Assert.IsTrue(throws);
+        }
+
         #region Helpers
         private void SerializePet(ref Utf8JsonWriter writer, Pet pet)
         {
@@ -224,6 +364,30 @@ private void SerializePet(ref Utf8JsonWriter writer, Pet pet)
 
             writer.WriteEndObject();
         }
+
+        public class AddHeaderPolicy : HttpPipelinePolicy
+        {
+            private string _headerName;
+            private string _headerVaue;
+
+            public AddHeaderPolicy(string headerName, string headerValue) : base()
+            {
+                _headerName = headerName;
+                _headerVaue = headerValue;
+            }
+
+            public override void Process(HttpMessage message, ReadOnlyMemory<HttpPipelinePolicy> pipeline)
+            {
+                message.Request.Headers.Add(_headerName, _headerVaue);
+                ProcessNext(message, pipeline);
+            }
+
+            public override ValueTask ProcessAsync(HttpMessage message, ReadOnlyMemory<HttpPipelinePolicy> pipeline)
+            {
+                message.Request.Headers.Add(_headerName, _headerVaue);
+                return ProcessNextAsync(message, pipeline);
+            }
+        }
         #endregion
     }
 }

sdk/core/Azure.Core/api/Azure.Core.net461.cs

@@ -190,6 +190,7 @@ public partial class RequestContext
         public RequestContext() { }
         public System.Threading.CancellationToken CancellationToken { get { throw null; } set { } }
         public Azure.ErrorOptions ErrorOptions { get { throw null; } set { } }
+        public void AddPolicy(Azure.Core.Pipeline.HttpPipelinePolicy policy, Azure.Core.HttpPipelinePosition position) { }
         public static implicit operator Azure.RequestContext (Azure.ErrorOptions options) { throw null; }
     }
     public partial class RequestFailedException : System.Exception, System.Runtime.Serialization.ISerializable
@@ -755,6 +756,7 @@ public HttpPipeline(Azure.Core.Pipeline.HttpPipelineTransport transport, Azure.C
         public static System.IDisposable CreateClientRequestIdScope(string? clientRequestId) { throw null; }
         public static System.IDisposable CreateHttpMessagePropertiesScope(System.Collections.Generic.IDictionary<string, object?> messageProperties) { throw null; }
         public Azure.Core.HttpMessage CreateMessage() { throw null; }
+        public Azure.Core.HttpMessage CreateMessage(Azure.RequestContext context) { throw null; }
         public Azure.Core.Request CreateRequest() { throw null; }
         public void Send(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { }
         public System.Threading.Tasks.ValueTask SendAsync(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { throw null; }

sdk/core/Azure.Core/api/Azure.Core.net5.0.cs

@@ -190,6 +190,7 @@ public partial class RequestContext
         public RequestContext() { }
         public System.Threading.CancellationToken CancellationToken { get { throw null; } set { } }
         public Azure.ErrorOptions ErrorOptions { get { throw null; } set { } }
+        public void AddPolicy(Azure.Core.Pipeline.HttpPipelinePolicy policy, Azure.Core.HttpPipelinePosition position) { }
         public static implicit operator Azure.RequestContext (Azure.ErrorOptions options) { throw null; }
     }
     public partial class RequestFailedException : System.Exception, System.Runtime.Serialization.ISerializable
@@ -755,6 +756,7 @@ public HttpPipeline(Azure.Core.Pipeline.HttpPipelineTransport transport, Azure.C
         public static System.IDisposable CreateClientRequestIdScope(string? clientRequestId) { throw null; }
         public static System.IDisposable CreateHttpMessagePropertiesScope(System.Collections.Generic.IDictionary<string, object?> messageProperties) { throw null; }
         public Azure.Core.HttpMessage CreateMessage() { throw null; }
+        public Azure.Core.HttpMessage CreateMessage(Azure.RequestContext context) { throw null; }
         public Azure.Core.Request CreateRequest() { throw null; }
         public void Send(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { }
         public System.Threading.Tasks.ValueTask SendAsync(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { throw null; }

sdk/core/Azure.Core/api/Azure.Core.netcoreapp2.1.cs

@@ -190,6 +190,7 @@ public partial class RequestContext
         public RequestContext() { }
         public System.Threading.CancellationToken CancellationToken { get { throw null; } set { } }
         public Azure.ErrorOptions ErrorOptions { get { throw null; } set { } }
+        public void AddPolicy(Azure.Core.Pipeline.HttpPipelinePolicy policy, Azure.Core.HttpPipelinePosition position) { }
         public static implicit operator Azure.RequestContext (Azure.ErrorOptions options) { throw null; }
     }
     public partial class RequestFailedException : System.Exception, System.Runtime.Serialization.ISerializable
@@ -755,6 +756,7 @@ public HttpPipeline(Azure.Core.Pipeline.HttpPipelineTransport transport, Azure.C
         public static System.IDisposable CreateClientRequestIdScope(string? clientRequestId) { throw null; }
         public static System.IDisposable CreateHttpMessagePropertiesScope(System.Collections.Generic.IDictionary<string, object?> messageProperties) { throw null; }
         public Azure.Core.HttpMessage CreateMessage() { throw null; }
+        public Azure.Core.HttpMessage CreateMessage(Azure.RequestContext context) { throw null; }
         public Azure.Core.Request CreateRequest() { throw null; }
         public void Send(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { }
         public System.Threading.Tasks.ValueTask SendAsync(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { throw null; }

sdk/core/Azure.Core/api/Azure.Core.netstandard2.0.cs

@@ -190,6 +190,7 @@ public partial class RequestContext
         public RequestContext() { }
         public System.Threading.CancellationToken CancellationToken { get { throw null; } set { } }
         public Azure.ErrorOptions ErrorOptions { get { throw null; } set { } }
+        public void AddPolicy(Azure.Core.Pipeline.HttpPipelinePolicy policy, Azure.Core.HttpPipelinePosition position) { }
         public static implicit operator Azure.RequestContext (Azure.ErrorOptions options) { throw null; }
     }
     public partial class RequestFailedException : System.Exception, System.Runtime.Serialization.ISerializable
@@ -755,6 +756,7 @@ public HttpPipeline(Azure.Core.Pipeline.HttpPipelineTransport transport, Azure.C
         public static System.IDisposable CreateClientRequestIdScope(string? clientRequestId) { throw null; }
         public static System.IDisposable CreateHttpMessagePropertiesScope(System.Collections.Generic.IDictionary<string, object?> messageProperties) { throw null; }
         public Azure.Core.HttpMessage CreateMessage() { throw null; }
+        public Azure.Core.HttpMessage CreateMessage(Azure.RequestContext context) { throw null; }
         public Azure.Core.Request CreateRequest() { throw null; }
         public void Send(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { }
         public System.Threading.Tasks.ValueTask SendAsync(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { throw null; }

sdk/core/Azure.Core/src/HttpMessage.cs

@@ -81,6 +81,8 @@ public Response Response
         /// </summary>
         public TimeSpan? NetworkTimeout { get; set; }
 
+        internal List<(HttpPipelinePosition Position, HttpPipelinePolicy Policy)>? Policies { get; set; }
+
         /// <summary>
         /// Gets a property that modifies the pipeline behavior. Please refer to individual policies documentation on what properties it supports.
         /// </summary>