1 files changed, 3 insertions, 3 deletions
diff --git a/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java b/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java
index 70f9c90da..0c5a6f5b6 100644
--- a/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java
+++ b/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java
@@ -51,15 +51,15 @@ public class DigitalWellBeingToastTest extends AbstractQuickStepTest {
             mLauncher.pressHome();
             final DigitalWellBeingToast toast = getToast();
 
-            assertTrue("Toast is not visible", toast.isShown());
-            assertEquals("Toast text: ", "5 minutes left today", toast.getTextView().getText());
+            assertTrue("Toast is not visible", toast.hasLimit());
+            assertEquals("Toast text: ", "5 minutes left today", toast.getText());
 
             // Unset time limit for app.
             runWithShellPermission(
                     () -> usageStatsManager.unregisterAppUsageLimitObserver(observerId));
 
             mLauncher.pressHome();
-            assertFalse("Toast is visible", getToast().isShown());
+            assertFalse("Toast is visible", getToast().hasLimit());
         } finally {
             runWithShellPermission(
                     () -> usageStatsManager.unregisterAppUsageLimitObserver(observerId));