Extend RunOnSecondaryUserTargetPreparer to allow concurrent secondary users

This CL introduced a parameter into RunOnSecondaryUserTargetPreparer
to allow to start a secondary user on a secondary display, and run
concurrently with the current user.

Bug:327704045
Test: m && make tradefed-tests && m -j tradefed && tools/tradefederation/core/tradefed.sh run singleCommand host -n --class com.android.tradefed.targetprep.RunOnSecondaryUserTargetPreparerTest

Change-Id: Ie8653733a5b00ad9a48f5d3b7a510b81eda9edce

2 files changed, 203 insertions, 28 deletions

diff --git a/javatests/com/android/tradefed/targetprep/RunOnSecondaryUserTargetPreparerTest.java b/javatests/com/android/tradefed/targetprep/RunOnSecondaryUserTargetPreparerTest.java
index aa1e53460..2f49ae0db 100644
--- a/javatests/com/android/tradefed/targetprep/RunOnSecondaryUserTargetPreparerTest.java
+++ b/javatests/com/android/tradefed/targetprep/RunOnSecondaryUserTargetPreparerTest.java
@@ -46,7 +46,9 @@ import org.mockito.junit.MockitoRule;
 
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 
 @RunWith(JUnit4.class)
 public class RunOnSecondaryUserTargetPreparerTest {
@@ -67,7 +69,7 @@ public class RunOnSecondaryUserTargetPreparerTest {
         ArrayList<Integer> userIds = new ArrayList<>();
         userIds.add(0);
 
-        when(mTestInfo.getDevice().getMaxNumberOfUsersSupported()).thenReturn(2);
+        when(mTestInfo.getDevice().getMaxNumberOfUsersSupported()).thenReturn(3);
         when(mTestInfo.getDevice().listUsers()).thenReturn(userIds);
         when(mTestInfo.getDevice().getApiLevel()).thenReturn(30);
         when(mTestInfo.getDevice().isHeadlessSystemUserMode()).thenReturn(false);
@@ -409,4 +411,143 @@ public class RunOnSecondaryUserTargetPreparerTest {
         verify(mTestInfo.properties(), never())
                 .put(eq(RunOnSecondaryUserTargetPreparer.SKIP_TESTS_REASON_KEY), any());
     }
+
+    @Test
+    public void setUp_hasOneSecondaryUser_createSecondaryUserOnSecondaryDisplay() throws Exception {
+        when(mTestInfo.getDevice().isHeadlessSystemUserMode()).thenReturn(true);
+        when(mTestInfo.getDevice().isVisibleBackgroundUsersSupported()).thenReturn(true);
+        mOptionSetter.setOptionValue(RunOnSecondaryUserTargetPreparer.START_BACKGROUND_USER,
+                "true");
+
+        ArrayList<Integer> userIds = new ArrayList<>();
+        int systemUser = 0;
+        int secondaryUser1 = 100;
+        int secondaryUser2 = 101;
+        userIds.add(systemUser);
+        userIds.add(secondaryUser1);
+        when(mTestInfo.getDevice().listUsers()).thenReturn(userIds);
+        when(mTestInfo.getDevice().getCurrentUser()).thenReturn(secondaryUser1);
+        when(mTestInfo.getDevice().createUser(any(), anyBoolean(), anyBoolean(), anyBoolean()))
+                .thenReturn(secondaryUser2);
+        when(mTestInfo.getDevice().isUserSecondary(secondaryUser1)).thenReturn(true);
+        when(mTestInfo.getDevice().isUserSecondary(secondaryUser2)).thenReturn(true);
+
+        Map<Integer, UserInfo> userInfos = new HashMap<>();
+        userInfos.put(
+                systemUser,
+                new UserInfo(systemUser, "system", /* flag= */ 0, /* isRunning= */ true));
+        userInfos.put(
+                secondaryUser1,
+                new UserInfo(secondaryUser1, "current", /* flag= */ 0, /* isRunning= */ true));
+        when(mTestInfo.getDevice().getUserInfos()).thenReturn(userInfos);
+
+        int secondaryDisplayId = 200;
+        Set<Integer> secondaryDisplayIdSet = new HashSet<>();
+        secondaryDisplayIdSet.add(secondaryDisplayId);
+        when(mTestInfo.getDevice().listDisplayIdsForStartingVisibleBackgroundUsers())
+                .thenReturn(secondaryDisplayIdSet);
+
+        mPreparer.setUp(mTestInfo);
+
+        verify(mTestInfo.getDevice())
+                .createUser(
+                        "secondary",
+                        /* guest= */ false,
+                        /* ephemeral= */ false,
+                        /* forTesting= */ true);
+        verify(mTestInfo.getDevice())
+                .startVisibleBackgroundUser(
+                        secondaryUser2, secondaryDisplayId, /* waitFlag= */ true);
+    }
+
+    @Test
+    public void setUp_hasOneSecondaryUser_doNothing() throws Exception {
+        when(mTestInfo.getDevice().isHeadlessSystemUserMode()).thenReturn(true);
+        when(mTestInfo.getDevice().isVisibleBackgroundUsersSupported()).thenReturn(true);
+        // START_BACKGROUND_USER of this test is false by default.
+
+        ArrayList<Integer> userIds = new ArrayList<>();
+        int systemUser = 0;
+        int secondaryUser1 = 100;
+        int secondaryUser2 = 101;
+        userIds.add(systemUser);
+        userIds.add(secondaryUser1);
+        when(mTestInfo.getDevice().listUsers()).thenReturn(userIds);
+        when(mTestInfo.getDevice().getCurrentUser()).thenReturn(secondaryUser1);
+        when(mTestInfo.getDevice().createUser(any(), anyBoolean(), anyBoolean(), anyBoolean()))
+                .thenReturn(secondaryUser2);
+        when(mTestInfo.getDevice().isUserSecondary(secondaryUser1)).thenReturn(true);
+        when(mTestInfo.getDevice().isUserSecondary(secondaryUser2)).thenReturn(true);
+
+        Map<Integer, UserInfo> userInfos = new HashMap<>();
+        userInfos.put(
+                systemUser,
+                new UserInfo(systemUser, "system", /* flag= */ 0, /* isRunning= */ true));
+        userInfos.put(
+                secondaryUser1,
+                new UserInfo(secondaryUser1, "current", /* flag= */ 0, /* isRunning= */ true));
+        when(mTestInfo.getDevice().getUserInfos()).thenReturn(userInfos);
+
+        int secondaryDisplayId = 200;
+        Set<Integer> secondaryDisplayIdSet = new HashSet<>();
+        secondaryDisplayIdSet.add(secondaryDisplayId);
+        when(mTestInfo.getDevice().listDisplayIdsForStartingVisibleBackgroundUsers())
+                .thenReturn(secondaryDisplayIdSet);
+
+        mPreparer.setUp(mTestInfo);
+
+        verify(mTestInfo.getDevice(), never())
+                .createUser(
+                        "secondary",
+                        /* guest= */ false,
+                        /* ephemeral= */ false,
+                        /* forTesting= */ true);
+        verify(mTestInfo.getDevice(), never())
+                .startVisibleBackgroundUser(
+                        secondaryUser2, secondaryDisplayId, /* waitFlag= */ true);
+    }
+
+    @Test
+    public void setUp_hasTwoSecondaryUsers_startSecondaryUserOnSecondaryDisplay() throws Exception {
+        when(mTestInfo.getDevice().isHeadlessSystemUserMode()).thenReturn(true);
+        when(mTestInfo.getDevice().isVisibleBackgroundUsersSupported()).thenReturn(true);
+        mOptionSetter.setOptionValue(RunOnSecondaryUserTargetPreparer.START_BACKGROUND_USER,
+                "true");
+
+        ArrayList<Integer> userIds = new ArrayList<>();
+        int systemUser = 0;
+        int secondaryUser1 = 100;
+        int secondaryUser2 = 101;
+        userIds.add(systemUser);
+        userIds.add(secondaryUser1);
+        userIds.add(secondaryUser2);
+        when(mTestInfo.getDevice().listUsers()).thenReturn(userIds);
+        when(mTestInfo.getDevice().getCurrentUser()).thenReturn(secondaryUser1);
+        when(mTestInfo.getDevice().isUserSecondary(secondaryUser1)).thenReturn(true);
+        when(mTestInfo.getDevice().isUserSecondary(secondaryUser2)).thenReturn(true);
+
+        Map<Integer, UserInfo> userInfos = new HashMap<>();
+        userInfos.put(
+                systemUser,
+                new UserInfo(systemUser, "system", /* flag= */ 0, /* isRunning= */ true));
+        userInfos.put(
+                secondaryUser1,
+                new UserInfo(secondaryUser1, "current", /* flag= */ 0, /* isRunning= */ true));
+        userInfos.put(
+                secondaryUser2,
+                new UserInfo(secondaryUser2, "secondary", /* flag= */ 0, /* isRunning= */ true));
+        when(mTestInfo.getDevice().getUserInfos()).thenReturn(userInfos);
+
+        int secondaryDisplayId = 200;
+        Set<Integer> secondaryDisplayIdSet = new HashSet<>();
+        secondaryDisplayIdSet.add(secondaryDisplayId);
+        when(mTestInfo.getDevice().listDisplayIdsForStartingVisibleBackgroundUsers())
+                .thenReturn(secondaryDisplayIdSet);
+
+        mPreparer.setUp(mTestInfo);
+
+        verify(mTestInfo.getDevice())
+                .startVisibleBackgroundUser(
+                        secondaryUser2, secondaryDisplayId, /* waitFlag= */ true);
+    }
 }
diff --git a/src/com/android/tradefed/targetprep/RunOnSecondaryUserTargetPreparer.java b/src/com/android/tradefed/targetprep/RunOnSecondaryUserTargetPreparer.java
index 862fde904..dff85e439 100644
--- a/src/com/android/tradefed/targetprep/RunOnSecondaryUserTargetPreparer.java
+++ b/src/com/android/tradefed/targetprep/RunOnSecondaryUserTargetPreparer.java
@@ -32,15 +32,17 @@ import java.util.Collections;
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 /**
- * An {@link ITargetPreparer} that creates a secondary user in setup, and marks that tests should be
- * run in that user.
+ * An {@link ITargetPreparer} to ensure that the test runs as a secondary user. In addition, if
+ * the option {@link START_BACKGROUND_USER} is {@code true} and the current user is already
+ * a secondary user, it will ensure that there is a visble background secondary user run on a
+ * secondary display.
  *
- * <p>In teardown, the secondary user is removed.
- *
- * <p>If a secondary user already exists, it will be used rather than creating a new one, and it
- * will not be removed in teardown.
+ * <p>If the target secondary user doesn't exist, it will create a new one and remove it in
+ * teardown. Otherwise, it will be used rather than creating a new one, and it will not be removed
+ * in teardown.
  *
  * <p>If the device does not have capacity to create a new user when one is required, then the
  * instrumentation argument skip-tests-reason will be set, and the user will not be changed. Tests
@@ -53,6 +55,8 @@ public class RunOnSecondaryUserTargetPreparer extends BaseTargetPreparer {
 
     @VisibleForTesting static final String SKIP_TESTS_REASON_KEY = "skip-tests-reason";
 
+    @VisibleForTesting static final String START_BACKGROUND_USER = "start-background-user";
+
     private int userIdToDelete = -1;
     private int originalUserId;
 
@@ -64,13 +68,26 @@ public class RunOnSecondaryUserTargetPreparer extends BaseTargetPreparer {
             importance = Option.Importance.IF_UNSET)
     private List<String> mTestPackages = new ArrayList<>();
 
+    @Option(
+            name = START_BACKGROUND_USER,
+            description =
+                    "If true and the current user is a secondary user, it will create a "
+                            + "background secondary user (if such user doesn't exist) and "
+                            + "start the background user on secondary display")
+    private boolean mStartBackgroundUser;
+
     @Override
     public void setUp(TestInformation testInfo)
             throws TargetSetupError, DeviceNotAvailableException {
         removeNonForTestingUsers(testInfo.getDevice());
 
+        originalUserId = testInfo.getDevice().getCurrentUser();
         // This must be a for-testing user because we removed the not-for-testing ones
-        int secondaryUserId = getSecondaryUserId(testInfo.getDevice());
+        int secondaryUserId = getTargetSecondaryUserId(testInfo.getDevice());
+
+        if (secondaryUserId == originalUserId) {
+            return;
+        }
 
         if (secondaryUserId == -1) {
             if (!assumeTrue(
@@ -85,33 +102,49 @@ public class RunOnSecondaryUserTargetPreparer extends BaseTargetPreparer {
         }
 
         // The wait flag is only supported on Android 29+
-        testInfo.getDevice()
-                .startUser(
-                        secondaryUserId, /* waitFlag= */ testInfo.getDevice().getApiLevel() >= 29);
+        boolean waitFlag = testInfo.getDevice().getApiLevel() >= 29;
+        if (!testInfo.getDevice().isUserSecondary(originalUserId)) {
 
-        originalUserId = testInfo.getDevice().getCurrentUser();
-
-        if (originalUserId != secondaryUserId) {
+            testInfo.getDevice().startUser(secondaryUserId, waitFlag);
             testInfo.getDevice().switchUser(secondaryUserId);
+            testInfo.properties().put(RUN_TESTS_AS_USER_KEY, Integer.toString(secondaryUserId));
+        } else {
+            Set<Integer> secondaryDisplayIdSet =
+                    testInfo.getDevice().listDisplayIdsForStartingVisibleBackgroundUsers();
+            if (!assumeTrue(
+                    !secondaryDisplayIdSet.isEmpty(),
+                    "This device has no secondary display",
+                    testInfo)) {
+                return;
+            }
+            int secondaryDisplayId = secondaryDisplayIdSet.stream().findFirst().get();
+            testInfo.getDevice()
+                    .startVisibleBackgroundUser(secondaryUserId, secondaryDisplayId, waitFlag);
         }
-
         for (String pkg : mTestPackages) {
             testInfo.getDevice()
                     .executeShellCommand(
                             "pm install-existing --user " + secondaryUserId + " " + pkg);
         }
 
-        testInfo.properties().put(RUN_TESTS_AS_USER_KEY, Integer.toString(secondaryUserId));
+        testInfo.getDevice().executeShellCommand("pm list packages --user all -U");
     }
 
