Add Delayed Response (#5454)

* add adjustments for new Transport property PlaybackResponseTime being honored in RecordingOptions * add chunking algo for splitting up a body across multiple writes * add test to ensure that chunking algo doesn't accidentally omit bytes during send Co-authored-by: Mike Harder <[email protected]>
Azure · Feb 24, 2023 · ffbfbee · ffbfbee
1 parent d695d83
commit ffbfbee
Show file tree

Hide file tree

Showing 4 changed files with 129 additions and 6 deletions.
diff --git a/tools/test-proxy/Azure.Sdk.Tools.TestProxy.Tests/RecordingHandlerTests.cs b/tools/test-proxy/Azure.Sdk.Tools.TestProxy.Tests/RecordingHandlerTests.cs
@@ -779,6 +779,46 @@ public async Task RecordingHandlerIsThreadSafe()
             Assert.Equal(requestCount, session.Session.Entries.Count);
         }
 
+        #region ByteManipulation
+
+        private const string longBody = @"Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt
+            ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut
+            aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore
+            eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt
+            mollit anim id est laborum.";
+
+        [Theory]
+        [InlineData("", 1)]
+        [InlineData("small body", 5)]
+        [InlineData("this is a body", 3)]
+        [InlineData("This is a little bit longer of a body that we are dividing in 2", 2)]
+        [InlineData(longBody, 5)]
+        [InlineData(longBody, 1)]
+        [InlineData(longBody, 10)]
+        public void TestGetBatches(string input, int batchCount)
+        {
+            RecordingHandler testRecordingHandler = new RecordingHandler(Directory.GetCurrentDirectory());
+            var bodyData = Encoding.UTF8.GetBytes(input);
+
+            var chunks = testRecordingHandler.GetBatches(bodyData, batchCount);
+
+            int bodyPosition = 0;
+
+            // ensure that all bytes are accounted for across the batches
+            foreach(var chunk in chunks)
+            {
+                for (int j = 0; j < chunk.Length; j++)
+                {
+                    Assert.Equal(chunk[j], bodyData[bodyPosition]);
+                    bodyPosition++;
+                }
+            }
+
+            Assert.Equal(bodyPosition, bodyData.Length);
+        }
+
+        #endregion
+
         #region SetRecordingOptions
         [Theory]
         [InlineData("{ \"HandleRedirects\": \"true\"}", true)]

diff --git a/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/ModifiableRecordSession.cs b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/ModifiableRecordSession.cs
@@ -1,4 +1,4 @@
-using System;
+using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Net.Http;
@@ -27,6 +27,8 @@ public ModifiableRecordSession(RecordSession session)
 
         public string SourceRecordingId { get; set; }
 
+        public int PlaybackResponseTime { get; set; }
+
         public void ResetExtensions()
         {
             AdditionalTransforms.Clear();

diff --git a/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/TransportCustomizations.cs b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/TransportCustomizations.cs
@@ -31,5 +31,11 @@ public class TransportCustomizations
         /// Each certificate pair contained within this list should be added to the clientHandler for the server or an individual recording.
         /// </summary>
         public List<PemPair> Certificates { get; set; }
+
+        /// <summary>
+        /// During playback, a response is normally returned all at once. By offering this response time, we can
+        /// "stretch" the writing of the response bytes over a time range of milliseconds.
+        /// </summary>
+        public int PlaybackResponseTime { get; set; } = 0;
     }
 }
diff --git a/tools/test-proxy/Azure.Sdk.Tools.TestProxy/RecordingHandler.cs b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/RecordingHandler.cs
@@ -268,7 +268,10 @@ public async Task HandleRecordRequestAsync(string recordingId, HttpRequest incom
             if (entry.Response.Body?.Length > 0)
             {
                 var bodyData = CompressionUtilities.CompressBody(entry.Response.Body, entry.Response.Headers);
-                outgoingResponse.ContentLength = bodyData.Length;
+
+                if (entry.Response.Headers.ContainsKey("Content-Length")){
+                    outgoingResponse.ContentLength = bodyData.Length;
+                }
                 await outgoingResponse.Body.WriteAsync(bodyData).ConfigureAwait(false);
             }
         }
@@ -464,8 +467,65 @@ public async Task HandlePlaybackRequest(string recordingId, HttpRequest incoming
             {
                 var bodyData = CompressionUtilities.CompressBody(match.Response.Body, match.Response.Headers);
 
-                outgoingResponse.ContentLength = bodyData.Length;
+                if (match.Response.Headers.ContainsKey("Content-Length"))
+                {
+                    outgoingResponse.ContentLength = bodyData.Length;
+                }
+
+                await WriteBodyBytes(bodyData, session.PlaybackResponseTime, outgoingResponse);
+            }
+        }
+
+        public byte[][] GetBatches(byte[] bodyData, int batchCount)
+        {
+            if (bodyData.Length == 0 || bodyData.Length < batchCount)
+            {
+                var result = new byte[1][];
+                result[0] = bodyData;
+
+                return result;
+            }
+
+            int chunkLength = bodyData.Length / batchCount;
+            int remainder = (bodyData.Length % batchCount);
+            var batches = new byte[batchCount + (remainder > 0 ? 1 : 0)][];
+
+            for(int i = 0; i < batches.Length; i++)
+            {
+                var calculatedChunkLength = ((i == batches.Length - 1) && (batches.Length > 1) && (remainder > 0)) ? remainder : chunkLength;
+                var batch = new byte[calculatedChunkLength];
+                Array.Copy(bodyData, i * chunkLength, batch, 0, calculatedChunkLength);
+
+                batches[i] = batch;
+            }
+
+            return batches;
+        }
+
+        public async Task WriteBodyBytes(byte[] bodyData, int playbackResponseTime, HttpResponse outgoingResponse)
+        {
+            if (playbackResponseTime > 0)
+            {
+                int batchCount = 10;
+                int sleepLength = playbackResponseTime / batchCount;
+
+                byte[][] chunks = GetBatches(bodyData, batchCount);
 
+                for(int i = 0; i < chunks.Length; i++)
+                {
+                    var chunk = chunks[i];
+
+                    await outgoingResponse.Body.WriteAsync(chunk).ConfigureAwait(false);
+
+                    if (i != chunks.Length - 1)
+                    {
+                        await Task.Delay(sleepLength);
+                    }
+                }
+
+            }
+            else
+            {
                 await outgoingResponse.Body.WriteAsync(bodyData).ConfigureAwait(false);
             }
         }
@@ -710,10 +770,25 @@ public void SetTransportOptions(TransportCustomizations customizations, string s
             {
                 var customizedClientHandler = GetTransport(customizations.AllowAutoRedirect, customizations);
 
-                RecordingSessions[sessionId].Client = new HttpClient(customizedClientHandler)
+                if (RecordingSessions.TryGetValue(sessionId, out var recordingSession))
                 {
-                    Timeout = timeoutSpan
-                };
+                    recordingSession.Client = new HttpClient(customizedClientHandler)
+                    {
+                        Timeout = timeoutSpan
+                    };
+                }
+
+                if (customizations.PlaybackResponseTime > 0)
+                {
+                    if (PlaybackSessions.TryGetValue(sessionId, out var playbackSession))
+                    {
+                        playbackSession.PlaybackResponseTime = customizations.PlaybackResponseTime;
+                    }
+                    else
+                    {
+                        throw new HttpException(HttpStatusCode.BadRequest, $"Unable to set a transport customization on a recording session that is not active. Id: \"{sessionId}\"");
+                    }
+                }
             }
             else
             {