Continuation Sanitizer Bugfix

tools/test-proxy/Azure.Sdk.Tools.TestProxy.Tests/Azure.Sdk.Tools.TestProxy.Tests.csproj

@@ -27,6 +27,9 @@
     <None Update="Test.RecordEntries\post_delete_get_content.json">
       <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
     </None>
+    <None Update="Test.RecordEntries\requests_with_continuation.json">
+      <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
+    </None>
     <None Update="Test.RecordEntries\request_with_subscriptionid.json">
       <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
     </None>

tools/test-proxy/Azure.Sdk.Tools.TestProxy.Tests/SanitizerTests.cs

@@ -330,21 +330,66 @@ public void BodyKeySanitizerQuietlyExits()
 
 
         [Fact]
-        public void ContinuationSanitizerMultipleStepsNoKey()
+        public void ContinuationSanitizerSingleReplace()
         {
-
+            var session = TestHelpers.LoadRecordSession("Test.RecordEntries/requests_with_continuation.json");
+            var continueSanitizer = new ContinuationSanitizer("correlationId", "guid", resetAfterFirst: "false");
+            var targetKey = "correlationId";
+            var originalRequestGuid = session.Session.Entries[0].Response.Headers[targetKey].First();
+
+            session.Session.Sanitize(continueSanitizer);
+
+            var firstRequest = session.Session.Entries[0].Response.Headers[targetKey].First();
+            var firstResponse = session.Session.Entries[1].Request.Headers[targetKey].First();
+            var secondRequest = session.Session.Entries[2].Response.Headers[targetKey].First();
+            var secondResponse = session.Session.Entries[3].Request.Headers[targetKey].First();
+
+            Assert.NotEqual(originalRequestGuid, firstRequest);
+            Assert.Equal(firstRequest, firstResponse);
+            Assert.NotEqual(firstRequest, secondRequest);
+            Assert.Equal(firstResponse, secondResponse);
         }
 
         [Fact]
-        public void ContinuationSanitizerSimpleRequestResponse()
+        public void ContinuationSanitizerMultipleReplace()
         {
-
+            var session = TestHelpers.LoadRecordSession("Test.RecordEntries/requests_with_continuation.json");
+            var continueSanitizer = new ContinuationSanitizer("correlationId", "guid", resetAfterFirst: "true");
+            var targetKey = "correlationId";
+            var originalSendGuid = session.Session.Entries[0].Response.Headers[targetKey].First();
+
+            session.Session.Sanitize(continueSanitizer);
+
+            var firstRequest = session.Session.Entries[0].Response.Headers[targetKey].First();
+            var firstResponse = session.Session.Entries[1].Request.Headers[targetKey].First();
+            var secondRequest = session.Session.Entries[2].Response.Headers[targetKey].First();
+            var secondResponse = session.Session.Entries[3].Request.Headers[targetKey].First();
+
+            Assert.NotEqual(originalSendGuid, firstRequest);
+            Assert.Equal(firstRequest, firstResponse);
+            Assert.NotEqual(firstResponse, secondRequest);
+            Assert.Equal(secondRequest, secondResponse);
         }
 
         [Fact]
-        public void ContinuationSanitizerMultipleSteps()
+        public void ContinuationSanitizerNonExistentKey()
         {
-
+            var session = TestHelpers.LoadRecordSession("Test.RecordEntries/requests_with_continuation.json");
+            var continueSanitizer = new ContinuationSanitizer("non-existent-key", "guid", resetAfterFirst: "true");
+            var targetKey = "correlationId";
+            var originalSendGuid = session.Session.Entries[0].Response.Headers[targetKey].First();
+
+            session.Session.Sanitize(continueSanitizer);
+
+            var firstRequest = session.Session.Entries[0].Response.Headers[targetKey].First();
+            var firstResponse = session.Session.Entries[1].Request.Headers[targetKey].First();
+            var secondRequest = session.Session.Entries[2].Response.Headers[targetKey].First();
+            var secondResponse = session.Session.Entries[3].Request.Headers[targetKey].First();
+
+            Assert.Equal(originalSendGuid, firstRequest);
+            Assert.Equal(firstRequest, firstResponse);
+            Assert.NotEqual(firstResponse, secondRequest);
+            Assert.Equal(secondRequest, secondResponse);
         }
     }
 }

tools/test-proxy/Azure.Sdk.Tools.TestProxy.Tests/Test.RecordEntries/requests_with_continuation.json

