Fix test flakiness caused by not reading in a loop (#3023)

Fixes #3021

test/Kestrel.InMemory.FunctionalTests/RequestTests.cs

@@ -67,10 +67,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
                 var buffer = new byte[1024];
                 try
                 {
+                    await context.Request.Body.ReadUntilLengthAsync(buffer, 6, cts.Token).DefaultTimeout();
 
-                    int read = await context.Request.Body.ReadAsync(buffer, 0, buffer.Length, cts.Token);
+                    Assert.Equal("Hello ", Encoding.ASCII.GetString(buffer, 0, 6));
-
-                    Assert.Equal("Hello ", Encoding.UTF8.GetString(buffer, 0, read));
 
                     helloTcs.TrySetResult(null);
                 }
@@ -82,7 +81,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
 
                 try
                 {
-                    await context.Request.Body.ReadAsync(buffer, 0, buffer.Length, cts.Token);
+                    await context.Request.Body.ReadAsync(buffer, 0, 1, cts.Token).DefaultTimeout();
 
                     context.Response.ContentLength = 12;
                     await context.Response.WriteAsync("Read success");
@@ -133,12 +132,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
             {
                 var stream = await context.Features.Get<IHttpUpgradeFeature>().UpgradeAsync();
                 var data = new byte[3];
-                var bytesRead = 0;
 
-                while (bytesRead < 3)
+                await stream.ReadUntilLengthAsync(data, 3).DefaultTimeout();
-                {
-                    bytesRead += await stream.ReadAsync(data, bytesRead, data.Length - bytesRead);
-                }
 
                 dataRead = Encoding.ASCII.GetString(data, 0, 3) == "abc";
             }, new TestServiceContext(LoggerFactory)))
@@ -284,9 +279,18 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
                            "",
                            "");
 
-                        var read = await connection.Reader.ReadAsync(buffer, 0, identifierLength);
+                        var offset = 0;
-                        Assert.Equal(identifierLength, read);
+
-                        var id = new string(buffer, 0, read);
+                        while (offset < identifierLength)
+                        {
+                            var read = await connection.Reader.ReadAsync(buffer, offset, identifierLength - offset);
+                            offset += read;
+
+                            Assert.NotEqual(0, read);
+                        }
+
+                        Assert.Equal(identifierLength, offset);
+                        var id = new string(buffer, 0, offset);
                         Assert.DoesNotContain(id, usedIds.ToArray());
                         usedIds.Add(id);
                     }
@@ -670,13 +674,10 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
                 var duplexStream = await upgradeFeature.UpgradeAsync();
 
                 var buffer = new byte[message.Length];
-                var read = 0;
-                while (read < message.Length)
-                {
-                    read += await duplexStream.ReadAsync(buffer, read, buffer.Length - read).DefaultTimeout();
-                }
 
-                await duplexStream.WriteAsync(buffer, 0, read);
+                await duplexStream.ReadUntilLengthAsync(buffer, message.Length).DefaultTimeout();
+
+                await duplexStream.WriteAsync(buffer, 0, buffer.Length);
             }, testContext))
             {
                 using (var connection = server.CreateConnection())
@@ -1094,14 +1095,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
                 Assert.Equal(CoreStrings.SynchronousReadsDisallowed, ioEx2.Message);
 
                 var buffer = new byte[5];
-                var offset = 0;
+                var read = await context.Request.Body.ReadUntilEndAsync(buffer).DefaultTimeout();
-                while (offset < 5)
-                {
-                    offset += await context.Request.Body.ReadAsync(buffer, offset, 5 - offset);
-                }
 
-                Assert.Equal(0, await context.Request.Body.ReadAsync(new byte[1], 0, 1));
+                Assert.Equal("Hello", Encoding.ASCII.GetString(buffer, 0, read));
-                Assert.Equal("Hello", Encoding.ASCII.GetString(buffer));
             }, testContext))
             {
                 using (var connection = server.CreateConnection())

test/shared/StreamExtensions.cs

@@ -0,0 +1,45 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.Threading;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace System.IO
+{
+    public static class StreamFillBufferExtensions
+    {
+        public static async Task<int> ReadUntilEndAsync(this Stream stream, byte[] buffer, CancellationToken cancellationToken = default)
+        {
+            var offset = 0;
+
+            while (offset < buffer.Length)
+            {
+                var read = await stream.ReadAsync(buffer, offset, buffer.Length - offset, cancellationToken);
+                offset += read;
+
+                if (read == 0)
+                {
+                    return offset;
+                }
+            }
+
+            Assert.Equal(0, await stream.ReadAsync(new byte[1], 0, 1, cancellationToken));
+
+            return offset;
+        }
+
+        public static async Task ReadUntilLengthAsync(this Stream stream, byte[] buffer, int length, CancellationToken cancellationToken = default)
+        {
+            var offset = 0;
+
+            while (offset < length)
+            {
+                var read = await stream.ReadAsync(buffer, offset, length - offset, cancellationToken);
+                offset += read;
+
+                Assert.NotEqual(0, read);
+            }
+        }
+    }
+}

test/shared/TestApp.cs

@@ -14,16 +14,10 @@ namespace Microsoft.AspNetCore.Testing
             var request = httpContext.Request;
             var response = httpContext.Response;
             var buffer = new byte[httpContext.Request.ContentLength ?? 0];
-            var bytesRead = 0;
-
-            while (bytesRead < buffer.Length)
-            {
-                var count = await request.Body.ReadAsync(buffer, bytesRead, buffer.Length - bytesRead);
-                bytesRead += count;
-            }
 
             if (buffer.Length > 0)
             {
+                await request.Body.ReadUntilEndAsync(buffer).DefaultTimeout();
                 await response.Body.WriteAsync(buffer, 0, buffer.Length);
             }
         }