-    /**
-     * Get the id of a secondary user currently on the device. -1 if there is
-     * none.
-     */
-    private static int getSecondaryUserId(ITestDevice device)
-            throws DeviceNotAvailableException {
+    /** Get the id of a target secondary user currently on the device. -1 if there is none. */
+    private int getTargetSecondaryUserId(ITestDevice device) throws DeviceNotAvailableException {
         for (Map.Entry<Integer, UserInfo> userInfo : device.getUserInfos().entrySet()) {
-            if (userInfo.getValue().isSecondary()) {
+            if (!userInfo.getValue().isSecondary()) {
+                continue;
+            }
+            // If mStartBackgroundUser is true and the current user is a secondary user,
+            // we need the target secondary user to be a non-current user (For example, on AAOS
+            // the current user is user 10, if mStartBackgroundUser is true, we need to create user
+            // 11). Otherwise, any secondary user is fine.
+            if (mStartBackgroundUser && device.isUserSecondary(originalUserId)) {
+                if (userInfo.getValue().userId() != originalUserId) {
+                    return userInfo.getKey();
+                }
+            } else {
                 return userInfo.getKey();
             }
         }
@@ -161,22 +194,23 @@ public class RunOnSecondaryUserTargetPreparer extends BaseTargetPreparer {
     /**
      * Remove all non for-testing users.
      *
-     * <p>For a headless device, it would remove every non for-testing user except the first
-     * secondary user and the system user.
+     * <p>For a headless device, if {@code mStartBackgroundUser} is true, it would remove every non
+     * for-testing user except the first two secondary users and the system user; otherwise, it
+     * would remove every non for-testing user except the first secondary user and the system user.
      *
      * <p>For a non-headless device, it would remove every non for-testing user except the system
      * user.
      *
      * <p>A communal profile is never removed.
      */
-    private static void removeNonForTestingUsers(ITestDevice device)
-            throws DeviceNotAvailableException {
+    private void removeNonForTestingUsers(ITestDevice device) throws DeviceNotAvailableException {
         Map<Integer, UserInfo> userInfoMap = device.getUserInfos();
 
         List<UserInfo> userInfos = new ArrayList<>(userInfoMap.values());
         Collections.sort(userInfos, Comparator.comparing(UserInfo::userId));
 
-        int maxSkippedUsers = device.isHeadlessSystemUserMode() ? 1 : 0;
+        int maxSkippedUsers =
+                device.isHeadlessSystemUserMode() ? (mStartBackgroundUser ? 2 : 1) : 0;
         int skippedUsers = 0;
 
         for (UserInfo userInfo : userInfos) {