diff --git a/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java b/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java
index a7c33a9..3e84a76 100644
--- a/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java
+++ b/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java
@@ -68,7 +68,7 @@
 
     private DigitalWellBeingToast getToast() {
         executeOnLauncher(launcher -> launcher.getStateManager().goToState(OVERVIEW));
-        waitForState("Launcher internal state didn't switch to Overview", OVERVIEW);
+        waitForState("Launcher internal state didn't switch to Overview", () -> OVERVIEW);
         final TaskView task = getOnceNotNull("No latest task", launcher -> getLatestTask(launcher));
 
         return getFromLauncher(launcher -> {
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
index 71d77fc..2804c0e 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
@@ -82,7 +82,8 @@
     @Ignore // Enable after b/131115533
     public void testPressRecentAppsLauncherAndGetOverview() throws RemoteException {
         mDevice.pressRecentApps();
-        waitForState("Launcher internal state didn't switch to Overview", LauncherState.OVERVIEW);
+        waitForState("Launcher internal state didn't switch to Overview",
+                () -> LauncherState.OVERVIEW);
 
         assertNotNull("getOverview() returned null", mLauncher.getOverview());
     }
@@ -93,7 +94,8 @@
     public void testWorkspaceSwitchToAllApps() {
         assertNotNull("switchToAllApps() returned null",
                 mLauncher.getWorkspace().switchToAllApps());
-        assertTrue("Launcher internal state is not All Apps", isInState(LauncherState.ALL_APPS));
+        assertTrue("Launcher internal state is not All Apps",
+                isInState(() -> LauncherState.ALL_APPS));
     }
 
     @Test
@@ -112,7 +114,7 @@
         // mLauncher.pressHome() also tests an important case of pressing home while in background.
         Overview overview = mLauncher.pressHome().switchToOverview();
         assertTrue("Launcher internal state didn't switch to Overview",
-                isInState(LauncherState.OVERVIEW));
+                isInState(() -> LauncherState.OVERVIEW));
         executeOnLauncher(
                 launcher -> assertTrue("Don't have at least 3 tasks", getTaskCount(launcher) >= 3));
 
@@ -121,14 +123,16 @@
                 0, getCurrentOverviewPage(launcher)));
 
         overview.flingForward();
-        assertTrue("Launcher internal state is not Overview", isInState(LauncherState.OVERVIEW));
+        assertTrue("Launcher internal state is not Overview",
+                isInState(() -> LauncherState.OVERVIEW));
         final Integer currentTaskAfterFlingForward = getFromLauncher(
                 launcher -> getCurrentOverviewPage(launcher));
         executeOnLauncher(launcher -> assertTrue("Current task in Overview is still 0",
                 currentTaskAfterFlingForward > 0));
 
         overview.flingBackward();
-        assertTrue("Launcher internal state is not Overview", isInState(LauncherState.OVERVIEW));
+        assertTrue("Launcher internal state is not Overview",
+                isInState(() -> LauncherState.OVERVIEW));
         executeOnLauncher(launcher -> assertTrue("Flinging back in Overview did nothing",
                 getCurrentOverviewPage(launcher) < currentTaskAfterFlingForward));
 
@@ -147,7 +151,7 @@
         // Test dismissing a task.
         overview = mLauncher.pressHome().switchToOverview();
         assertTrue("Launcher internal state didn't switch to Overview",
-                isInState(LauncherState.OVERVIEW));
+                isInState(() -> LauncherState.OVERVIEW));
         final Integer numTasks = getFromLauncher(launcher -> getTaskCount(launcher));
         task = overview.getCurrentTask();
         assertNotNull("overview.getCurrentTask() returned null (2)", task);
@@ -162,29 +166,29 @@
             final AllAppsFromOverview allApps = overview.switchToAllApps();
             assertNotNull("overview.switchToAllApps() returned null (1)", allApps);
             assertTrue("Launcher internal state is not All Apps (1)",
-                    isInState(LauncherState.ALL_APPS));
+                    isInState(() -> LauncherState.ALL_APPS));
 
             overview = allApps.switchBackToOverview();
             assertNotNull("allApps.switchBackToOverview() returned null", overview);
             assertTrue("Launcher internal state didn't switch to Overview",
-                    isInState(LauncherState.OVERVIEW));
+                    isInState(() -> LauncherState.OVERVIEW));
 
             // Test UIDevice.pressBack()
             overview.switchToAllApps();
             assertNotNull("overview.switchToAllApps() returned null (2)", allApps);
             assertTrue("Launcher internal state is not All Apps (2)",
-                    isInState(LauncherState.ALL_APPS));
+                    isInState(() -> LauncherState.ALL_APPS));
             mDevice.pressBack();
             mLauncher.getOverview();
         }
 
         // Test UIDevice.pressHome, once we are in AllApps.
         mDevice.pressHome();
