diff --git a/tools/test-proxy/Azure.Sdk.Tools.TestProxy.Tests/RecordingHandlerTests.cs b/tools/test-proxy/Azure.Sdk.Tools.TestProxy.Tests/RecordingHandlerTests.cs index ecb358b91c5..725bc3f82a2 100644 --- 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 index f7eef9a7b42..9de32d7326f 100644 --- 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 index 388f271866e..4740e62896a 100644 --- 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. /// public List Certificates { get; set; } + + /// + /// 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. + /// + 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 index 7544c55114c..eb2871238d8 100644 --- 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 {