More improvements for wellbeing tests

Change-Id: I3c5866660bfb96378cdc80c1afc70ec438fc9add

src/com/android/launcher3/LauncherStateManager.java

@@ -140,6 +140,10 @@ public class LauncherStateManager {
         return mState;
     }
 
+    public LauncherState getCurrentStableState() {
+        return mCurrentStableState;
+    }
+
     public StateHandler[] getStateHandlers() {
         if (mStateHandlers == null) {
             mStateHandlers = UiFactory.getStateHandler(mLauncher);

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",