Fix adding multiple headers and query string parameters (Azure#8986)

* Fix adding multiple headers and query string parameters

* Add mock tests

sdk/appconfiguration/Azure.Data.AppConfiguration/src/ConfigurationClient.cs

@@ -84,8 +84,8 @@ public ConfigurationClient(Uri endpoint, TokenCredential credential, Configurati
 
         private static HttpPipeline CreatePipeline(ConfigurationClientOptions options, HttpPipelinePolicy authenticationPolicy)
             => HttpPipelineBuilder.Build(options,
-                new HttpPipelinePolicy[] { new CustomHeadersPolicy() },
-                new HttpPipelinePolicy[] { new ApiVersionPolicy(options.GetVersionString()), authenticationPolicy, new SyncTokenPolicy() },
+                new HttpPipelinePolicy[] { new CustomHeadersPolicy(), new ApiVersionPolicy(options.GetVersionString()) },
+                new HttpPipelinePolicy[] { authenticationPolicy, new SyncTokenPolicy() },
                 new ResponseClassifier());
 
         private static string GetDefaultScope(Uri uri)
@@ -686,7 +686,7 @@ private Request CreateGetRequest(string key, string label, DateTimeOffset accept
             if (acceptDateTime != default)
             {
                 var dateTime = acceptDateTime.UtcDateTime.ToString(AcceptDateTimeFormat, CultureInfo.InvariantCulture);
-                request.Headers.Add(AcceptDatetimeHeader, dateTime);
+                request.Headers.SetValue(AcceptDatetimeHeader, dateTime);
             }
 
             if (requestOptions != null)
@@ -773,7 +773,7 @@ private Request CreateBatchRequest(SettingSelector selector, string pageLink)
             if (selector.AcceptDateTime.HasValue)
             {
                 var dateTime = selector.AcceptDateTime.Value.UtcDateTime.ToString(AcceptDateTimeFormat, CultureInfo.InvariantCulture);
-                request.Headers.Add(AcceptDatetimeHeader, dateTime);
+                request.Headers.SetValue(AcceptDatetimeHeader, dateTime);
             }
 
             return request;
@@ -854,7 +854,7 @@ private Request CreateGetRevisionsRequest(SettingSelector selector, string pageL
             if (selector.AcceptDateTime.HasValue)
             {
                 var dateTime = selector.AcceptDateTime.Value.UtcDateTime.ToString(AcceptDateTimeFormat, CultureInfo.InvariantCulture);
-                request.Headers.Add(AcceptDatetimeHeader, dateTime);
+                request.Headers.SetValue(AcceptDatetimeHeader, dateTime);
             }
 
             return request;

sdk/appconfiguration/Azure.Data.AppConfiguration/src/SyncTokenPolicy.cs

@@ -21,6 +21,7 @@ public SyncTokenPolicy()
 
         public override void OnSendingRequest(HttpMessage message)
         {
+            message.Request.Headers.Remove(SyncTokenHeader);
             foreach (SyncToken token in _syncTokens.Values)
             {
                 message.Request.Headers.Add(SyncTokenHeader, token.ToString());

sdk/appconfiguration/Azure.Data.AppConfiguration/tests/ConfigurationMockTests.cs

@@ -534,6 +534,32 @@ public async Task RequestHasSpecificApiVersion()
             StringAssert.Contains("api-version=1.0", request.Uri.ToUri().ToString());
         }
 
+        [Test]
+        public async Task RequestHasSpecificApiVersionOnlyOnceOnRetry()
+        {
+            var response = new MockResponse(200);
+            response.SetContent(SerializationHelpers.Serialize(s_testSetting, SerializeSetting));
+
+            var mockTransport = new MockTransport(new MockResponse(503), response);
+            var options = new ConfigurationClientOptions(ConfigurationClientOptions.ServiceVersion.V1_0);
+            options.Diagnostics.ApplicationId = "test_application";
+            options.Transport = mockTransport;
+
+            ConfigurationClient client = CreateClient<ConfigurationClient>(s_connectionString, options);
+
+            await client.AddConfigurationSettingAsync(s_testSetting);
+            MockRequest request = mockTransport.Requests[0];
+            MockRequest retriedRequest = mockTransport.Requests[1];
+
+            const string expectedApiString = "api-version=1.0";
+            StringAssert.Contains(expectedApiString, request.Uri.Query);
+            StringAssert.Contains(expectedApiString, retriedRequest.Uri.Query);
+
+            var apiStringFirstIndex = retriedRequest.Uri.Query.IndexOf(expectedApiString, StringComparison.Ordinal);
+            var apiStringLastIndex = retriedRequest.Uri.Query.LastIndexOf(expectedApiString, StringComparison.Ordinal);
+            Assert.AreEqual(apiStringFirstIndex, apiStringLastIndex);
+        }
+
         [Test]
         public async Task AuthorizationHeaderFormat()
         {
@@ -724,12 +750,13 @@ public async Task CustomHeadersAreAdded()
 
             AssertRequestCommon(request);
             Assert.IsTrue(request.Headers.TryGetValue("x-ms-client-request-id", out string clientRequestId));
-            Assert.AreEqual(clientRequestId, "CustomRequestId");
             Assert.IsTrue(request.Headers.TryGetValue("x-ms-correlation-request-id", out string correlationRequestId));
-            Assert.AreEqual(correlationRequestId, "CorrelationRequestId");
             Assert.IsTrue(request.Headers.TryGetValue("correlation-context", out string correlationContext));
-            Assert.AreEqual(correlationContext, "CorrelationContextValue1,CorrelationContextValue2");
             Assert.IsFalse(request.Headers.TryGetValue("x-ms-random-id", out string randomId));
+
+            Assert.AreEqual(clientRequestId, "CustomRequestId");
+            Assert.AreEqual(correlationRequestId, "CorrelationRequestId");
+            Assert.AreEqual(correlationContext, "CorrelationContextValue1,CorrelationContextValue2");
         }
 
         [Test]
@@ -741,7 +768,7 @@ public async Task AuthorizationHeadersAddedOnceWithRetries()
             var mockTransport = new MockTransport(new MockResponse(503), response);
             ConfigurationClient service = CreateTestService(mockTransport);
 
-            ConfigurationSetting setting = await service.GetConfigurationSettingAsync(s_testSetting.Key, s_testSetting.Label);
+            await service.GetConfigurationSettingAsync(s_testSetting.Key, s_testSetting.Label);
 
             var retriedRequest = mockTransport.Requests[1];
 
@@ -754,6 +781,39 @@ public async Task AuthorizationHeadersAddedOnceWithRetries()
             Assert.AreEqual(1, authorizationHeaders.Count());
         }
 
+        [Test]
+        public async Task CustomHeadersAreAddedOnceWithRetries()
+        {
+            var response = new MockResponse(200);
+            response.SetContent(SerializationHelpers.Serialize(s_testSetting, SerializeSetting));
+
+            var mockTransport = new MockTransport(new MockResponse(503), response);
+            ConfigurationClient service = CreateTestService(mockTransport);
+
+            var activity = new Activity("Azure.CustomDiagnosticHeaders");
+
+            activity.Start();
+            activity.AddTag("x-ms-client-request-id", "CustomRequestId");
+            activity.AddTag("x-ms-correlation-request-id", "CorrelationRequestId");
+            activity.AddTag("correlation-context", "CorrelationContextValue1,CorrelationContextValue2");
+            activity.AddTag("x-ms-random-id", "RandomValue");
+
+            await service.SetConfigurationSettingAsync(s_testSetting);
+            activity.Stop();
+
+            var retriedRequest = mockTransport.Requests[1];
+
+            AssertRequestCommon(retriedRequest);
+            Assert.IsTrue(retriedRequest.Headers.TryGetValues("x-ms-client-request-id", out var clientRequestIds));
+            Assert.IsTrue(retriedRequest.Headers.TryGetValues("x-ms-correlation-request-id", out var correlationRequestIds));
+            Assert.IsTrue(retriedRequest.Headers.TryGetValues("correlation-context", out var correlationContexts));
+            Assert.IsFalse(retriedRequest.Headers.TryGetValues("x-ms-random-id", out var randomId));
+
+            Assert.AreEqual(1, clientRequestIds.Count());
+            Assert.AreEqual(1, correlationRequestIds.Count());
+            Assert.AreEqual(1, correlationContexts.Count());
+        }
+
 
         private void AssertContent(byte[] expected, MockRequest request, bool compareAsString = true)
         {