Merge "Updated UserUtils" into rvc-dev

2 files changed, 59 insertions, 29 deletions

diff --git a/libraries/sts-common-util/host-side/src/com/android/sts/common/UserUtils.java b/libraries/sts-common-util/host-side/src/com/android/sts/common/UserUtils.java
index 01d7c964f..7f8d8bc86 100644
--- a/libraries/sts-common-util/host-side/src/com/android/sts/common/UserUtils.java
+++ b/libraries/sts-common-util/host-side/src/com/android/sts/common/UserUtils.java
@@ -20,6 +20,9 @@ import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.util.CommandResult;
 import com.android.tradefed.util.CommandStatus;
 
+import java.util.HashMap;
+import java.util.Map;
+
 /** Util to manage secondary user */
 public class UserUtils {
 
@@ -34,28 +37,32 @@ public class UserUtils {
         private boolean mIsPreCreateOnly; // User type : --pre-created-only
         private boolean mIsRestricted; // User type : --restricted
         private boolean mSwitch; // Switch to newly created user
-        private boolean
-                mDisallowAppInstall; // Disallow app installation in secondary user explicitly
         private int mProfileOf; // Userid associated with managed user
         private int mTestUserId;
+        private Map<String, String> mUserRestrictions; // Map of user-restrictions for new user
 
         /**
          * Create an instance of secondary user.
          *
          * @param device the device {@link ITestDevice} to use.
-         * @throws IllegalArgumentException when {@code device} is null.
+         * @throws Exception
          */
-        public SecondaryUser(ITestDevice device) throws IllegalArgumentException {
+        public SecondaryUser(ITestDevice device) throws Exception {
             // Device should not be null
             if (device == null) {
                 throw new IllegalArgumentException("Device should not be null");
             }
 
+            // Check if device supports multiple users
+            if (!device.isMultiUserSupported()) {
+                throw new IllegalStateException("Device does not support multiple users");
+            }
+
             mDevice = device;
             mName = "testUser"; /* Default username */
+            mUserRestrictions = new HashMap<String, String>();
 
             // Set default value for all flags as false
-            mDisallowAppInstall = false;
             mIsDemo = false;
             mIsEphemeral = false;
             mIsForTesting = false;
@@ -67,17 +74,6 @@ public class UserUtils {
         }
 
         /**
-         * Disallow app installation in secondary user explicitly. This requires root UID, system
-         * UID, or MANAGE_USERS permission.
-         *
-         * @return this object for method chaining.
-         */
-        public SecondaryUser disallowAppInstallation() {
-            mDisallowAppInstall = true;
-            return this;
-        }
-
-        /**
          * Set the user type as demo.
          *
          * @return this object for method chaining.
@@ -177,6 +173,17 @@ public class UserUtils {
         }
 
         /**
+         * Set user-restrictions on newly created secondary user.
+         * Note: Setting user-restrictions requires enabling root.
+         *
+         * @return this object for method chaining.
+         */
+        public SecondaryUser withUserRestrictions(Map<String, String> restrictions) {
+            mUserRestrictions.putAll(restrictions);
+            return this;
+        }
+
+        /**
          * Create a secondary user and if required, switch to it. Returns an Autocloseable that
          * removes the secondary user.
          *
@@ -233,20 +240,26 @@ public class UserUtils {
                         String.format("Failed to start the user: %s", mTestUserId));
             }
 
-            if (mDisallowAppInstall) {
-                // Enable/Disable app installation in secondary user
-                final CommandResult userRestrictionCmdOutput =
-                        mDevice.executeShellV2Command(
+            // Add user-restrictions to newly created secondary user
+            if (!mUserRestrictions.isEmpty()) {
+                if (!mDevice.isAdbRoot()) {
+                    throw new IllegalStateException("Setting user-restriction requires root");
+                }
+
+                for (Map.Entry<String, String> entry : mUserRestrictions.entrySet()) {
+                    final CommandResult cmdOutput =
+                            mDevice.executeShellV2Command(
+                                    String.format(
+                                            "pm set-user-restriction --user %d %s %s",
+                                            mTestUserId, entry.getKey(), entry.getValue()));
+                    if (cmdOutput.getStatus() != CommandStatus.SUCCESS) {
+                        asSecondaryUser.close();
+                        throw new IllegalStateException(
                                 String.format(
-                                        "pm set-user-restriction --user %d no_install_apps %d",
-                                        mTestUserId, 1));
-                if (userRestrictionCmdOutput.getStatus() != CommandStatus.SUCCESS) {
-                    asSecondaryUser.close();
-                    throw new IllegalStateException(
-                            String.format(
-                                    "Failed to set user restriction 'no_install_apps' with message:"
-                                            + " %s",
-                                    userRestrictionCmdOutput.toString()));
+                                        "Failed to set user restriction %s value %s with"
+                                                + " message %s",
+                                        entry.getKey(), entry.getValue(), cmdOutput.toString()));
+                    }
                 }
             }
 
diff --git a/libraries/sts-common-util/host-side/tests/src/com/android/sts/common/UserUtilsTest.java b/libraries/sts-common-util/host-side/tests/src/com/android/sts/common/UserUtilsTest.java
index 5da5d301b..aa58f1022 100644
--- a/libraries/sts-common-util/host-side/tests/src/com/android/sts/common/UserUtilsTest.java
+++ b/libraries/sts-common-util/host-side/tests/src/com/android/sts/common/UserUtilsTest.java
@@ -29,6 +29,8 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.util.Map;
+
 /** Unit tests for {@link UserUtils}. */
 @RunWith(DeviceJUnit4ClassRunner.class)
 public class UserUtilsTest extends BaseHostJUnit4Test {
@@ -79,4 +81,19 @@ public class UserUtilsTest extends BaseHostJUnit4Test {
                 "device should still be root after cleanup if started with root",
                 getDevice().isAdbRoot());
     }
+
+    @Test
+    public void testUserUtilsUserRestriction() throws Exception {
+        assertTrue("must test with rootable device", getDevice().enableAdbRoot());
+        try (AutoCloseable user =
+                new UserUtils.SecondaryUser(getDevice())
+                    .name(TEST_USER_NAME)
+                    .withUserRestrictions(Map.of("test_restriction", "1"))
+                    .withUser()) {
+            // Exception is thrown if any error occurs while setting user restriction above
+        }
+        assertTrue(
+                "device should still be root after cleanup if started with root",
+                getDevice().isAdbRoot());
+    }
 }