More improvements for wellbeing tests

Change-Id: I3c5866660bfb96378cdc80c1afc70ec438fc9add

1 files changed, 14 insertions, 1 deletions

diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 4a9df989c..ab2808d44 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -280,7 +280,7 @@ public abstract class AbstractLauncherUiTest {
     // the results of that gesture because the wait can hide flakeness.
     protected void waitForState(String message, LauncherState state) {
         waitForLauncherCondition(message,
-                launcher -> launcher.getStateManager().getState() == state);
+                launcher -> launcher.getStateManager().getCurrentStableState() == state);
     }
 
     protected void waitForResumed(String message) {
@@ -301,6 +301,19 @@ public abstract class AbstractLauncherUiTest {
         Wait.atMost(message, () -> getFromLauncher(condition), timeout);
     }
 
+    // Cannot be used in TaplTests after injecting any gesture using Tapl because this can hide
+    // flakiness.
+    protected void waitForLauncherCondition(
+            String message,
+            Runnable testThreadAction, Function<Launcher, Boolean> condition,
+            long timeout) {
+        if (!TestHelpers.isInLauncherProcess()) return;
+        Wait.atMost(message, () -> {
+            testThreadAction.run();
+            return getFromLauncher(condition);
+        }, timeout);
+    }
+
     protected LauncherActivityInfo getSettingsApp() {
         return LauncherAppsCompat.getInstance(mTargetContext)
                 .getActivityList("com.android.settings",