-        waitForState("Launcher internal state didn't switch to Home", LauncherState.NORMAL);
+        waitForState("Launcher internal state didn't switch to Home", () -> LauncherState.NORMAL);
 
         // Test dismissing all tasks.
         mLauncher.getWorkspace().switchToOverview().dismissAllTasks();
-        waitForState("Launcher internal state didn't switch to Home", LauncherState.NORMAL);
+        waitForState("Launcher internal state didn't switch to Home", () -> LauncherState.NORMAL);
         executeOnLauncher(
                 launcher -> assertEquals("Still have tasks after dismissing all",
                         0, getTaskCount(launcher)));
@@ -202,7 +206,8 @@
     public void testAppIconLaunchFromAllAppsFromOverview() throws Exception {
         final AllApps allApps =
                 mLauncher.getWorkspace().switchToOverview().switchToAllApps();
-        assertTrue("Launcher internal state is not All Apps", isInState(LauncherState.ALL_APPS));
+        assertTrue("Launcher internal state is not All Apps",
+                isInState(() -> LauncherState.ALL_APPS));
 
         TaplTestsLauncher3.runIconLaunchFromAllAppsTest(this, allApps);
     }
@@ -214,7 +219,7 @@
         assertNotNull("Workspace.switchToOverview() returned null",
                 mLauncher.pressHome().switchToOverview());
         assertTrue("Launcher internal state didn't switch to Overview",
-                isInState(LauncherState.OVERVIEW));
+                isInState(() -> LauncherState.OVERVIEW));
     }
 
     @Test
@@ -226,7 +231,7 @@
 
         assertNotNull("Background.switchToOverview() returned null", background.switchToOverview());
         assertTrue("Launcher internal state didn't switch to Overview",
-                isInState(LauncherState.OVERVIEW));
+                isInState(() -> LauncherState.OVERVIEW));
     }
 
     private Background getAndAssertBackground() {
@@ -249,9 +254,11 @@
         TaplTestsLauncher3.runAllAppsTest(this, mLauncher.getAllApps());
 
         // Testing pressHome.
-        assertTrue("Launcher internal state is not All Apps", isInState(LauncherState.ALL_APPS));
+        assertTrue("Launcher internal state is not All Apps",
+                isInState(() -> LauncherState.ALL_APPS));
         assertNotNull("pressHome returned null", mLauncher.pressHome());
-        assertTrue("Launcher internal state is not Home", isInState(LauncherState.NORMAL));
+        assertTrue("Launcher internal state is not Home",
+                isInState(() -> LauncherState.NORMAL));
         assertNotNull("getHome returned null", mLauncher.getWorkspace());
     }
 
diff --git a/tests/src/com/android/launcher3/compat/PromiseIconUiTest.java b/tests/src/com/android/launcher3/compat/PromiseIconUiTest.java
index efbd9c9..33066e4 100644
--- a/tests/src/com/android/launcher3/compat/PromiseIconUiTest.java
+++ b/tests/src/com/android/launcher3/compat/PromiseIconUiTest.java
@@ -49,7 +49,7 @@
         super.setUp();
         mDevice.pressHome();
         waitForLauncherCondition("Launcher didn't start", launcher -> launcher != null);
