Allow device owner type to be set more than once

It is only allowed if the device owner is a test only package.

Bug: 216852998
Test: atest com.android.server.devicepolicy.DevicePolicyManagerTest
Test: Manual test (called the API twice in ManagedProvisioning)

Change-Id: I8851f613fd4e01ec6eff1c95681dd2b474b6c4ad
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 3bda7bf0..7be5d34 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -18271,12 +18271,16 @@
     private void setDeviceOwnerTypeLocked(ComponentName admin,
             @DeviceOwnerType int deviceOwnerType) {
         String packageName = admin.getPackageName();
+        boolean isAdminTestOnly;
 
         verifyDeviceOwnerTypePreconditionsLocked(admin);
-        Preconditions.checkState(!mOwners.isDeviceOwnerTypeSetForDeviceOwner(packageName),
-                "The device owner type has already been set for " + packageName);
 
-        mOwners.setDeviceOwnerType(packageName, deviceOwnerType);
+        isAdminTestOnly = isAdminTestOnlyLocked(admin, mOwners.getDeviceOwnerUserId());
+        Preconditions.checkState(isAdminTestOnly
+                        || !mOwners.isDeviceOwnerTypeSetForDeviceOwner(packageName),
+                "Test only admins can only set the device owner type more than once");
+
+        mOwners.setDeviceOwnerType(packageName, deviceOwnerType, isAdminTestOnly);
     }
 
     @Override
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/Owners.java b/services/devicepolicy/java/com/android/server/devicepolicy/Owners.java
index 3584728..fe8f223 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/Owners.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/Owners.java
@@ -637,13 +637,15 @@
         }
     }
 
-    void setDeviceOwnerType(String packageName, @DeviceOwnerType int deviceOwnerType) {
+    void setDeviceOwnerType(String packageName, @DeviceOwnerType int deviceOwnerType,
+            boolean isAdminTestOnly) {
         synchronized (mLock) {
             if (!hasDeviceOwner()) {
                 Slog.e(TAG, "Attempting to set a device owner type when there is no device owner");
                 return;
-            } else if (isDeviceOwnerTypeSetForDeviceOwner(packageName)) {
-                Slog.e(TAG, "Device owner type for " + packageName + " has already been set");
+            } else if (!isAdminTestOnly && isDeviceOwnerTypeSetForDeviceOwner(packageName)) {
+                Slog.e(TAG, "Setting the device owner type more than once is only allowed"
+                        + " for test only admins");
                 return;
             }
 
diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
index 64ce6b2..c877bd1 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
@@ -7737,30 +7737,20 @@
 
         dpm.setDeviceOwnerType(admin1, DEVICE_OWNER_TYPE_FINANCED);
 
-        int returnedDeviceOwnerType = dpm.getDeviceOwnerType(admin1);
-        assertThat(dpms.mOwners.hasDeviceOwner()).isTrue();
-        assertThat(returnedDeviceOwnerType).isEqualTo(DEVICE_OWNER_TYPE_FINANCED);
-
+        assertThat(dpm.getDeviceOwnerType(admin1)).isEqualTo(DEVICE_OWNER_TYPE_FINANCED);
         initializeDpms();
-
-        returnedDeviceOwnerType = dpm.getDeviceOwnerType(admin1);
-        assertThat(dpms.mOwners.hasDeviceOwner()).isTrue();
-        assertThat(returnedDeviceOwnerType).isEqualTo(DEVICE_OWNER_TYPE_FINANCED);
+        assertThat(dpm.getDeviceOwnerType(admin1)).isEqualTo(DEVICE_OWNER_TYPE_FINANCED);
     }
 
     @Test
-    public void testSetDeviceOwnerType_asDeviceOwner_throwsExceptionWhenSetDeviceOwnerTypeAgain()
+    public void testSetDeviceOwnerType_asDeviceOwner_setDeviceOwnerTypeTwice_success()
             throws Exception {
         setDeviceOwner();
+        dpm.setDeviceOwnerType(admin1, DEVICE_OWNER_TYPE_DEFAULT);
 
         dpm.setDeviceOwnerType(admin1, DEVICE_OWNER_TYPE_FINANCED);
 
-        int returnedDeviceOwnerType = dpm.getDeviceOwnerType(admin1);
-        assertThat(dpms.mOwners.hasDeviceOwner()).isTrue();
-        assertThat(returnedDeviceOwnerType).isEqualTo(DEVICE_OWNER_TYPE_FINANCED);
-
-        assertThrows(IllegalStateException.class,
-                () -> dpm.setDeviceOwnerType(admin1, DEVICE_OWNER_TYPE_DEFAULT));
+        assertThat(dpm.getDeviceOwnerType(admin1)).isEqualTo(DEVICE_OWNER_TYPE_FINANCED);
     }
 
     @Test
diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/OwnersTest.java b/services/tests/servicestests/src/com/android/server/devicepolicy/OwnersTest.java
index 02a8ae8..9a5254d 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/OwnersTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/OwnersTest.java
@@ -93,7 +93,7 @@
             assertThat(owners.getProfileOwnerUserRestrictionsNeedsMigration(21)).isFalse();
 
             owners.setDeviceOwnerType(owners.getDeviceOwnerPackageName(),
-                    DEVICE_OWNER_TYPE_FINANCED);
+                    DEVICE_OWNER_TYPE_FINANCED, /* isAdminTestOnly= */ false);
             // There is no device owner, so the default owner type should be returned.
             assertThat(owners.getDeviceOwnerType(owners.getDeviceOwnerPackageName())).isEqualTo(
                     DEVICE_OWNER_TYPE_DEFAULT);
@@ -367,7 +367,7 @@
             owners.setDeviceOwnerUserRestrictionsMigrated();
 
             owners.setDeviceOwnerType(owners.getDeviceOwnerPackageName(),
-                    DEVICE_OWNER_TYPE_FINANCED);
+                    DEVICE_OWNER_TYPE_FINANCED, /* isAdminTestOnly= */ false);
             assertThat(owners.getDeviceOwnerType(owners.getDeviceOwnerPackageName())).isEqualTo(
                     DEVICE_OWNER_TYPE_FINANCED);
 
@@ -399,7 +399,7 @@
             owners.setProfileOwnerUserRestrictionsMigrated(11);
 
             owners.setDeviceOwnerType(owners.getDeviceOwnerPackageName(),
-                    DEVICE_OWNER_TYPE_DEFAULT);
+                    DEVICE_OWNER_TYPE_DEFAULT, /* isAdminTestOnly= */ false);
             // The previous device owner type should persist.
             assertThat(owners.getDeviceOwnerType(owners.getDeviceOwnerPackageName())).isEqualTo(
                     DEVICE_OWNER_TYPE_FINANCED);
@@ -585,7 +585,8 @@
         assertThat(owners.getProfileOwnerFile(11).exists()).isTrue();
 
         String previousDeviceOwnerPackageName = owners.getDeviceOwnerPackageName();
-        owners.setDeviceOwnerType(previousDeviceOwnerPackageName, DEVICE_OWNER_TYPE_FINANCED);
+        owners.setDeviceOwnerType(previousDeviceOwnerPackageName, DEVICE_OWNER_TYPE_FINANCED,
+                /* isAdminTestOnly= */ false);
         assertThat(owners.getDeviceOwnerType(previousDeviceOwnerPackageName)).isEqualTo(
                 DEVICE_OWNER_TYPE_FINANCED);
         owners.setDeviceOwnerProtectedPackages(