diff --git a/src/common/UITestAutomation/FindHelper.cs b/src/common/UITestAutomation/FindHelper.cs
index 38b3d3f1f9..db05d66229 100644
--- a/src/common/UITestAutomation/FindHelper.cs
+++ b/src/common/UITestAutomation/FindHelper.cs
@@ -19,19 +19,6 @@ namespace Microsoft.PowerToys.UITest
     {
         public static ReadOnlyCollection<T>? FindAll<T, TW>(Func<IReadOnlyCollection<TW>> findElementsFunc, WindowsDriver<WindowsElement>? driver, int timeoutMS)
             where T : Element, new()
-        {
-            var items = findElementsFunc();
-            var res = items.Select(item =>
-            {
-                var element = item as WindowsElement;
-                return NewElement<T>(element, driver, timeoutMS);
-            }).Where(item => item.IsMatchingTarget()).ToList();
-
-            return new ReadOnlyCollection<T>(res);
-        }
-
-        public static ReadOnlyCollection<T>? FindAll<T, TW>(Func<ReadOnlyCollection<TW>> findElementsFunc, WindowsDriver<WindowsElement>? driver, int timeoutMS)
-            where T : Element, new()
         {
             var items = FindElementsWithRetry(findElementsFunc, timeoutMS);
             var res = items.Select(item =>
@@ -43,21 +30,21 @@ namespace Microsoft.PowerToys.UITest
             return new ReadOnlyCollection<T>(res);
         }
 
-        private static ReadOnlyCollection<TW> FindElementsWithRetry<TW>(Func<ReadOnlyCollection<TW>> findElementsFunc, int timeoutMS = 120000)
+        private static ReadOnlyCollection<TW> FindElementsWithRetry<TW>(Func<IReadOnlyCollection<TW>> findElementsFunc, int timeoutMS)
         {
-            int retryIntervalMS = 500;
-            int elapsedTime = 0;
+            var timeout = TimeSpan.FromMilliseconds(timeoutMS);
+            var retryIntervalMS = TimeSpan.FromMilliseconds(500);
+            DateTime startTime = DateTime.Now;
 
-            while (elapsedTime < timeoutMS)
+            while (DateTime.Now - startTime < timeout)
             {
                 var items = findElementsFunc();
                 if (items.Count > 0)
                 {
-                    return items;
+                    return new ReadOnlyCollection<TW>((IList<TW>)items);
                 }
 
                 Task.Delay(retryIntervalMS).Wait();
-                elapsedTime += retryIntervalMS;
             }
 
             return new ReadOnlyCollection<TW>(new List<TW>());
diff --git a/src/common/UITestAutomation/Session.cs b/src/common/UITestAutomation/Session.cs
index 064c085f22..2ba4b44365 100644
--- a/src/common/UITestAutomation/Session.cs
+++ b/src/common/UITestAutomation/Session.cs
@@ -496,7 +496,7 @@ namespace Microsoft.PowerToys.UITest
                         break;
                     }
 
-                    Thread.Sleep(retryInterval);
+                    Task.Delay(retryInterval).Wait();
                 }
 
                 if (matchingWindows == null || matchingWindows.Count == 0 || matchingWindows[0].HWnd == IntPtr.Zero)
diff --git a/src/modules/Hosts/Hosts.UITests/HostModuleTests.cs b/src/modules/Hosts/Hosts.UITests/HostModuleTests.cs
index beae5dc5c7..70cfe12746 100644
--- a/src/modules/Hosts/Hosts.UITests/HostModuleTests.cs
+++ b/src/modules/Hosts/Hosts.UITests/HostModuleTests.cs
@@ -283,25 +283,25 @@ namespace Hosts.UITests
         private void CloseWarningDialog()
         {
             // Find 'Accept' button which come in 'Warning' dialog
-            if (this.FindAll("Warning").Count > 0 &&
-                this.FindAll<Button>("Accept").Count > 0)
+            if (this.FindAll("Warning", 1000).Count > 0 &&
+                this.FindAll<Button>("Accept", 1000).Count > 0)
             {
                 // Hide Warning dialog if any
-                this.Find<Button>("Accept").Click();
+                this.Find<Button>("Accept", 1000).Click();
             }
         }
 
         private void RemoveAllEntries()
         {
             // Delete all existing host-override rules
-            foreach (var deleteBtn in this.FindAll<Button>("Delete"))
+            foreach (var deleteBtn in this.FindAll<Button>("Delete", 1000))
             {
                 deleteBtn.Click();
-                this.Find<Button>("Yes").Click();
+                this.Find<Button>("Yes", 1000).Click();
             }
 
             // Should have no row left, and no more delete button
-            Assert.IsTrue(this.FindAll<Button>("Delete").Count == 0);
+            Assert.IsTrue(this.FindAll<Button>("Delete", 1000).Count == 0);
         }
     }
 }