From 67c742b64acb3167bb5fd83ff3df4982252569eb Mon Sep 17 00:00:00 2001
From: scbedd <45376673+scbedd@users.noreply.github.com>
Date: Tue, 14 Feb 2023 18:10:44 -0800
Subject: [PATCH] add function and test to batch the body data

---
 .../RecordingHandlerTests.cs                  | 22 +++++++++++----
 .../RecordingHandler.cs                       | 27 ++++++++++++++++---
 2 files changed, 41 insertions(+), 8 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
index 7a0c8de9eda..05330b4d519 100644
--- a/tools/test-proxy/Azure.Sdk.Tools.TestProxy.Tests/RecordingHandlerTests.cs
+++ b/tools/test-proxy/Azure.Sdk.Tools.TestProxy.Tests/RecordingHandlerTests.cs
@@ -780,14 +780,27 @@ public async Task RecordingHandlerIsThreadSafe()
         }
 
         #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("", 100, 1)]
-        public void TestGetBatches(string input, int playbackResponseTime, int batchCount)
+        [InlineData("", 1)]
+        [InlineData("small body", 10)]
+        [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, playbackResponseTime, batchCount);
+            var chunks = testRecordingHandler.GetBatches(bodyData, batchCount);
 
             int bodyPosition = 0;
 
@@ -797,9 +810,8 @@ public void TestGetBatches(string input, int playbackResponseTime, int batchCoun
                 for (int j = 0; j < chunk.Length; j++)
                 {
                     Assert.Equal(chunk[j], bodyData[bodyPosition]);
+                    bodyPosition++;
                 }
-
-                bodyPosition++;
             }
         }
 
diff --git a/tools/test-proxy/Azure.Sdk.Tools.TestProxy/RecordingHandler.cs b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/RecordingHandler.cs
index 7ae16ea5778..fd5cbaf4b34 100644
--- a/tools/test-proxy/Azure.Sdk.Tools.TestProxy/RecordingHandler.cs
+++ b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/RecordingHandler.cs
@@ -473,9 +473,30 @@ public async Task HandlePlaybackRequest(string recordingId, HttpRequest incoming
             }
         }
 
-        public byte[][] GetBatches(byte[] bodyData, int playbackResponseTime, int batchCount)
+        public byte[][] GetBatches(byte[] bodyData, int batchCount)
         {
-            return null;
+            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 batch = new byte[chunkLength];
+
+                Array.Copy(bodyData, i * chunkLength, batch, 0, chunkLength);
+
+                batches[i] = batch;
+            }
+
+            return batches;
         }
 
         public async Task WriteBodyBytes(byte[] bodyData, int playbackResponseTime, HttpResponse outgoingResponse)
@@ -485,7 +506,7 @@ public async Task WriteBodyBytes(byte[] bodyData, int playbackResponseTime, Http
                 int batchCount = 10;
                 int sleepLength = playbackResponseTime / batchCount;
 
-                byte[][] chunks = GetBatches(bodyData, playbackResponseTime, batchCount);
+                byte[][] chunks = GetBatches(bodyData, batchCount);
 
                 for(int i = 0; i < chunks.Length; i++)
                 {