From 637363c858412402dfcffd1a24b70037a01d1d82 Mon Sep 17 00:00:00 2001
From: Yasin Kilicdere <tyk@google.com>
Date: Tue, 6 Jun 2023 11:54:09 +0100
Subject: [PATCH] Add assumptions to testRemoveUserWhenPossible_withProfiles.

This test method creates a user and adds profiles to it. Then after
removing the user it asserts that the profiles are also removed. But,
in case the profiles fail to be created, test was failing with a null
pointer exception. This CL adds necessary assumptions to make sure it
is possible to add profiles to the secondary user and skips otherwise
instead of failing.

Bug: 261468057
Test: atest com.android.server.pm.UserManagerTest#testRemoveUserWhenPossible_withProfiles
Change-Id: I3c31b8c66f596192f6ddd1a3b04dac34ac8fe681
---
 .../android/server/pm/UserManagerTest.java    | 39 ++++++++++++++-----
 1 file changed, 30 insertions(+), 9 deletions(-)

diff --git a/services/tests/servicestests/src/com/android/server/pm/UserManagerTest.java b/services/tests/servicestests/src/com/android/server/pm/UserManagerTest.java
index 6bcda3fbcf43d..dc86e5fbe1f3d 100644
--- a/services/tests/servicestests/src/com/android/server/pm/UserManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/pm/UserManagerTest.java
@@ -53,6 +53,8 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -545,23 +547,42 @@ public final class UserManagerTest {
     public void testRemoveUserWhenPossible_withProfiles() throws Exception {
         assumeHeadlessModeEnabled();
         assumeCloneEnabled();
-        final UserInfo parentUser = createUser("Human User", /* flags= */ 0);
-        final UserInfo cloneProfileUser = createProfileForUser("Clone Profile user",
+        final List<String> profileTypesToCreate = Arrays.asList(
                 UserManager.USER_TYPE_PROFILE_CLONE,
-                parentUser.id);
+                UserManager.USER_TYPE_PROFILE_MANAGED
+        );
 
-        final UserInfo workProfileUser = createProfileForUser("Work Profile user",
-                UserManager.USER_TYPE_PROFILE_MANAGED,
-                parentUser.id);
+        final UserInfo parentUser = createUser("Human User", /* flags= */ 0);
+        assertWithMessage("Could not create parent user")
+                .that(parentUser).isNotNull();
+
+        final List<Integer> profileIds = new ArrayList<>();
+        for (String profileType : profileTypesToCreate) {
+            final String name = profileType.substring(profileType.lastIndexOf('.') + 1);
+            if (mUserManager.canAddMoreProfilesToUser(profileType, parentUser.id)) {
+                final UserInfo profile = createProfileForUser(name, profileType, parentUser.id);
+                assertWithMessage("Could not create " + name)
+                        .that(profile).isNotNull();
+                profileIds.add(profile.id);
+            } else {
+                Slog.w(TAG, "Can not add " + name + " to user #" + parentUser.id);
+            }
+        }
+
+        // Test shouldn't pass or fail unless it's allowed to add profiles to secondary users.
+        assumeTrue("Not possible to create any profiles to user #" + parentUser.id,
+                profileIds.size() > 0);
 
         assertThat(mUserManager.removeUserWhenPossible(parentUser.getUserHandle(),
                 /* overrideDevicePolicy= */ false))
                 .isEqualTo(UserManager.REMOVE_RESULT_REMOVED);
         waitForUserRemoval(parentUser.id);
 
-        assertThat(hasUser(parentUser.id)).isFalse();
-        assertThat(hasUser(cloneProfileUser.id)).isFalse();
-        assertThat(hasUser(workProfileUser.id)).isFalse();
+        assertWithMessage("Parent user still exists")
+                .that(hasUser(parentUser.id)).isFalse();
+        profileIds.forEach(id ->
+                assertWithMessage("Profile still exists")
+                        .that(hasUser(id)).isFalse());
     }
 
     /** Tests creating a FULL user via specifying userType. */