diff --git a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
index 5606ac28e7..b786c8b2fc 100644
--- a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
+++ b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
@@ -35,6 +35,7 @@ import androidx.test.uiautomator.UiDevice;
 
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.tapl.TestHelpers;
+import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.rule.FailureWatcher;
 import com.android.systemui.shared.system.QuickStepContract;
 
@@ -57,6 +58,8 @@ public class NavigationModeSwitchRule implements TestRule {
 
     static final String TAG = "QuickStepOnOffRule";
 
+    public static final int WAIT_TIME_MS = 10000;
+
     public enum Mode {
         THREE_BUTTON, TWO_BUTTON, ZERO_BUTTON, ALL
     }
@@ -118,8 +121,8 @@ public class NavigationModeSwitchRule implements TestRule {
                         if (mode == THREE_BUTTON || mode == ALL) {
                             evaluateWithThreeButtons();
                         }
-                    } catch (Exception e) {
-                        Log.e(TAG, "Exception", e);
+                    } catch (Throwable e) {
+                        Log.e(TAG, "Error", e);
                         throw e;
                     } finally {
                         assertTrue("Couldn't set overlay",
@@ -195,19 +198,14 @@ public class NavigationModeSwitchRule implements TestRule {
                                 currentSysUiNavigationMode() == expectedMode);
                     }
 
-                    for (int i = 0; i != 100; ++i) {
-                        if (mLauncher.getNavigationModel() == expectedMode) break;
-                        Thread.sleep(100);
-                    }
-                    assertTrue("Couldn't switch to " + overlayPackage,
-                            mLauncher.getNavigationModel() == expectedMode);
+                    Wait.atMost("Couldn't switch to " + overlayPackage,
+                            () -> mLauncher.getNavigationModel() == expectedMode, WAIT_TIME_MS,
+                            mLauncher);
 
-                    for (int i = 0; i != 100; ++i) {
-                        if (mLauncher.getNavigationModeMismatchError() == null) break;
-                        Thread.sleep(100);
-                    }
-                    final String error = mLauncher.getNavigationModeMismatchError();
-                    assertTrue("Switching nav mode: " + error, error == null);
+                    Wait.atMost(() -> "Switching nav mode: "
+                                    + mLauncher.getNavigationModeMismatchError(),
+                            () -> mLauncher.getNavigationModeMismatchError() == null, WAIT_TIME_MS,
+                            mLauncher);
 
                     return true;
                 }
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 4243ed041f..056ab9fcbd 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -259,7 +259,7 @@ public abstract class AbstractLauncherUiTest {
     protected <T> T getOnUiThread(final Callable<T> callback) {
         try {
             return mMainThreadExecutor.submit(callback).get();
-        } catch (Exception e) {
+        } catch (Throwable e) {
             throw new RuntimeException(e);
         }
     }
diff --git a/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java b/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java
index 80bb3edddd..1a68122217 100644
--- a/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java
+++ b/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java
@@ -38,8 +38,8 @@ class PortraitLandscapeRunner implements TestRule {
 
                     evaluateInPortrait();
                     evaluateInLandscape();
-                } catch (Exception e) {
-                    Log.e(TAG, "Exception", e);
+                } catch (Throwable e) {
+                    Log.e(TAG, "Error", e);
                     throw e;
                 } finally {
                     mTest.mDevice.setOrientationNatural();
diff --git a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
index 0472ce1c6f..62e2a530ff 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
@@ -103,11 +103,11 @@ public class AddConfigWidgetTest extends AbstractLauncherUiTest {
 
         setResult(acceptConfig);
         if (acceptConfig) {
-            Wait.atMost(null, new WidgetSearchCondition(), DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
+            Wait.atMost("", new WidgetSearchCondition(), DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
             assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
         } else {
             // Verify that the widget id is deleted.
-            Wait.atMost(null, () -> mAppWidgetManager.getAppWidgetInfo(mWidgetId) == null,
+            Wait.atMost("", () -> mAppWidgetManager.getAppWidgetInfo(mWidgetId) == null,
                     DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
         }
     }
diff --git a/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java b/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java
index d909ad7158..59b861cfb3 100644
--- a/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java
@@ -170,7 +170,7 @@ public class RequestPinItemTest extends AbstractLauncherUiTest {
 
         // Go back to home
         mLauncher.pressHome();
-        Wait.atMost(null, new ItemSearchCondition(itemMatcher), DEFAULT_ACTIVITY_TIMEOUT,
+        Wait.atMost("", new ItemSearchCondition(itemMatcher), DEFAULT_ACTIVITY_TIMEOUT,
                 mLauncher);
     }
 
diff --git a/tests/src/com/android/launcher3/util/Wait.java b/tests/src/com/android/launcher3/util/Wait.java
index 2663d02760..2ab1e00252 100644
--- a/tests/src/com/android/launcher3/util/Wait.java
+++ b/tests/src/com/android/launcher3/util/Wait.java
@@ -7,6 +7,8 @@ import com.android.launcher3.tapl.LauncherInstrumentation;
 
 import org.junit.Assert;
 
+import java.util.function.Supplier;
+
 /**
  * A utility class for waiting for a condition to be true.
  */
@@ -16,10 +18,16 @@ public class Wait {
 
     public static void atMost(String message, Condition condition, long timeout,
             LauncherInstrumentation launcher) {
+        atMost(() -> message, condition, timeout, DEFAULT_SLEEP_MS, launcher);
+    }
+
+    public static void atMost(Supplier<String> message, Condition condition, long timeout,
+            LauncherInstrumentation launcher) {
         atMost(message, condition, timeout, DEFAULT_SLEEP_MS, launcher);
     }
 
-    public static void atMost(String message, Condition condition, long timeout, long sleepMillis,
+    public static void atMost(Supplier<String> message, Condition condition, long timeout,
+            long sleepMillis,
             LauncherInstrumentation launcher) {
         final long startTime = SystemClock.uptimeMillis();
         long endTime = startTime + timeout;
@@ -45,6 +53,6 @@ public class Wait {
         }
         Log.d("Wait", "atMost: timed out: " + SystemClock.uptimeMillis());
         launcher.checkForAnomaly();
-        Assert.fail(message);
+        Assert.fail(message.get());
     }
 }