Fix a flaky Extensions.ML test. (#4458)

* Fix a flaky Extensions.ML test. Make the reload model tests more resistant to timing changes. * PR feedback.
dotnet · Nov 14, 2019 · f3e0f6b · f3e0f6b
1 parent c1e190a
commit f3e0f6b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 20 deletions.
diff --git a/test/Microsoft.Extensions.ML.Tests/FileLoaderTests.cs b/test/Microsoft.Extensions.ML.Tests/FileLoaderTests.cs
@@ -4,7 +4,7 @@
 
 using System;
 using System.IO;
-using System.Threading.Tasks;
+using System.Threading;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
@@ -48,10 +48,8 @@ public void can_load_model()
             Assert.True(prediction.Sentiment);
         }
 
-        //TODO: This is a quick test to give coverage of the main scenarios. Refactoring and re-implementing of tests should happen.
-        //Right now this screams of probably flakeyness
         [Fact]
-        public async Task can_reload_model()
+        public void can_reload_model()
         {
             var services = new ServiceCollection()
                 .AddOptions()
@@ -61,16 +59,14 @@ public async Task can_reload_model()
             var loaderUnderTest = ActivatorUtilities.CreateInstance<FileLoaderMock>(sp);
             loaderUnderTest.Start("testdata.txt", true);
 
-            var changed = false;
-            var changeTokenRegistration = ChangeToken.OnChange(
+            using AutoResetEvent changed = new AutoResetEvent(false);
+            using IDisposable changeTokenRegistration = ChangeToken.OnChange(
                         () => loaderUnderTest.GetReloadToken(),
-                        () => changed = true);
+                        () => changed.Set());
 
             File.WriteAllText("testdata.txt", "test");
 
-            await Task.Delay(1000);
-
-            Assert.True(changed);
+            Assert.True(changed.WaitOne(1000), "FileLoader ChangeToken didn't fire before the allotted time.");
         }
 
 

diff --git a/test/Microsoft.Extensions.ML.Tests/UriLoaderTests.cs b/test/Microsoft.Extensions.ML.Tests/UriLoaderTests.cs
@@ -40,12 +40,12 @@ public void can_reload_model()
             var loaderUnderTest = ActivatorUtilities.CreateInstance<UriLoaderMock>(sp);
             loaderUnderTest.Start(new Uri("http://microsoft.com"), TimeSpan.FromMilliseconds(1));
 
-            var changed = false;
-            var changeTokenRegistration = ChangeToken.OnChange(
+            using AutoResetEvent changed = new AutoResetEvent(false);
+            using IDisposable changeTokenRegistration = ChangeToken.OnChange(
                         () => loaderUnderTest.GetReloadToken(),
-                        () => changed = true);
-            Thread.Sleep(30);
-            Assert.True(changed);
+                        () => changed.Set());
+
+            Assert.True(changed.WaitOne(1000), "UriLoader ChangeToken didn't fire before the allotted time.");
         }
 
         [Fact]
@@ -62,12 +62,12 @@ public void no_reload_no_change()
 
             loaderUnderTest.Start(new Uri("http://microsoft.com"), TimeSpan.FromMilliseconds(1));
 
-            var changed = false;
-            var changeTokenRegistration = ChangeToken.OnChange(
+            using AutoResetEvent changed = new AutoResetEvent(false);
+            using IDisposable changeTokenRegistration = ChangeToken.OnChange(
                         () => loaderUnderTest.GetReloadToken(),
-                        () => changed = true);
-            Thread.Sleep(30);
-            Assert.False(changed);
+                        () => changed.Set());
+
+            Assert.False(changed.WaitOne(100), "UriLoader ChangeToken fired but shouldn't have.");
         }
     }