@@ -0,0 +1,79 @@
+{
+  "Entries": [
+    {
+      "RequestUri": "http://localhost:3000/request",
+      "RequestMethod": "GET",
+      "RequestHeaders": {
+        "Accept-Encoding": "gzip,deflate",
+        "Connection": "keep-alive"
+      },
+      "RequestBody": null,
+      "StatusCode": 200,
+      "ResponseHeaders": {
+        "Connection": "keep-alive",
+        "Content-Length": "0",
+        "Date": "Wed, 20 Oct 2021 22:45:39 GMT",
+        "correlationId": "a34cafcb-ea7f-4d0f-b5ca-83c714761a50"
+      },
+      "ResponseBody": null
+    },
+    {
+      "RequestUri": "http://localhost:3000/send",
+      "RequestMethod": "POST",
+      "RequestHeaders": {
+        "Accept-Encoding": "gzip,deflate",
+        "Connection": "keep-alive",
+        "correlationId": "a34cafcb-ea7f-4d0f-b5ca-83c714761a50"
+      },
+      "RequestBody": null,
+      "StatusCode": 200,
+      "ResponseHeaders": {
+        "Connection": "keep-alive",
+        "Content-Length": "13",
+        "Content-Type": "application/json; charset=utf-8",
+        "Date": "Wed, 20 Oct 2021 22:45:39 GMT"
+      },
+      "ResponseBody": {
+        "val": "abc"
+      }
+    },
+    {
+      "RequestUri": "http://localhost:3000/request",
+      "RequestMethod": "GET",
+      "RequestHeaders": {
+        "Accept-Encoding": "gzip,deflate",
+        "Connection": "keep-alive"
+      },
+      "RequestBody": null,
+      "StatusCode": 200,
+      "ResponseHeaders": {
+        "Connection": "keep-alive",
+        "Content-Length": "0",
+        "Date": "Wed, 20 Oct 2021 22:45:39 GMT",
+        "correlationId": "b24f75a9-b830-11eb-b949-10e7c6392c5a"
+      },
+      "ResponseBody": null
+    },
+    {
+      "RequestUri": "http://localhost:3000/send",
+      "RequestMethod": "GET",
+      "RequestHeaders": {
+        "Accept-Encoding": "gzip,deflate",
+        "Connection": "keep-alive",
+        "correlationId": "b24f75a9-b830-11eb-b949-10e7c6392c5a"
+      },
+      "RequestBody": null,
+      "StatusCode": 200,
+      "ResponseHeaders": {
+        "Connection": "keep-alive",
+        "Content-Length": "13",
+        "Content-Type": "application/json; charset=utf-8",
+        "Date": "Wed, 20 Oct 2021 22:45:39 GMT"
+      },
+      "ResponseBody": {
+        "val": "abc"
+      }
+    }
+  ],
+  "Variables": {}
+}

tools/test-proxy/Azure.Sdk.Tools.TestProxy/Sanitizers/ContinuationSanitizer.cs

@@ -10,9 +10,9 @@
 namespace Azure.Sdk.Tools.TestProxy.Sanitizers
 {
     /// <summary>
-    /// Used to sanitize "session" variables. A "session" variable is one that is returned as a result of Response A, then used as INPUT of Response B.
+    /// Used to sanitize "session" variables. A "session" variable is one that is returned as a result of Response A, then used as INPUT of Response B. 
     /// 
-    /// The value inserted defaults to a new guid if one is not provided. This sanitizer applies at the session level.
+    /// The value inserted defaults to a new guid if one is not provided. This sanitizer applies at the session level, and looks at the HEADERS of the request/response pairs.
     /// </summary>
     public class ContinuationSanitizer : RecordedTestSanitizer
     {
@@ -33,15 +33,17 @@ public static string GuidReplacer()
         /// <summary>
         /// This sanitizer is applied at the session level, and is used to anonymize private keys in response/request pairs.
         /// 
-        /// For instance, a request hands back a "sessionId" that needs to be present in the next request.
+        /// For instance, request/response pair A has a RESPONSE that contains a "sessionId" header whos value must be present in the REQUEST of the next request/response pair B.
         /// 
-        /// Supports "all further requests get this key" as well as "single response/request pair". Defaults to maintaining same key 
-        /// for rest of recording.
+        /// This sanitizer supports "all further requests get this the value of this first key" as well as as the more standard "response value gets used in next request" that is described 
+        /// in the scenario before. 
+        /// 
+        /// Defaults to maintaining same key for rest of recording.
         /// </summary>
         /// <param name="key">The name of the header whos value will be replaced from response -> next request.</param>
-        /// <param name="method">The method by which the value of the targeted key will be replaced. Defaults to guid replacement.</param>
-        /// <param name="resetAfterFirst">Do we need multiple pairs replaced? Or do we want to replace each value with the same value.</param>
-        public ContinuationSanitizer(string key, string method, string resetAfterFirst = "false")
+        /// <param name="method">The method by which the value of the targeted key will be replaced. Currently only supports generating a new guid for the value replacement..</param>
+        /// <param name="resetAfterFirst">Do we need multiple pairs replaced? Or do we want to stop after the first set.</param>
+        public ContinuationSanitizer(string key, string method, string resetAfterFirst = "true")
         {
             _targetKey = key;
             _method = method;
@@ -72,7 +74,7 @@ public override void Sanitize(RecordSession session)
 
                 if (currentEntry.Request.Headers.ContainsKey(_targetKey) && !String.IsNullOrWhiteSpace(newValue))
                 {
-                    currentEntry.Response.Headers[_targetKey] = new string[] { newValue };
+                    currentEntry.Request.Headers[_targetKey] = new string[] { newValue };
 
                     if (_resetAfterFirst)
                     {