From c9526d5c27bb072b447dc0011fa635f4bc754df4 Mon Sep 17 00:00:00 2001
From: Krzysztof Nozderko <krzysztof.nozderko@snowflake.com>
Date: Wed, 3 Apr 2024 16:51:59 +0200
Subject: [PATCH] Fix flaky test TestWaitUntilResourceAvailable

---
 .../UnitTests/Session/WaitingQueueTest.cs     | 32 ++++++++++++-------
 Snowflake.Data.Tests/Util/TestRepeater.cs     | 32 +++++++++++++++++++
 2 files changed, 52 insertions(+), 12 deletions(-)
 create mode 100644 Snowflake.Data.Tests/Util/TestRepeater.cs

diff --git a/Snowflake.Data.Tests/UnitTests/Session/WaitingQueueTest.cs b/Snowflake.Data.Tests/UnitTests/Session/WaitingQueueTest.cs
index 0c567234b..3de525b06 100644
--- a/Snowflake.Data.Tests/UnitTests/Session/WaitingQueueTest.cs
+++ b/Snowflake.Data.Tests/UnitTests/Session/WaitingQueueTest.cs
@@ -1,8 +1,10 @@
+using System;
 using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 using NUnit.Framework;
 using Snowflake.Data.Core.Session;
+using Snowflake.Data.Tests.Util;
 
 namespace Snowflake.Data.Tests.UnitTests.Session
 {
@@ -48,22 +50,28 @@ public void TestWaitForTheResourceUntilCancellation()
         public void TestWaitUntilResourceAvailable()
         {
             // arrange
-            var queue = new WaitingQueue();
-            var watch = new Stopwatch();
-            Task.Run(() =>
+            var tests = TestRepeater<Tuple<bool, long>>.Test(3, () =>
             {
-                Thread.Sleep(50);
-                queue.OnResourceIncrease();
+                var queue = new WaitingQueue();
+                var watch = new Stopwatch();
+                Task.Run(() =>
+                {
+                    Thread.Sleep(50);
+                    queue.OnResourceIncrease();
+                });
+
+                // act
+                watch.Start();
+                var result = queue.Wait(30000, CancellationToken.None);
+                watch.Stop();
+                return Tuple.Create(result, watch.ElapsedMilliseconds);
             });
-            
-            // act
-            watch.Start();
-            var result = queue.Wait(30000, CancellationToken.None);
-            watch.Stop();
 
             // assert
-            Assert.IsTrue(result);
-            Assert.That(watch.ElapsedMilliseconds, Is.InRange(50, 1500));
+            tests.ForEach(t => Assert.IsTrue(t.Item1));
+            tests.ForEach(t => Assert.GreaterOrEqual(t.Item2, 50));
+            tests.SkipLargest(t => t.Item2) // some execution can be randomly delayed so we skip the largest value
+                .ForEach(t => Assert.That(t.Item2, Is.InRange(50, 1500)));
         }
 
         [Test]
diff --git a/Snowflake.Data.Tests/Util/TestRepeater.cs b/Snowflake.Data.Tests/Util/TestRepeater.cs
new file mode 100644
index 000000000..31c29b60f
--- /dev/null
+++ b/Snowflake.Data.Tests/Util/TestRepeater.cs
@@ -0,0 +1,32 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+
+namespace Snowflake.Data.Tests.Util
+{
+    public class TestRepeater<T>
+    {
+        private readonly List<T> _result;
+
+        private TestRepeater(List<T> result)
+        {
+            _result = result;
+        }
+
+        public void ForEach(Action<T> action) => _result.ForEach(action);
+        
+        public TestRepeater<T> SkipLargest<TKey>(Func<T, TKey> keySelector)
+        {
+            var resultsWithoutLargest = _result.OrderBy(keySelector).SkipLast(1).ToList();
+            return new TestRepeater<T>(resultsWithoutLargest);
+        }
+        
+        public static TestRepeater<T> Test(int times, Func<T> testFunction)
+        {
+            var resultList = Enumerable.Repeat(0, times)
+                .Select(_ => testFunction())
+                .ToList();
+            return new TestRepeater<T>(resultList);
+        }
+    }
+}