add adjustments for PlaybackResponseTime

Azure · Feb 13, 2023 · 6bed92f · 6bed92f
1 parent 5ba9d1d
commit 6bed92f
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 5 deletions.
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 playack, 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
@@ -7,6 +7,7 @@
 using Azure.Sdk.Tools.TestProxy.Vendored;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Features;
 using Microsoft.Extensions.Primitives;
 using System;
 using System.Collections.Concurrent;
@@ -466,10 +467,15 @@ public async Task HandlePlaybackRequest(string recordingId, HttpRequest incoming
 
                 outgoingResponse.ContentLength = bodyData.Length;
 
-                await outgoingResponse.Body.WriteAsync(bodyData).ConfigureAwait(false);
+                await WriteBodyBytes(bodyData, session.PlaybackResponseTime, outgoingResponse);
             }
         }
 
+        public async Task WriteBodyBytes(byte[] bodyData, int playbackResponseTime, HttpResponse outgoingResponse)
+        {
+            await outgoingResponse.Body.WriteAsync(bodyData).ConfigureAwait(false);
+        }
+
         public static async Task<(RecordEntry, byte[])> CreateEntryAsync(HttpRequest request)
         {
             var entry = new RecordEntry();
@@ -710,10 +716,29 @@ 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
+                    };
+                }
+                else
+                {
+                    throw new HttpException(HttpStatusCode.BadRequest, $"Unable to set a transport customization on a recording session that is not active. Id: \"{sessionId}\"");
+                }
+
+                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
             {