Add explicit and persistent user provisioning state.

Add explicit modelling of provisioning state so that integration
of management provisioning flows with packages such as setup-wizard
are cleaner, and can be more direct. Previously we relied upon
USER_SETUP_COMPLETE secure setting and HOME intents to signal intent,
but this is not very clear and can be fragile.

Bug: 25858670
Change-Id: Idc56a040f710c3aee281db420f21717da3960722

5 files changed, 370 insertions, 7 deletions

diff --git a/cmds/dpm/src/com/android/commands/dpm/Dpm.java b/cmds/dpm/src/com/android/commands/dpm/Dpm.java
index 6dc3cd1d48e0..b83484d99824 100644
--- a/cmds/dpm/src/com/android/commands/dpm/Dpm.java
+++ b/cmds/dpm/src/com/android/commands/dpm/Dpm.java
@@ -18,6 +18,7 @@ package com.android.commands.dpm;
 
 import android.app.ActivityManagerNative;
 import android.app.IActivityManager;
+import android.app.admin.DevicePolicyManager;
 import android.app.admin.IDevicePolicyManager;
 import android.content.ComponentName;
 import android.content.Context;
@@ -143,6 +144,10 @@ public final class Dpm extends BaseCommand {
             mDevicePolicyManager.removeActiveAdmin(mComponent, UserHandle.USER_SYSTEM);
             throw e;
         }
+
+        mDevicePolicyManager.setUserProvisioningState(
+                DevicePolicyManager.STATE_USER_SETUP_FINALIZED, mUserId);
+
         System.out.println("Success: Device owner set to package " + mComponent);
         System.out.println("Active admin set to component " + mComponent.toShortString());
     }
@@ -161,6 +166,10 @@ public final class Dpm extends BaseCommand {
             mDevicePolicyManager.removeActiveAdmin(mComponent, mUserId);
             throw e;
         }
+
+        mDevicePolicyManager.setUserProvisioningState(
+                DevicePolicyManager.STATE_USER_SETUP_FINALIZED, mUserId);
+
         System.out.println("Success: Active admin and profile owner set to "
                 + mComponent.toShortString() + " for user " + mUserId);
     }
@@ -180,4 +189,4 @@ public final class Dpm extends BaseCommand {
             throw new IllegalArgumentException ("Invalid integer argument '" + argument + "'", e);
         }
     }
-}
\ No newline at end of file
+}
diff --git a/core/java/android/app/admin/DevicePolicyManager.java b/core/java/android/app/admin/DevicePolicyManager.java
index 25c54fa522d9..b098d045d19f 100644
--- a/core/java/android/app/admin/DevicePolicyManager.java
+++ b/core/java/android/app/admin/DevicePolicyManager.java
@@ -16,6 +16,7 @@
 
 package android.app.admin;
 
+import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SdkConstant;
@@ -53,6 +54,8 @@ import org.xmlpull.v1.XmlPullParserException;
 
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
 import java.net.InetSocketAddress;
 import java.net.Proxy;
 import java.security.KeyFactory;
