commit b647866cd033263c80e7a790ad20a865daa649bc
author vadimt <vadimt@google.com> Mon Apr 22 11:32:43 2019 -0700
committer vadimt <vadimt@google.com> Mon Apr 22 11:32:43 2019 -0700
tree d090e4c5c0d6ea4f74a9e81b0f56e1f388cee02f
parent 8b0cb4113fc34ced9a9de71e40f448615ec8210d

Adopt shell permission only to set and unset time limit.

We don't want this permission to "help" Launcher with calling other
APIs.

Bug: 130914022
Change-Id: I649dede958aa6e4a67ccf332151a15cad53b0e9f

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

diff --git a/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java b/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java
index aa11384..8798157 100644
--- a/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java
+++ b/quickstep/tests/src/com/android/quickstep/DigitalWellBeingToastTest.java
@@ -31,14 +31,14 @@
                 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));
+            runWithShellPermission(() ->
+                    usageStatsManager.registerAppUsageLimitObserver(observerId, packages,
+                            Duration.ofSeconds(600), Duration.ofSeconds(300),
+                            PendingIntent.getActivity(mTargetContext, -1, new Intent(), 0)));
 
             mLauncher.pressHome();
             final DigitalWellBeingToast toast = getToast();
@@ -47,13 +47,14 @@
             assertEquals("Toast text: ", "5 minutes left today", toast.getTextView().getText());
 
             // Unset time limit for app.
-            usageStatsManager.unregisterAppUsageLimitObserver(observerId);
+            runWithShellPermission(
+                    () -> usageStatsManager.unregisterAppUsageLimitObserver(observerId));
 
             mLauncher.pressHome();
             assertFalse("Toast is visible", getToast().isShown());
         } finally {
-            usageStatsManager.unregisterAppUsageLimitObserver(observerId);
-            getInstrumentation().getUiAutomation().dropShellPermissionIdentity();
+            runWithShellPermission(
+                    () -> usageStatsManager.unregisterAppUsageLimitObserver(observerId));
         }
     }
 
@@ -73,4 +74,14 @@
     private TaskView getLatestTask(Launcher launcher) {
         return launcher.<RecentsView>getOverviewPanel().getTaskViewAt(0);
     }
+
+    private void runWithShellPermission(Runnable action) {
+        getInstrumentation().getUiAutomation().adoptShellPermissionIdentity();
+        try {
+            action.run();
+        } finally {
+            getInstrumentation().getUiAutomation().dropShellPermissionIdentity();
+        }
+
+    }
 }