diff --git a/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java b/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java
index ec3d49a..a7c33a9 100644
--- a/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java
+++ b/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java
@@ -69,10 +69,9 @@
     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();
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index d0b2a0a..bee3a49 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -24,7 +24,6 @@
 
 import static java.lang.System.exit;
 
-import android.app.Instrumentation;
 import android.content.BroadcastReceiver;
 import android.content.ComponentName;
 import android.content.Context;
@@ -277,6 +276,12 @@
 
     // 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(
             String message, Function<Launcher, Boolean> condition, long timeout) {
         if (!TestHelpers.isInLauncherProcess()) return;
@@ -285,6 +290,20 @@
 
     // 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(
             String message,
             Runnable testThreadAction, Function<Launcher, Boolean> condition,