Merge "Adding a test method to get a Launcher object when it becomes non-null" into ub-launcher3-qt-qpr1-dev am: b34f951e7e

am: c698c0cdbb

Change-Id: I0b62ddc7a7bc9fb31020a87750d3093d5d83ebb7

quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java

@@ -69,10 +69,9 @@ public class DigitalWellBeingToastTest extends AbstractQuickStepTest {
     private DigitalWellBeingToast getToast() {
         executeOnLauncher(launcher -> launcher.getStateManager().goToState(OVERVIEW));
         waitForState("Launcher internal state didn't switch to Overview", OVERVIEW);
-        waitForLauncherCondition("No latest task", launcher -> getLatestTask(launcher) != null);
+        final TaskView task = getOnceNotNull("No latest task", launcher -> getLatestTask(launcher));
 
         return getFromLauncher(launcher -> {
-            final TaskView task = getLatestTask(launcher);
             assertTrue("Latest task is not Calculator",
                     CALCULATOR_PACKAGE.equals(task.getTask().getTopComponent().getPackageName()));
             return task.getDigitalWellBeingToast();

tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java

@@ -267,6 +267,12 @@ public abstract class AbstractLauncherUiTest {
         waitForLauncherCondition(message, condition, DEFAULT_ACTIVITY_TIMEOUT);
     }
 
+    // Cannot be used in TaplTests after injecting any gesture using Tapl because this can hide
+    // flakiness.
+    protected <T> T getOnceNotNull(String message, Function<Launcher, T> f) {
+        return getOnceNotNull(message, f, DEFAULT_ACTIVITY_TIMEOUT);
+    }
+
     // Cannot be used in TaplTests after injecting any gesture using Tapl because this can hide
     // flakiness.
     protected void waitForLauncherCondition(
@@ -275,6 +281,20 @@ 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 <T> T getOnceNotNull(String message, Function<Launcher, T> f, long timeout) {
+        if (!TestHelpers.isInLauncherProcess()) return null;
+
+        final Object[] output = new Object[1];
+        Wait.atMost(message, () -> {
+            final Object fromLauncher = getFromLauncher(f);
+            output[0] = fromLauncher;
+            return fromLauncher != null;
+        }, timeout);
+        return (T) output[0];
+    }
+
     // Cannot be used in TaplTests after injecting any gesture using Tapl because this can hide
     // flakiness.
     protected void waitForLauncherCondition(