From 935969aa9142da6d87d58480d8ebc573827e6607 Mon Sep 17 00:00:00 2001
From: Ivan Zlatanov <ivan@zlatanov.net>
Date: Mon, 19 Apr 2021 11:20:00 +0300
Subject: [PATCH] Addressing flakiness of a couple of tests.

---
 .../tests/WebSocketDeflateTests.cs            | 26 ++++++++++++-------
 .../tests/ZLibStreamTests.cs                  |  8 +++---
 2 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/src/libraries/System.Net.WebSockets/tests/WebSocketDeflateTests.cs b/src/libraries/System.Net.WebSockets/tests/WebSocketDeflateTests.cs
index 74bb00e2ef22a5..33ac5122516cc5 100644
--- a/src/libraries/System.Net.WebSockets/tests/WebSocketDeflateTests.cs
+++ b/src/libraries/System.Net.WebSockets/tests/WebSocketDeflateTests.cs
@@ -4,6 +4,7 @@
 using System.Buffers;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.IO;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
@@ -493,19 +494,23 @@ public async Task ReceiveInvalidCompressedData()
             Assert.Equal(WebSocketState.Aborted, client.State);
         }
 
-        [Fact]
-        public async Task PayloadShouldHaveSimilarSizeWhenSplitIntoSegments()
+        [Theory]
+        [MemberData(nameof(SupportedWindowBits))]
+        public async Task PayloadShouldHaveSimilarSizeWhenSplitIntoSegments(int windowBits)
         {
-            WebSocketTestStream stream = new();
+            MemoryStream stream = new();
             WebSocket client = WebSocket.CreateFromStream(stream, new WebSocketCreationOptions
             {
                 DangerousDeflateOptions = new WebSocketDeflateOptions()
+                {
+                    ClientMaxWindowBits = windowBits
+                }
             });
 
             // We're using a frame size that is close to the sliding window size for the deflate
-            const int frameSize = 32_000;
+            int frameSize = 2 << windowBits;
 
-            byte[] message = new byte[frameSize * 100];
+            byte[] message = new byte[frameSize * 10];
             Random random = new(0);
 
             for (int i = 0; i < message.Length; ++i)
@@ -515,15 +520,18 @@ public async Task PayloadShouldHaveSimilarSizeWhenSplitIntoSegments()
 
             await client.SendAsync(message, WebSocketMessageType.Binary, true, CancellationToken);
 
-            int payloadLength = stream.Remote.Available;
-            stream.Remote.Clear();
+            long payloadLength = stream.Length;
+            stream.SetLength(0);
 
-            for (var i = 0; i < message.Length; i += frameSize)
+            for (int i = 0; i < message.Length; i += frameSize)
             {
                 await client.SendAsync(message.AsMemory(i, frameSize), WebSocketMessageType.Binary, i + frameSize == message.Length, CancellationToken);
             }
 
-            Assert.Equal(0.999, Math.Round(payloadLength * 1.0 / stream.Remote.Available, 3));
+            double difference = Math.Round(1 - payloadLength * 1.0 / stream.Length, 3);
+
+            // The difference should not be more than 10% in either direction
+            Assert.InRange(difference, -0.1, 0.1);
         }
 
         [Theory]
diff --git a/src/libraries/System.Net.WebSockets/tests/ZLibStreamTests.cs b/src/libraries/System.Net.WebSockets/tests/ZLibStreamTests.cs
index eb0ad93e3770e5..d12bf696f9be9a 100644
--- a/src/libraries/System.Net.WebSockets/tests/ZLibStreamTests.cs
+++ b/src/libraries/System.Net.WebSockets/tests/ZLibStreamTests.cs
@@ -13,7 +13,7 @@ public class ZLibStreamTests
         [Fact]
         public async Task PoolShouldReuseTheSameInstance()
         {
-            var pool = new Pool(timeoutMilliseconds: 100);
+            var pool = new Pool(timeoutMilliseconds: 25);
 
             object inflater = pool.GetInflater();
             for ( var i = 0; i < 10_000; ++i)
@@ -29,7 +29,7 @@ public async Task PoolShouldReuseTheSameInstance()
             pool.ReturnInflater(inflater);
 
             Assert.Equal(1, pool.ActiveCount);
-            await Task.Delay(250);
+            await Task.Delay(200);
 
             // After timeout elapses we should not have any active instances
             Assert.Equal(0, pool.ActiveCount);
@@ -39,7 +39,7 @@ public async Task PoolShouldReuseTheSameInstance()
         [PlatformSpecific(~TestPlatforms.Browser)] // There is no concurrency in browser
         public async Task PoolingConcurrently()
         {
-            var pool = new Pool(timeoutMilliseconds: 100);
+            var pool = new Pool(timeoutMilliseconds: 25);
             var parallelOptions = new ParallelOptions
             {
                 MaxDegreeOfParallelism = 16
@@ -60,7 +60,7 @@ public async Task PoolingConcurrently()
 
             Assert.True(pool.ActiveCount >= 2);
             Assert.True(pool.ActiveCount <= parallelOptions.MaxDegreeOfParallelism * 2);
-            await Task.Delay(250);
+            await Task.Delay(200);
             Assert.Equal(0, pool.ActiveCount);
         }