From 736adde1c70b5c819a1ceb3af8d5fd791cee719a Mon Sep 17 00:00:00 2001
From: vadimt <vadimt@google.com>
Date: Wed, 6 Feb 2019 15:35:09 -0800
Subject: [PATCH] Adding wellbeing toast test

Bug: 130914022
Change-Id: I777440884c122ec040e5e00bfbed41f2eb86b5b0
---
 .../views/DigitalWellBeingToast.java          |  5 +-
 .../com/android/quickstep/views/TaskView.java |  4 +
 .../quickstep/DigitalWellBeingToastTest.java  | 76 +++++++++++++++++++
 .../launcher3/ui/AbstractLauncherUiTest.java  |  2 +-
 4 files changed, 85 insertions(+), 2 deletions(-)
 create mode 100644 quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java

diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/DigitalWellBeingToast.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/DigitalWellBeingToast.java
index b9791bfbc6..204dd56872 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/DigitalWellBeingToast.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/DigitalWellBeingToast.java
@@ -32,7 +32,6 @@ import android.os.UserHandle;
 import android.util.AttributeSet;
 import android.util.Log;
 import android.view.View;
-import android.widget.ImageView;
 import android.widget.LinearLayout;
 import android.widget.TextView;
 
@@ -69,6 +68,10 @@ public final class DigitalWellBeingToast extends LinearLayout {
         mLauncherApps = context.getSystemService(LauncherApps.class);
     }
 
+    public TextView getTextView() {
+        return mText;
+    }
+
     @Override
     protected void onFinishInflate() {
         super.onFinishInflate();
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java
index ec0112d522..05fe2d0d1e 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java
@@ -212,6 +212,10 @@ public class TaskView extends FrameLayout implements PageCallbacks, Reusable {
         return mMenuView;
     }
 
+    public DigitalWellBeingToast getDigitalWellBeingToast() {
+        return mDigitalWellBeingToast;
+    }
+
     /**
      * Updates this task view to the given {@param task}.
      */
diff --git a/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java b/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java
new file mode 100644
index 0000000000..aa113844f6
--- /dev/null
+++ b/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java
@@ -0,0 +1,76 @@
+package com.android.quickstep;
+
+import static androidx.test.InstrumentationRegistry.getInstrumentation;
+
+import static com.android.launcher3.LauncherState.OVERVIEW;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import android.app.PendingIntent;
+import android.app.usage.UsageStatsManager;
+import android.content.Intent;
+
+import com.android.launcher3.Launcher;
+import com.android.quickstep.views.DigitalWellBeingToast;
+import com.android.quickstep.views.RecentsView;
+import com.android.quickstep.views.TaskView;
+
+import org.junit.Test;
+
+import java.time.Duration;
+
+public class DigitalWellBeingToastTest extends AbstractQuickStepTest {
+    private static final String CALCULATOR_PACKAGE =
+            resolveSystemApp(Intent.CATEGORY_APP_CALCULATOR);
+
+    @Test
+    public void testToast() throws Exception {
+        final UsageStatsManager usageStatsManager =
+                mTargetContext.getSystemService(UsageStatsManager.class);
+        final int observerId = 0;
+
+        getInstrumentation().getUiAutomation().adoptShellPermissionIdentity();
+        try {
+            final String[] packages = new String[]{CALCULATOR_PACKAGE};
+
+            // Set time limit for app.
+            usageStatsManager.registerAppUsageLimitObserver(observerId, packages,
+                    Duration.ofSeconds(600), Duration.ofSeconds(300),
+                    PendingIntent.getActivity(mTargetContext, -1, new Intent(), 0));
+
+            mLauncher.pressHome();
+            final DigitalWellBeingToast toast = getToast();
+
+            assertTrue("Toast is not visible", toast.isShown());
+            assertEquals("Toast text: ", "5 minutes left today", toast.getTextView().getText());
+
+            // Unset time limit for app.
+            usageStatsManager.unregisterAppUsageLimitObserver(observerId);
+
+            mLauncher.pressHome();
+            assertFalse("Toast is visible", getToast().isShown());
+        } finally {
+            usageStatsManager.unregisterAppUsageLimitObserver(observerId);
+            getInstrumentation().getUiAutomation().dropShellPermissionIdentity();
+        }
+    }
+
+    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);
+
+        return getFromLauncher(launcher -> {
+            final TaskView task = getLatestTask(launcher);
+            assertTrue("Latest task is not Calculator",
+                    CALCULATOR_PACKAGE.equals(task.getTask().getTopComponent().getPackageName()));
+            return task.getDigitalWellBeingToast();
+        });
+    }
+
+    private TaskView getLatestTask(Launcher launcher) {
+        return launcher.<RecentsView>getOverviewPanel().getTaskViewAt(0);
+    }
+}
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index ab2808d443..71e7880318 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -361,7 +361,7 @@ public abstract class AbstractLauncherUiTest {
                 mDevice.wait(Until.hasObject(By.pkg(packageName).depth(0)), LONG_WAIT_TIME_MS));
     }
 
-    protected String resolveSystemApp(String category) {
+    protected static String resolveSystemApp(String category) {
         return getInstrumentation().getContext().getPackageManager().resolveActivity(
                 new Intent(Intent.ACTION_MAIN).addCategory(category),
                 PackageManager.MATCH_SYSTEM_ONLY).