-        waitForState("Launcher internal state didn't switch to Home", LauncherState.NORMAL);
+        waitForState("Launcher internal state didn't switch to Home", () -> LauncherState.NORMAL);
         mSessionId = -1;
     }
 
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 19997eb..25670e7 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -84,6 +84,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 /**
  * Base class for all instrumentation tests providing various utility methods.
@@ -281,9 +282,9 @@
 
     // Cannot be used in TaplTests between a Tapl call injecting a gesture and a tapl call expecting
     // the results of that gesture because the wait can hide flakeness.
-    protected void waitForState(String message, LauncherState state) {
+    protected void waitForState(String message, Supplier<LauncherState> state) {
         waitForLauncherCondition(message,
-                launcher -> launcher.getStateManager().getCurrentStableState() == state);
+                launcher -> launcher.getStateManager().getCurrentStableState() == state.get());
     }
 
     protected void waitForResumed(String message) {
@@ -430,9 +431,9 @@
         return !launcher.hasBeenResumed();
     }
 
-    protected boolean isInState(LauncherState state) {
+    protected boolean isInState(Supplier<LauncherState> state) {
         if (!TestHelpers.isInLauncherProcess()) return true;
-        return getFromLauncher(launcher -> launcher.getStateManager().getState() == state);
+        return getFromLauncher(launcher -> launcher.getStateManager().getState() == state.get());
     }
 
     protected int getAllAppsScroll(Launcher launcher) {
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
index 61f5150..2c1883b 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -73,7 +73,8 @@
         test.clearLauncherData();
         test.mDevice.pressHome();
         test.waitForLauncherCondition("Launcher didn't start", launcher -> launcher != null);
-        test.waitForState("Launcher internal state didn't switch to Home", LauncherState.NORMAL);
+        test.waitForState("Launcher internal state didn't switch to Home",
+                () -> LauncherState.NORMAL);
         test.waitForResumed("Launcher internal state is still Background");
         // Check that we switched to home.
         test.mLauncher.getWorkspace();
@@ -143,7 +144,7 @@
 
             assertTrue(
                     "Launcher internal state is not All Apps",
-                    test.isInState(LauncherState.ALL_APPS));
+                    test.isInState(() -> LauncherState.ALL_APPS));
 
             // Test flinging forward and backward.
             test.executeOnLauncher(launcher -> assertEquals(
@@ -152,7 +153,7 @@
 
             allApps.flingForward();
             assertTrue("Launcher internal state is not All Apps",
-                    test.isInState(LauncherState.ALL_APPS));
+                    test.isInState(() -> LauncherState.ALL_APPS));
             final Integer flingForwardY = test.getFromLauncher(
                     launcher -> test.getAllAppsScroll(launcher));
             test.executeOnLauncher(
@@ -162,7 +163,7 @@
             allApps.flingBackward();
             assertTrue(
                     "Launcher internal state is not All Apps",
-                    test.isInState(LauncherState.ALL_APPS));
+                    test.isInState(() -> LauncherState.ALL_APPS));
             final Integer flingBackwardY = test.getFromLauncher(
                     launcher -> test.getAllAppsScroll(launcher));
             test.executeOnLauncher(launcher -> assertTrue("flingBackward() didn't scroll App Apps",
@@ -179,7 +180,7 @@
 
             assertTrue(
                     "Launcher internal state is not All Apps",
-                    test.isInState(LauncherState.ALL_APPS));
+                    test.isInState(() -> LauncherState.ALL_APPS));
         } finally {
             allApps.unfreeze();
         }
@@ -190,7 +191,8 @@
     public void testWorkspaceSwitchToAllApps() {
         assertNotNull("switchToAllApps() returned null",
                 mLauncher.getWorkspace().switchToAllApps());
-        assertTrue("Launcher internal state is not All Apps", isInState(LauncherState.ALL_APPS));
+        assertTrue("Launcher internal state is not All Apps",
+                isInState(() -> LauncherState.ALL_APPS));
     }
 
     @Test
@@ -216,7 +218,7 @@
 
         // Test flinging workspace.
         workspace.flingBackward();
-        assertTrue("Launcher internal state is not Home", isInState(LauncherState.NORMAL));
+        assertTrue("Launcher internal state is not Home", isInState(() -> LauncherState.NORMAL));
         executeOnLauncher(
                 launcher -> assertEquals("Flinging back didn't switch workspace to page #0",
                         0, getCurrentWorkspacePage(launcher)));
@@ -225,7 +227,7 @@
         executeOnLauncher(
                 launcher -> assertEquals("Flinging forward didn't switch workspace to page #1",
                         1, getCurrentWorkspacePage(launcher)));
-        assertTrue("Launcher internal state is not Home", isInState(LauncherState.NORMAL));
+        assertTrue("Launcher internal state is not Home", isInState(() -> LauncherState.NORMAL));
 
         // Test starting a workspace app.
         final AppIcon app = workspace.getWorkspaceAppIcon("Chrome");
@@ -251,7 +253,8 @@
     @PortraitLandscape
     public void testAppIconLaunchFromAllAppsFromHome() throws Exception {
         final AllApps allApps = mLauncher.getWorkspace().switchToAllApps();
-        assertTrue("Launcher internal state is not All Apps", isInState(LauncherState.ALL_APPS));
+        assertTrue("Launcher internal state is not All Apps",
+                isInState(() -> LauncherState.ALL_APPS));
 
         runIconLaunchFromAllAppsTest(this, allApps);
     }