@@ -280,6 +283,21 @@ public class DevicePolicyManager {
         = "android.app.action.PROVISION_MANAGED_SHAREABLE_DEVICE";
 
     /**
+     * Activity action: Finalizes management provisioning, should be used after user-setup
+     * has been completed and {@link #getUserProvisioningState()} returns one of:
+     * <ul>
+     * <li>{@link #STATE_USER_SETUP_INCOMPLETE}</li>
+     * <li>{@link #STATE_USER_SETUP_COMPLETE}</li>
+     * <li>{@link #STATE_USER_PROFILE_COMPLETE}</li>
+     * </ul>
+     *
+     * @hide
+     */
+    @SdkConstant(SdkConstantType.ACTIVITY_INTENT_ACTION)
+    public static final String ACTION_PROVISION_FINALIZATION
+            = "android.app.action.PROVISION_FINALIZATION";
+
+    /**
      * A {@link android.os.Parcelable} extra of type {@link android.os.PersistableBundle} that
      * allows a mobile device management application or NFC programmer application which starts
      * managed provisioning to pass data to the management application instance after provisioning.
@@ -861,6 +879,44 @@ public class DevicePolicyManager {
     public static final int PERMISSION_GRANT_STATE_DENIED = 2;
 
     /**
+     * No management for current user in-effect. This is the default.
+     * @hide
+     */
+    public static final int STATE_USER_UNMANAGED = 0;
+
+    /**
+     * Management partially setup, user setup needs to be completed.
+     * @hide
+     */
+    public static final int STATE_USER_SETUP_INCOMPLETE = 1;
+
+    /**
+     * Management partially setup, user setup completed.
+     * @hide
+     */
+    public static final int STATE_USER_SETUP_COMPLETE = 2;
+
+    /**
+     * Management setup and active on current user.
+     * @hide
+     */
+    public static final int STATE_USER_SETUP_FINALIZED = 3;
+
+    /**
+     * Management partially setup on a managed profile.
+     * @hide
+     */
+    public static final int STATE_USER_PROFILE_COMPLETE = 4;
+
+    /**
+     * @hide
+     */
+    @IntDef({STATE_USER_UNMANAGED, STATE_USER_SETUP_INCOMPLETE, STATE_USER_SETUP_COMPLETE,
+            STATE_USER_SETUP_FINALIZED, STATE_USER_PROFILE_COMPLETE})
+    @Retention(RetentionPolicy.SOURCE)
+    public @interface UserProvisioningState {}
+
+    /**
      * Return true if the given administrator component is currently
      * active (enabled) in the system.
      */
@@ -5203,6 +5259,40 @@ public class DevicePolicyManager {
     }
 
     /**
+     * @return the {@link UserProvisioningState} for the current user - for unmanaged users will
+     *         return {@link #STATE_USER_UNMANAGED}
+     * @hide
+     */
+    @UserProvisioningState
+    public int getUserProvisioningState() {
+        if (mService != null) {
+            try {
+                return mService.getUserProvisioningState();
+            } catch (RemoteException e) {
+                Log.w(TAG, REMOTE_EXCEPTION_MESSAGE, e);
+            }
+        }
+        return STATE_USER_UNMANAGED;
+    }
+
+    /**
+     * Set the {@link UserProvisioningState} for the supplied user, if they are managed.
+     *
+     * @param state to store
+     * @param userHandle for user
+     * @hide
+     */
+    public void setUserProvisioningState(@UserProvisioningState int state, int userHandle) {
+        if (mService != null) {
+            try {
+                mService.setUserProvisioningState(state, userHandle);
+            } catch (RemoteException e) {
+                Log.w(TAG, REMOTE_EXCEPTION_MESSAGE, e);
+            }
+        }
+    }
+
+    /**
      * @hide
      * Indicates the entity that controls the device or profile owner. A user/profile is considered
      * affiliated if it is managed by the same entity as the device.
diff --git a/core/java/android/app/admin/IDevicePolicyManager.aidl b/core/java/android/app/admin/IDevicePolicyManager.aidl
index 2b378a4fbaff..25cadf93ad6f 100644
--- a/core/java/android/app/admin/IDevicePolicyManager.aidl
+++ b/core/java/android/app/admin/IDevicePolicyManager.aidl
@@ -268,6 +268,9 @@ interface IDevicePolicyManager {
     int getOrganizationColor(in ComponentName admin);
     int getOrganizationColorForUser(int userHandle);
 
+    int getUserProvisioningState();
+    void setUserProvisioningState(int state, int userHandle);
+
     void setAffiliationIds(in ComponentName admin, in List<String> ids);
     boolean isAffiliatedUser();
 }
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 74d46592a70c..696604163337 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -196,6 +196,7 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
 
     private static final String ATTR_PERMISSION_PROVIDER = "permission-provider";
     private static final String ATTR_SETUP_COMPLETE = "setup-complete";
+    private static final String ATTR_PROVISIONING_STATE = "provisioning-state";
     private static final String ATTR_PERMISSION_POLICY = "permission-policy";
 
     private static final String ATTR_DELEGATED_CERT_INSTALLER = "delegated-cert-installer";
@@ -358,6 +359,7 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
         int mPasswordOwner = -1;
         long mLastMaximumTimeToLock = -1;
         boolean mUserSetupComplete = false;
+        int mUserProvisioningState;
         int mPermissionPolicy;
 
         final ArrayMap<ComponentName, ActiveAdmin> mAdminMap = new ArrayMap<>();
@@ -2008,6 +2010,10 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
                 out.attribute(null, ATTR_SETUP_COMPLETE,
                         Boolean.toString(true));
             }
+            if (policy.mUserProvisioningState != DevicePolicyManager.STATE_USER_UNMANAGED) {
+                out.attribute(null, ATTR_PROVISIONING_STATE,
+                        Integer.toString(policy.mUserProvisioningState));
+            }
             if (policy.mPermissionPolicy != DevicePolicyManager.PERMISSION_POLICY_PROMPT) {
                 out.attribute(null, ATTR_PERMISSION_POLICY,
                         Integer.toString(policy.mPermissionPolicy));
@@ -2145,6 +2151,10 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
             if (userSetupComplete != null && Boolean.toString(true).equals(userSetupComplete)) {
                 policy.mUserSetupComplete = true;
             }
+            String provisioningState = parser.getAttributeValue(null, ATTR_PROVISIONING_STATE);
+            if (!TextUtils.isEmpty(provisioningState)) {
+                policy.mUserProvisioningState = Integer.parseInt(provisioningState);
+            }
             String permissionPolicy = parser.getAttributeValue(null, ATTR_PERMISSION_POLICY);
             if (!TextUtils.isEmpty(permissionPolicy)) {
                 policy.mPermissionPolicy = Integer.parseInt(permissionPolicy);
@@ -5352,6 +5362,7 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
         policy.mDelegatedCertInstallerPackage = null;
         policy.mApplicationRestrictionsManagingPackage = null;
         policy.mStatusBarDisabled = false;
+        policy.mUserProvisioningState = DevicePolicyManager.STATE_USER_UNMANAGED;
         saveSettingsLocked(userId);
 
         final long ident = mInjector.binderClearCallingIdentity();
@@ -5379,6 +5390,98 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
     }
 
     @Override
+    public int getUserProvisioningState() {
+        if (!mHasFeature) {
+            return DevicePolicyManager.STATE_USER_UNMANAGED;
+        }
+        int userHandle = mInjector.userHandleGetCallingUserId();
+        return getUserProvisioningState(userHandle);
+    }
+
+    private int getUserProvisioningState(int userHandle) {
+        return getUserData(userHandle).mUserProvisioningState;
+    }
+
+    @Override
+    public void setUserProvisioningState(int newState, int userHandle) {
+        if (!mHasFeature) {
+            return;
+        }
+
+        if (userHandle != mOwners.getDeviceOwnerUserId() && !mOwners.hasProfileOwner(userHandle)
+                && getManagedUserId(userHandle) == -1) {
+            // No managed device, user or profile, so setting provisioning state makes no sense.
+            throw new IllegalStateException("Not allowed to change provisioning state unless a "
+                      + "device or profile owner is set.");
+        }
+
+        synchronized (this) {
+            boolean transitionCheckNeeded = true;
+
+            // Calling identity/permission checks.
+            final int callingUid = mInjector.binderGetCallingUid();
+            if (callingUid == Process.SHELL_UID || callingUid == Process.ROOT_UID) {
+                // ADB shell can only move directly from un-managed to finalized as part of directly
+                // setting profile-owner or device-owner.
+                if (getUserProvisioningState(userHandle) !=
+                        DevicePolicyManager.STATE_USER_UNMANAGED
+                        || newState != DevicePolicyManager.STATE_USER_SETUP_FINALIZED) {
+                    throw new IllegalStateException("Not allowed to change provisioning state "
+                            + "unless current provisioning state is unmanaged, and new state is "
+                            + "finalized.");
+                }
+                transitionCheckNeeded = false;
+            } else {
+                // For all other cases, caller must have MANAGE_PROFILE_AND_DEVICE_OWNERS.
+                mContext.enforceCallingOrSelfPermission(
+                        android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS, null);
+            }
+
+            final DevicePolicyData policyData = getUserData(userHandle);
+            if (transitionCheckNeeded) {
+                // Optional state transition check for non-ADB case.
+                checkUserProvisioningStateTransition(policyData.mUserProvisioningState, newState);
+            }
+            policyData.mUserProvisioningState = newState;
+            saveSettingsLocked(userHandle);
+        }
+    }
+
+    private void checkUserProvisioningStateTransition(int currentState, int newState) {
+        // Valid transitions for normal use-cases.
+        switch (currentState) {
+            case DevicePolicyManager.STATE_USER_UNMANAGED:
+                // Can move to any state from unmanaged (except itself as an edge case)..
+                if (newState != DevicePolicyManager.STATE_USER_UNMANAGED) {
+                    return;
+                }
+                break;
+            case DevicePolicyManager.STATE_USER_SETUP_INCOMPLETE:
+            case DevicePolicyManager.STATE_USER_SETUP_COMPLETE:
+                // Can only move to finalized from these states.
+                if (newState == DevicePolicyManager.STATE_USER_SETUP_FINALIZED) {
+                    return;
+                }
+                break;
+            case DevicePolicyManager.STATE_USER_PROFILE_COMPLETE:
+                // Current user has a managed-profile, but current user is not managed, so
+                // rather than moving to finalized state, go back to unmanaged once
+                // profile provisioning is complete.
+                if (newState == DevicePolicyManager.STATE_USER_UNMANAGED) {
+                    return;
+                }
+                break;
+            case DevicePolicyManager.STATE_USER_SETUP_FINALIZED:
+                // Cannot transition out of finalized.
+                break;
+        }
+
+        // Didn't meet any of the accepted state transition checks above, throw appropriate error.
+        throw new IllegalStateException("Cannot move to user provisioning state [" + newState + "] "
+                + "from state [" + currentState + "]");
+    }
+
+    @Override
     public void setProfileEnabled(ComponentName who) {
         if (!mHasFeature) {
             return;
@@ -5697,7 +5800,8 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
             for (int u = 0; u < userCount; u++) {
                 DevicePolicyData policy = getUserData(mUserData.keyAt(u));
                 pw.println();
-                pw.println("  Enabled Device Admins (User " + policy.mUserHandle + "):");
+                pw.println("  Enabled Device Admins (User " + policy.mUserHandle
+                        + ", provisioningState: " + policy.mUserProvisioningState + "):");
                 final int N = policy.mAdminList.size();
                 for (int i=0; i<N; i++) {
                     ActiveAdmin ap = policy.mAdminList.get(i);
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 536fb70d9068..72421ae1f959 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
@@ -15,9 +15,6 @@
  */
 package com.android.server.devicepolicy;
 
-import com.android.server.LocalServices;
-import com.android.server.SystemService;
-
 import android.Manifest.permission;
 import android.app.Activity;
 import android.app.admin.DeviceAdminReceiver;
@@ -27,7 +24,6 @@ import android.content.BroadcastReceiver;
 import android.content.ComponentName;
 import android.content.pm.PackageManager;
 import android.net.wifi.WifiInfo;
-import android.os.Build;
 import android.os.Build.VERSION_CODES;
 import android.os.Bundle;
 import android.os.Process;
@@ -37,11 +33,16 @@ import android.test.MoreAsserts;
 import android.util.ArraySet;
 import android.util.Pair;
 
+import com.android.server.LocalServices;
+import com.android.server.SystemService;
+
 import org.mockito.ArgumentCaptor;
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
 
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -64,13 +65,17 @@ import static org.mockito.Mockito.when;
  *
  m FrameworksServicesTests &&
  adb install \
-   -r out/target/product/hammerhead/data/app/FrameworksServicesTests/FrameworksServicesTests.apk &&
+   -r ${ANDROID_PRODUCT_OUT}/data/app/FrameworksServicesTests/FrameworksServicesTests.apk &&
  adb shell am instrument -e class com.android.server.devicepolicy.DevicePolicyManagerTest \
    -w com.android.frameworks.servicestests/android.support.test.runner.AndroidJUnitRunner
 
  (mmma frameworks/base/services/tests/servicestests/ for non-ninja build)
  */
 public class DevicePolicyManagerTest extends DpmTestBase {
+    private static final List<String> OWNER_SETUP_PERMISSIONS = Arrays.asList(
+            permission.MANAGE_DEVICE_ADMINS, permission.MANAGE_PROFILE_AND_DEVICE_OWNERS,
+            permission.MANAGE_USERS, permission.INTERACT_ACROSS_USERS_FULL);
+
     private DpmMockContext mContext;
     public DevicePolicyManager dpm;
     public DevicePolicyManagerServiceTestable dpms;
@@ -1543,4 +1548,156 @@ public class DevicePolicyManagerTest extends DpmTestBase {
         mContext.binder.callingUid = DpmMockContext.CALLER_SYSTEM_USER_UID;
         assertTrue(dpm.isAffiliatedUser());
     }
+
+    public void testGetUserProvisioningState_defaultResult() {
+        assertEquals(DevicePolicyManager.STATE_USER_UNMANAGED, dpm.getUserProvisioningState());
+    }
+
+    public void testSetUserProvisioningState_permission() throws Exception {
+        setupProfileOwner();
+        mContext.callerPermissions.add(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS);
+
+        exerciseUserProvisioningTransitions(DpmMockContext.CALLER_USER_HANDLE,
+                DevicePolicyManager.STATE_USER_SETUP_FINALIZED);
+    }
+
+    public void testSetUserProvisioningState_unprivileged() throws Exception {
+        setupProfileOwner();
+        try {
+            dpm.setUserProvisioningState(DevicePolicyManager.STATE_USER_SETUP_FINALIZED,
+                    DpmMockContext.CALLER_USER_HANDLE);
+            fail("Expected SecurityException");
+        } catch (SecurityException expected) {
+        }
+    }
+
+    public void testSetUserProvisioningState_noManagement() {
+        mContext.callerPermissions.add(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS);
+        try {
+            dpm.setUserProvisioningState(DevicePolicyManager.STATE_USER_SETUP_FINALIZED,
+                    DpmMockContext.CALLER_USER_HANDLE);
+            fail("IllegalStateException expected");
+        } catch (IllegalStateException e) {
+            MoreAsserts.assertContainsRegex("change provisioning state unless a .* owner is set",
+                    e.getMessage());
+        }
+        assertEquals(DevicePolicyManager.STATE_USER_UNMANAGED, dpm.getUserProvisioningState());
+    }
+
+    public void testSetUserProvisioningState_deviceOwnerFromSetupWizard() throws Exception {
+        mContext.binder.callingUid = DpmMockContext.CALLER_SYSTEM_USER_UID;
+        setupDeviceOwner();
+        mContext.callerPermissions.add(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS);
+
+        exerciseUserProvisioningTransitions(UserHandle.USER_SYSTEM,
+                DevicePolicyManager.STATE_USER_SETUP_COMPLETE,
+                DevicePolicyManager.STATE_USER_SETUP_FINALIZED);
+    }
+
+    public void testSetUserProvisioningState_deviceOwnerFromSetupWizardAlternative()
+            throws Exception {
+        mContext.binder.callingUid = DpmMockContext.CALLER_SYSTEM_USER_UID;
+        setupDeviceOwner();
+        mContext.callerPermissions.add(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS);
+
+        exerciseUserProvisioningTransitions(UserHandle.USER_SYSTEM,
+                DevicePolicyManager.STATE_USER_SETUP_INCOMPLETE,
+                DevicePolicyManager.STATE_USER_SETUP_FINALIZED);
+    }
+
+    public void testSetUserProvisioningState_deviceOwnerWithoutSetupWizard() throws Exception {
+        mContext.binder.callingUid = DpmMockContext.CALLER_SYSTEM_USER_UID;
+        setupDeviceOwner();
+        mContext.callerPermissions.add(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS);
+
+        exerciseUserProvisioningTransitions(UserHandle.USER_SYSTEM,
+                DevicePolicyManager.STATE_USER_SETUP_FINALIZED);
+    }
+
+    public void testSetUserProvisioningState_managedProfileFromSetupWizard_primaryUser()
+            throws Exception {
+        setupProfileOwner();
+        mContext.callerPermissions.add(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS);
+
+        exerciseUserProvisioningTransitions(DpmMockContext.CALLER_USER_HANDLE,
+                DevicePolicyManager.STATE_USER_PROFILE_COMPLETE,
+                DevicePolicyManager.STATE_USER_UNMANAGED);
+    }
+
+    public void testSetUserProvisioningState_managedProfileFromSetupWizard_managedProfile()
+            throws Exception {
+        setupProfileOwner();
+        mContext.callerPermissions.add(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS);
+
+        exerciseUserProvisioningTransitions(DpmMockContext.CALLER_USER_HANDLE,
+                DevicePolicyManager.STATE_USER_SETUP_COMPLETE,
+                DevicePolicyManager.STATE_USER_SETUP_FINALIZED);
+    }
+
+    public void testSetUserProvisioningState_managedProfileWithoutSetupWizard() throws Exception {
+        setupProfileOwner();
+        mContext.callerPermissions.add(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS);
+
+        exerciseUserProvisioningTransitions(DpmMockContext.CALLER_USER_HANDLE,
+                DevicePolicyManager.STATE_USER_SETUP_FINALIZED);
+    }
+
+    public void testSetUserProvisioningState_illegalTransitionOutOfFinalized1() throws Exception {
+        setupProfileOwner();
+        mContext.callerPermissions.add(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS);
+
+        try {
+            exerciseUserProvisioningTransitions(DpmMockContext.CALLER_USER_HANDLE,
+                    DevicePolicyManager.STATE_USER_SETUP_FINALIZED,
+                    DevicePolicyManager.STATE_USER_UNMANAGED);
+            fail("Expected IllegalStateException");
+        } catch (IllegalStateException e) {
+            MoreAsserts.assertContainsRegex("Cannot move to user provisioning state",
+                    e.getMessage());
+        }
+    }
+
+    public void testSetUserProvisioningState_illegalTransitionToAnotherInProgressState()
+            throws Exception {
+        setupProfileOwner();
+        mContext.callerPermissions.add(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS);
+
+        try {
+            exerciseUserProvisioningTransitions(DpmMockContext.CALLER_USER_HANDLE,
+                    DevicePolicyManager.STATE_USER_SETUP_INCOMPLETE,
+                    DevicePolicyManager.STATE_USER_SETUP_COMPLETE);
+            fail("Expected IllegalStateException");
+        } catch (IllegalStateException e) {
+            MoreAsserts.assertContainsRegex("Cannot move to user provisioning state",
+                    e.getMessage());
+        }
+    }
+
+    private void exerciseUserProvisioningTransitions(int userId, int... states) {
+        assertEquals(DevicePolicyManager.STATE_USER_UNMANAGED, dpm.getUserProvisioningState());
+        for (int state : states) {
+            dpm.setUserProvisioningState(state, userId);
+            assertEquals(state, dpm.getUserProvisioningState());
+        }
+    }
+
+    private void setupProfileOwner() throws Exception {
+        mContext.callerPermissions.addAll(OWNER_SETUP_PERMISSIONS);
+
+        setUpPackageManagerForAdmin(admin1, DpmMockContext.CALLER_UID);
+        dpm.setActiveAdmin(admin1, false);
+        assertTrue(dpm.setProfileOwner(admin1, null, DpmMockContext.CALLER_USER_HANDLE));
+
+        mContext.callerPermissions.removeAll(OWNER_SETUP_PERMISSIONS);
+    }
+
+    private void setupDeviceOwner() throws Exception {
+        mContext.callerPermissions.addAll(OWNER_SETUP_PERMISSIONS);
+
+        setUpPackageManagerForAdmin(admin1, DpmMockContext.CALLER_SYSTEM_USER_UID);
+        dpm.setActiveAdmin(admin1, false);
+        assertTrue(dpm.setDeviceOwner(admin1, null, UserHandle.USER_SYSTEM));
+
+        mContext.callerPermissions.removeAll(OWNER_SETUP_PERMISSIONS);
+    }
 }