Use DevicePolicyManager.setDeviceOwner on S+.

This allows us to enter the state of running on a secondary user on a
device with a device owner - except that currently
forceUpdateUserSetupComplete doesn't allow calling from secondary users.

Test: atest HarrierTest
Test: atest NeneTest
Bug: 177663307
Change-Id: Ib69f55906ff31c326510c9c15214fbec367a54f1
diff --git a/common/device-side/bedstead/harrier/src/main/java/com/android/bedstead/harrier/DeviceState.java b/common/device-side/bedstead/harrier/src/main/java/com/android/bedstead/harrier/DeviceState.java
index 8f9c208..6ffde6c 100644
--- a/common/device-side/bedstead/harrier/src/main/java/com/android/bedstead/harrier/DeviceState.java
+++ b/common/device-side/bedstead/harrier/src/main/java/com/android/bedstead/harrier/DeviceState.java
@@ -28,6 +28,7 @@
 
 import android.content.Context;
 import android.content.Intent;
+import android.os.Build;
 import android.os.Bundle;
 import android.util.Log;
 
@@ -69,6 +70,7 @@
 import com.android.bedstead.nene.users.UserBuilder;
 import com.android.bedstead.nene.users.UserReference;
 import com.android.bedstead.nene.utils.ShellCommand;
+import com.android.bedstead.nene.utils.Versions;
 import com.android.bedstead.remotedpc.RemoteDpc;
 import com.android.compatibility.common.util.BlockingBroadcastReceiver;
 
@@ -218,7 +220,8 @@
                     if (annotation instanceof EnsureHasDeviceOwner) {
                         EnsureHasDeviceOwner ensureHasDeviceOwnerAnnotation =
                                 (EnsureHasDeviceOwner) annotation;
-                        ensureHasDeviceOwner(ensureHasDeviceOwnerAnnotation.onUser());
+                        ensureHasDeviceOwner(ensureHasDeviceOwnerAnnotation.onUser(),
+                                ensureHasDeviceOwnerAnnotation.failureMode());
                     }
 
                     if (annotation instanceof EnsureHasNoDeviceOwner) {
@@ -911,7 +914,7 @@
         }
     }
 
-    private void ensureHasDeviceOwner(UserType onUser) {
+    private void ensureHasDeviceOwner(UserType onUser, FailureMode failureMode) {
         // TODO(scottjonathan): Should support non-remotedpc device owner (default to remotedpc)
         // TODO(scottjonathan): Should allow setting the device owner on a different user
         DeviceOwner currentDeviceOwner = sTestApis.devicePolicy().getDeviceOwner();
@@ -923,15 +926,24 @@
 
         UserReference instrumentedUser = sTestApis.users().instrumented();
 
-        // TODO(scottjonathan): Consider if we should restore these users
-        for (UserReference u : sTestApis.users().all()) {
-            if (u.equals(instrumentedUser)) {
-                continue;
-            }
 
-            removeAndRecordUser(u);
+        if (!Versions.meetsMinimumSdkVersionRequirement(Build.VERSION_CODES.S)) {
+            // Prior to S we can't set device owner if there are other users on the device
+            for (UserReference u : sTestApis.users().all()) {
+                if (u.equals(instrumentedUser)) {
+                    continue;
+                }
+                try {
+                    removeAndRecordUser(u);
+                } catch (NeneException e) {
+                    failOrSkip(
+                            "Error removing user to prepare for DeviceOwner: " + e.toString(),
+                            failureMode);
+                }
+            }
         }
 
+
         // TODO(scottjonathan): Remove accounts
         ensureHasNoProfileOwner(onUser);
 
diff --git a/common/device-side/bedstead/harrier/src/main/java/com/android/bedstead/harrier/annotations/enterprise/EnsureHasDeviceOwner.java b/common/device-side/bedstead/harrier/src/main/java/com/android/bedstead/harrier/annotations/enterprise/EnsureHasDeviceOwner.java
index 99863a9..77ae0f3 100644
--- a/common/device-side/bedstead/harrier/src/main/java/com/android/bedstead/harrier/annotations/enterprise/EnsureHasDeviceOwner.java
+++ b/common/device-side/bedstead/harrier/src/main/java/com/android/bedstead/harrier/annotations/enterprise/EnsureHasDeviceOwner.java
@@ -19,6 +19,7 @@
 import static com.android.bedstead.harrier.DeviceState.UserType.SYSTEM_USER;
 
 import com.android.bedstead.harrier.DeviceState;
+import com.android.bedstead.harrier.annotations.FailureMode;
 
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;
@@ -41,6 +42,9 @@
 public @interface EnsureHasDeviceOwner {
     /** Which user type the device owner should be installed on. */
     DeviceState.UserType onUser() default SYSTEM_USER;
+
+    /** Behaviour if the device owner cannot be set. */
+    FailureMode failureMode() default FailureMode.FAIL;
 }
 // TODO(scottjonathan): Is there a feature or something that we need to check to make sure DO is
 //  supported?
diff --git a/common/device-side/bedstead/harrier/src/main/java/com/android/bedstead/harrier/annotations/parameterized/IncludeRunOnNonAffiliatedDeviceOwnerSecondaryUser.java b/common/device-side/bedstead/harrier/src/main/java/com/android/bedstead/harrier/annotations/parameterized/IncludeRunOnNonAffiliatedDeviceOwnerSecondaryUser.java
index d32f9ba..3517f45 100644
--- a/common/device-side/bedstead/harrier/src/main/java/com/android/bedstead/harrier/annotations/parameterized/IncludeRunOnNonAffiliatedDeviceOwnerSecondaryUser.java
+++ b/common/device-side/bedstead/harrier/src/main/java/com/android/bedstead/harrier/annotations/parameterized/IncludeRunOnNonAffiliatedDeviceOwnerSecondaryUser.java
@@ -34,9 +34,6 @@
 @Target({ElementType.METHOD, ElementType.TYPE})
 @Retention(RetentionPolicy.RUNTIME)
 @ParameterizedAnnotation
-// TODO(scottjonathan): We can't set the device owner when we're already running on secondary user
-//  so we need to have a different @RequireRunOn annotation for a secondary user when a DO is
-//  already set
 @RequireRunOnSecondaryUser
 @EnsureHasDeviceOwner(onUser = SYSTEM_USER)
 public @interface IncludeRunOnNonAffiliatedDeviceOwnerSecondaryUser {
diff --git a/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/AdbDevicePolicyParser27.java b/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/AdbDevicePolicyParser27.java
index 3a20f6c..de24e0d 100644
--- a/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/AdbDevicePolicyParser27.java
+++ b/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/AdbDevicePolicyParser27.java
@@ -71,7 +71,7 @@
                         "ComponentInfo\\{", 2)[1].split("\\}", 2)[0]);
         int userId = Integer.parseInt(deviceOwnerSection.split(
                 "User ID: ", 2)[1].split("\n", 2)[0]);
-        return new DeviceOwner(mTestApis.users().find(userId),
+        return new DeviceOwner(mTestApis, mTestApis.users().find(userId),
                 mTestApis.packages().find(componentName.getPackageName()), componentName);
     }
 
@@ -114,7 +114,7 @@
                         "ComponentInfo\\{", 2)[1].split("\\}", 2)[0]);
         int userId = Integer.parseInt(
                 profileOwnerSection.split("\\(User ", 2)[1].split("\\)", 2)[0]);
-        return new ProfileOwner(mTestApis.users().find(userId),
+        return new ProfileOwner(mTestApis, mTestApis.users().find(userId),
                 mTestApis.packages().find(componentName.getPackageName()), componentName);
     }
 }
diff --git a/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/DeviceOwner.java b/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/DeviceOwner.java
index 3d9a4ea..c46791b 100644
--- a/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/DeviceOwner.java
+++ b/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/DeviceOwner.java
@@ -18,6 +18,7 @@
 
 import android.content.ComponentName;
 
+import com.android.bedstead.nene.TestApis;
 import com.android.bedstead.nene.exceptions.AdbException;
 import com.android.bedstead.nene.exceptions.NeneException;
 import com.android.bedstead.nene.packages.PackageReference;
@@ -32,10 +33,11 @@
  */
 public final class DeviceOwner extends DevicePolicyController {
 
-    DeviceOwner(UserReference user,
+    DeviceOwner(TestApis testApis,
+            UserReference user,
             PackageReference pkg,
             ComponentName componentName) {
-        super(user, pkg, componentName);
+        super(testApis, user, pkg, componentName);
     }
 
     @Override
diff --git a/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/DevicePolicy.java b/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/DevicePolicy.java
index 8ed118e..c9e9a9e 100644
--- a/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/DevicePolicy.java
+++ b/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/DevicePolicy.java
@@ -17,12 +17,19 @@
 package com.android.bedstead.nene.devicepolicy;
 
 import static android.Manifest.permission.INTERACT_ACROSS_USERS_FULL;
+import static android.Manifest.permission.WRITE_SECURE_SETTINGS;
 import static android.os.Build.VERSION.SDK_INT;
 
+import static com.android.bedstead.nene.permissions.Permissions.MANAGE_DEVICE_ADMINS;
+import static com.android.bedstead.nene.permissions.Permissions.MANAGE_PROFILE_AND_DEVICE_OWNERS;
+
+import android.app.admin.DevicePolicyManager;
 import android.content.ComponentName;
 import android.content.Intent;
 import android.content.pm.PackageManager;
 import android.content.pm.ResolveInfo;
+import android.os.Build;
+import android.provider.Settings;
 
 import androidx.annotation.Nullable;
 
@@ -36,10 +43,15 @@
 import com.android.bedstead.nene.users.UserReference;
 import com.android.bedstead.nene.utils.ShellCommand;
 import com.android.bedstead.nene.utils.ShellCommandUtils;
+import com.android.bedstead.nene.utils.Versions;
+import com.android.compatibility.common.util.PollingCheck;
 
+import java.util.Arrays;
 import java.util.Collection;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 
 /**
@@ -47,6 +59,8 @@
  */
 public final class DevicePolicy {
 
+    private static final String USER_SETUP_COMPLETE_KEY = "user_setup_complete";
+
     private final TestApis mTestApis;
     private final AdbDevicePolicyParser mParser;
 
@@ -75,49 +89,44 @@
                 .addOperand(profileOwnerComponent.flattenToShortString())
                 .validate(ShellCommandUtils::startsWithSuccess);
 
-        try {
-            command.execute();
-        } catch (AdbException e) {
-            // If it fails, we check for terminal failure states - and if not we retry because if
-            //  the profile owner was recently removed, it can take some time to be allowed to set
-            //  it again
-
-            ProfileOwner profileOwner = getProfileOwner(user);
-            if (profileOwner != null) {
-                // TODO(scottjonathan): Should we actually fail here if the component name is the
-                //  same?
-
-                throw new NeneException(
-                        "Could not set profile owner for user " + user
-                                + " as a profile owner is already set: " + profileOwner);
-            }
-
-            PackageReference pkg = mTestApis.packages().find(
-                    profileOwnerComponent.getPackageName());
-            if (!mTestApis.packages().installedForUser(user).contains(pkg)) {
-                throw new NeneException(
-                        "Could not set profile owner for user " + user
-                                + " as the package " + pkg + " is not installed");
-            }
-
-            if (!componentCanBeSetAsDeviceAdmin(profileOwnerComponent, user)) {
-                throw new NeneException("Could not set profile owner for user "
-                        + user + " as component " + profileOwnerComponent + " is not valid");
-            }
-
-            try {
-                command.executeUntilValid();
-            } catch (AdbException | InterruptedException e2) {
-                throw new NeneException("Could not set profile owner for user " + user
-                        + " component " + profileOwnerComponent, e2);
-            }
-        }
-
-        return new ProfileOwner(user,
+        // TODO(b/187925230): If it fails, we check for terminal failure states - and if not
+        //  we retry because if the profile owner was recently removed, it can take some time
+        //  to be allowed to set it again
+        retryIfNotTerminal(
+                () -> command.executeOrThrowNeneException("Could not set profile owner for user "
+                        + user + " component " + profileOwnerComponent),
+                () -> checkForTerminalProfileOwnerFailures(user, profileOwnerComponent));
+        return new ProfileOwner(mTestApis, user,
                 mTestApis.packages().find(
                         profileOwnerComponent.getPackageName()), profileOwnerComponent);
     }
 
+    private void checkForTerminalProfileOwnerFailures(
+            UserReference user, ComponentName profileOwnerComponent) {
+        ProfileOwner profileOwner = getProfileOwner(user);
+        if (profileOwner != null) {
+            // TODO(scottjonathan): Should we actually fail here if the component name is the
+            //  same?
+
+            throw new NeneException(
+                    "Could not set profile owner for user " + user
+                            + " as a profile owner is already set: " + profileOwner);
+        }
+
+        PackageReference pkg = mTestApis.packages().find(
+                profileOwnerComponent.getPackageName());
+        if (!mTestApis.packages().installedForUser(user).contains(pkg)) {
+            throw new NeneException(
+                    "Could not set profile owner for user " + user
+                            + " as the package " + pkg + " is not installed");
+        }
+
+        if (!componentCanBeSetAsDeviceAdmin(profileOwnerComponent, user)) {
+            throw new NeneException("Could not set profile owner for user "
+                    + user + " as component " + profileOwnerComponent + " is not valid");
+        }
+    }
+
     /**
      * Get the profile owner for a given {@link UserReference}.
      */
@@ -137,43 +146,148 @@
             throw new NullPointerException();
         }
 
-        // TODO: use setDeviceOwner on S+
+        if (!Versions.meetsMinimumSdkVersionRequirement(Build.VERSION_CODES.S)) {
+            return setDeviceOwnerPreS(user, deviceOwnerComponent);
+        }
 
+        DevicePolicyManager devicePolicyManager =
+                mTestApis.context().instrumentedContext()
+                        .getSystemService(DevicePolicyManager.class);
+
+        boolean userSetupComplete = getUserSetupComplete();
+
+        try {
+            setUserSetupComplete(false);
+
+            try (PermissionContext p =
+                         mTestApis.permissions().withPermission(
+                                 MANAGE_PROFILE_AND_DEVICE_OWNERS, MANAGE_DEVICE_ADMINS)) {
+                devicePolicyManager.setActiveAdmin(deviceOwnerComponent,
+                        /* refreshing= */ true, user.id());
+
+                // TODO(b/187925230): If it fails, we check for terminal failure states - and if not
+                //  we retry because if the DO/PO was recently removed, it can take some time
+                //  to be allowed to set it again
+                retryIfNotTerminal(
+                        () -> devicePolicyManager.setDeviceOwner(
+                                deviceOwnerComponent, "Nene", user.id()),
+                        () -> checkForTerminalDeviceOwnerFailures(
+                                user, deviceOwnerComponent, /* allowAdditionalUsers= */ true));
+            } catch (IllegalArgumentException | IllegalStateException | SecurityException e) {
+                throw new NeneException("Error setting device owner", e);
+            }
+        } finally {
+            setUserSetupComplete(userSetupComplete);
+        }
+
+        return new DeviceOwner(mTestApis, user,
+                mTestApis.packages().find(
+                        deviceOwnerComponent.getPackageName()), deviceOwnerComponent);
+    }
+
+    /**
+     * Runs {@code operation}. If it fails, runs {@code terminalCheck} and then retries
+     * {@code operation} until it does not fail or for a maximum of 30 seconds.
+     *
+     * <p>The {@code operation} is considered to be successful if it does not throw an exception
+     */
+    private void retryIfNotTerminal(
+            Runnable operation, Runnable terminalCheck,
+            Class<? extends RuntimeException>... exceptions) {
+        Set<Class<? extends RuntimeException>> exceptionSet =
+                new HashSet<>(Arrays.asList(exceptions));
+        try {
+            operation.run();
+        } catch (RuntimeException e) {
+            if (!exceptionSet.contains(e.getClass())) {
+                throw e;
+            }
+
+            terminalCheck.run();
+
+            try {
+                PollingCheck.waitFor(30_000, () -> {
+                    try {
+                        operation.run();
+                        return true;
+                    } catch (RuntimeException e2) {
+                        if (!exceptionSet.contains(e2.getClass())) {
+                            throw e2;
+                        }
+                        return false;
+                    }
+                });
+            } catch (AssertionError e3) {
+                operation.run();
+            }
+        }
+    }
+
+
+    private void setUserSetupComplete(boolean complete) {
+        DevicePolicyManager devicePolicyManager =
+                mTestApis.context().instrumentedContext()
+                        .getSystemService(DevicePolicyManager.class);
+        try (PermissionContext p = mTestApis.permissions().withPermission(
+                WRITE_SECURE_SETTINGS, MANAGE_PROFILE_AND_DEVICE_OWNERS)) {
+            Settings.Secure.putInt(mTestApis.context().instrumentedContext().getContentResolver(),
+                    USER_SETUP_COMPLETE_KEY, complete ? 1 : 0);
+            devicePolicyManager.forceUpdateUserSetupComplete(mTestApis.users().instrumented().id());
+        }
+    }
+
+    private boolean getUserSetupComplete() {
+        return Settings.Secure.getInt(
+                mTestApis.context().instrumentedContext().getContentResolver(),
+                USER_SETUP_COMPLETE_KEY, /* def= */ 0) == 1;
+    }
+
+    private DeviceOwner setDeviceOwnerPreS(UserReference user, ComponentName deviceOwnerComponent) {
         ShellCommand.Builder command = ShellCommand.builderForUser(
                 user, "dpm set-device-owner")
                 .addOperand(deviceOwnerComponent.flattenToShortString())
                 .validate(ShellCommandUtils::startsWithSuccess);
 
-        try {
-            command.execute();
-        } catch (AdbException e) {
-            // If it fails, we check for terminal failure states - and if not we retry because if
-            //  the device owner was recently removed, it can take some time to be allowed to set
-            //  it again
+        // TODO(b/187925230): If it fails, we check for terminal failure states - and if not
+        //  we retry because if the device owner was recently removed, it can take some time
+        //  to be allowed to set it again
+        retryIfNotTerminal(
+                () -> command.executeOrThrowNeneException("Could not set device owner for user "
+                        + user + " component " + deviceOwnerComponent),
+                () -> checkForTerminalDeviceOwnerFailures(
+                    user, deviceOwnerComponent, /* allowAdditionalUsers= */ false));
 
-            DeviceOwner deviceOwner = getDeviceOwner();
-            if (deviceOwner != null) {
-                // TODO(scottjonathan): Should we actually fail here if the component name is the
-                //  same?
+        return new DeviceOwner(mTestApis, user,
+                mTestApis.packages().find(
+                        deviceOwnerComponent.getPackageName()), deviceOwnerComponent);
+    }
 
-                throw new NeneException(
-                        "Could not set device owner for user " + user
-                                + " as a device owner is already set: " + deviceOwner);
-            }
+    private void checkForTerminalDeviceOwnerFailures(
+            UserReference user, ComponentName deviceOwnerComponent, boolean allowAdditionalUsers) {
+        DeviceOwner deviceOwner = getDeviceOwner();
+        if (deviceOwner != null) {
+            // TODO(scottjonathan): Should we actually fail here if the component name is the
+            //  same?
 
-            PackageReference pkg = mTestApis.packages().find(
-                    deviceOwnerComponent.getPackageName());
-            if (!mTestApis.packages().installedForUser(user).contains(pkg)) {
-                throw new NeneException(
-                        "Could not set device owner for user " + user
-                                + " as the package " + pkg + " is not installed");
-            }
+            throw new NeneException(
+                    "Could not set device owner for user " + user
+                            + " as a device owner is already set: " + deviceOwner);
+        }
 
-            if (!componentCanBeSetAsDeviceAdmin(deviceOwnerComponent, user)) {
-                throw new NeneException("Could not set device owner for user "
-                        + user + " as component " + deviceOwnerComponent + " is not valid");
-            }
+        PackageReference pkg = mTestApis.packages().find(
+                deviceOwnerComponent.getPackageName());
+        if (!mTestApis.packages().installedForUser(user).contains(pkg)) {
+            throw new NeneException(
+                    "Could not set device owner for user " + user
+                            + " as the package " + pkg + " is not installed");
+        }
 
+        if (!componentCanBeSetAsDeviceAdmin(deviceOwnerComponent, user)) {
+            throw new NeneException("Could not set device owner for user "
+                    + user + " as component " + deviceOwnerComponent + " is not valid");
+        }
+
+        if (!allowAdditionalUsers) {
             Collection<User> users = mTestApis.users().all();
 
             if (users.size() > 1) {
@@ -181,19 +295,8 @@
                         + user + " as there are already additional users on the device: " + users);
             }
 
-            // TODO(scottjonathan): Check accounts
-
-            try {
-                command.executeUntilValid();
-            } catch (AdbException | InterruptedException e2) {
-                throw new NeneException("Could not set device owner for user " + user
-                        + " component " + deviceOwnerComponent, e2);
-            }
         }
-
-        return new DeviceOwner(user,
-                mTestApis.packages().find(
-                        deviceOwnerComponent.getPackageName()), deviceOwnerComponent);
+        // TODO(scottjonathan): Check accounts
     }
 
     private boolean componentCanBeSetAsDeviceAdmin(ComponentName component, UserReference user) {
diff --git a/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/DevicePolicyController.java b/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/DevicePolicyController.java
index 9a8aff3..059712f 100644
--- a/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/DevicePolicyController.java
+++ b/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/DevicePolicyController.java
@@ -19,6 +19,7 @@
 import android.app.admin.DeviceAdminReceiver;
 import android.content.ComponentName;
 
+import com.android.bedstead.nene.TestApis;
 import com.android.bedstead.nene.packages.PackageReference;
 import com.android.bedstead.nene.users.UserReference;
 
@@ -29,15 +30,18 @@
  */
 public abstract class DevicePolicyController implements AutoCloseable {
 
+    protected final TestApis mTestApis;
     protected final UserReference mUser;
     protected final PackageReference mPackage;
     protected final ComponentName mComponentName;
 
-    DevicePolicyController(UserReference user, PackageReference pkg, ComponentName componentName) {
-        if (user == null || pkg == null || componentName == null) {
+    DevicePolicyController(TestApis testApis,
+            UserReference user, PackageReference pkg, ComponentName componentName) {
+        if (testApis == null || user == null || pkg == null || componentName == null) {
             throw new NullPointerException();
         }
 
+        mTestApis = testApis;
         mUser = user;
         mPackage = pkg;
         mComponentName = componentName;
diff --git a/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/ProfileOwner.java b/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/ProfileOwner.java
index 7fb21a8..6074421 100644
--- a/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/ProfileOwner.java
+++ b/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/devicepolicy/ProfileOwner.java
@@ -18,6 +18,7 @@
 
 import android.content.ComponentName;
 
+import com.android.bedstead.nene.TestApis;
 import com.android.bedstead.nene.exceptions.AdbException;
 import com.android.bedstead.nene.exceptions.NeneException;
 import com.android.bedstead.nene.packages.PackageReference;
@@ -31,10 +32,12 @@
  * A reference to a Profile Owner.
  */
 public final class ProfileOwner extends DevicePolicyController {
-    ProfileOwner(UserReference user,
+
+    ProfileOwner(TestApis testApis,
+            UserReference user,
             PackageReference pkg,
             ComponentName componentName) {
-        super(user, pkg, componentName);
+        super(testApis, user, pkg, componentName);
     }
 
     @Override
diff --git a/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/permissions/Permissions.java b/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/permissions/Permissions.java
index 441d06f..1ede7d8 100644
--- a/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/permissions/Permissions.java
+++ b/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/permissions/Permissions.java
@@ -39,6 +39,11 @@
 /** Permission manager for tests. */
 public class Permissions {
 
+    public static final String MANAGE_PROFILE_AND_DEVICE_OWNERS =
+            "android.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS";
+
+    public static final String MANAGE_DEVICE_ADMINS = "android.permission.MANAGE_DEVICE_ADMINS";
+
     private static final String LOG_TAG = "Permissions";
 
     private List<PermissionContextImpl> mPermissionContexts = new ArrayList<>();
diff --git a/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/utils/ShellCommand.java b/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/utils/ShellCommand.java
index 0bada6b..3245462 100644
--- a/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/utils/ShellCommand.java
+++ b/common/device-side/bedstead/nene/src/main/java/com/android/bedstead/nene/utils/ShellCommand.java
@@ -20,6 +20,7 @@
 import androidx.annotation.Nullable;
 
 import com.android.bedstead.nene.exceptions.AdbException;
+import com.android.bedstead.nene.exceptions.NeneException;
 import com.android.bedstead.nene.users.UserReference;
 
 import java.util.function.Function;
@@ -128,6 +129,18 @@
             return commandBuilder.toString();
         }
 
+        /**
+         * See {@link #execute()} except that any {@link AdbException} is wrapped in a
+         * {@link NeneException} with the message {@code errorMessage}.
+         */
+        public String executeOrThrowNeneException(String errorMessage) throws NeneException {
+            try {
+                return execute();
+            } catch (AdbException e) {
+                throw new NeneException(errorMessage, e);
+            }
+        }
+
         /** See {@link ShellCommandUtils#executeCommand(java.lang.String)}. */
         public String execute() throws AdbException {
             if (mOutputSuccessChecker != null) {
diff --git a/common/device-side/bedstead/nene/src/test/java/com/android/bedstead/nene/devicepolicy/DevicePolicyTest.java b/common/device-side/bedstead/nene/src/test/java/com/android/bedstead/nene/devicepolicy/DevicePolicyTest.java
index d26fb2c..839c69a 100644
--- a/common/device-side/bedstead/nene/src/test/java/com/android/bedstead/nene/devicepolicy/DevicePolicyTest.java
+++ b/common/device-side/bedstead/nene/src/test/java/com/android/bedstead/nene/devicepolicy/DevicePolicyTest.java
@@ -18,28 +18,43 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.junit.Assume.assumeFalse;
+import static org.junit.Assume.assumeTrue;
 import static org.testng.Assert.assertThrows;
 
 import android.content.ComponentName;
+import android.os.Build;
 
+import com.android.bedstead.harrier.BedsteadJUnit4;
+import com.android.bedstead.harrier.DeviceState;
+import com.android.bedstead.harrier.annotations.EnsureHasNoWorkProfile;
+import com.android.bedstead.harrier.annotations.EnsureHasSecondaryUser;
+import com.android.bedstead.harrier.annotations.enterprise.EnsureHasDeviceOwner;
+import com.android.bedstead.harrier.annotations.enterprise.EnsureHasNoDeviceOwner;
+import com.android.bedstead.harrier.annotations.enterprise.EnsureHasNoProfileOwner;
 import com.android.bedstead.nene.TestApis;
 import com.android.bedstead.nene.exceptions.NeneException;
 import com.android.bedstead.nene.users.UserReference;
 import com.android.bedstead.nene.users.UserType;
+import com.android.bedstead.nene.utils.Versions;
 import com.android.bedstead.testapp.TestApp;
 import com.android.bedstead.testapp.TestAppProvider;
 import com.android.eventlib.premade.EventLibDeviceAdminReceiver;
 
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
+import org.junit.ClassRule;
 import org.junit.Ignore;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
-@RunWith(JUnit4.class)
+@RunWith(BedsteadJUnit4.class)
 public class DevicePolicyTest {
 
+    @ClassRule @Rule
+    public static final DeviceState sDeviceState = new DeviceState();
+
     //  TODO(180478924): We shouldn't need to hardcode this
     private static final String DEVICE_ADMIN_TESTAPP_PACKAGE_NAME = "android.DeviceAdminTestApp";
     private static final ComponentName DPC_COMPONENT_NAME =
@@ -70,6 +85,8 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoWorkProfile
     public void setProfileOwner_profileOwnerIsSet() {
         UserReference profile = sTestApis.users().createUser()
                 .parent(sUser)
@@ -88,6 +105,8 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoWorkProfile
     public void setProfileOwner_profileOwnerIsAlreadySet_throwsException() {
         UserReference profile = sTestApis.users().createUser()
                 .parent(sUser)
@@ -106,6 +125,8 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoWorkProfile
     public void setProfileOwner_componentNameNotInstalled_throwsException() {
         UserReference profile = sTestApis.users().createUser()
                 .parent(sUser)
@@ -120,12 +141,16 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
     public void setProfileOwner_componentNameIsNotDPC_throwsException() {
         assertThrows(NeneException.class,
                 () -> sTestApis.devicePolicy().setProfileOwner(sUser, NOT_DPC_COMPONENT_NAME));
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
     public void setProfileOwner_nullUser_throwsException() {
         assertThrows(NullPointerException.class,
                 () -> sTestApis.devicePolicy().setProfileOwner(
@@ -133,6 +158,8 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
     public void setProfileOwner_nullComponentName_throwsException() {
         assertThrows(NullPointerException.class,
                 () -> sTestApis.devicePolicy().setProfileOwner(
@@ -140,6 +167,8 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
     public void setProfileOwner_userDoesNotExist_throwsException() {
         assertThrows(NeneException.class,
                 () -> sTestApis.devicePolicy().setProfileOwner(
@@ -147,6 +176,8 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoWorkProfile
     public void getProfileOwner_returnsProfileOwner() {
         UserReference profile = sTestApis.users().createUser()
                 .parent(sUser)
@@ -165,6 +196,8 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoWorkProfile
     public void getProfileOwner_noProfileOwner_returnsNull() {
         UserReference profile = sTestApis.users().createUser()
                 .parent(sUser)
@@ -186,6 +219,8 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
     public void setDeviceOwner_deviceOwnerIsSet() {
         DeviceOwner deviceOwner =
                 sTestApis.devicePolicy().setDeviceOwner(sUser, DPC_COMPONENT_NAME);
@@ -198,19 +233,15 @@
     }
 
     @Test
+    @EnsureHasDeviceOwner
     public void setDeviceOwner_deviceOwnerIsAlreadySet_throwsException() {
-        DeviceOwner deviceOwner =
-                sTestApis.devicePolicy().setDeviceOwner(sUser, DPC_COMPONENT_NAME);
-
-        try {
-            assertThrows(NeneException.class,
-                    () -> sTestApis.devicePolicy().setDeviceOwner(sUser, DPC_COMPONENT_NAME));
-        } finally {
-            deviceOwner.remove();
-        }
+        assertThrows(NeneException.class,
+                () -> sTestApis.devicePolicy().setDeviceOwner(sUser, DPC_COMPONENT_NAME));
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
     public void setDeviceOwner_componentNameNotInstalled_throwsException() {
         sTestApp.reference().uninstall(sUser);
         try {
@@ -222,20 +253,40 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
     public void setDeviceOwner_componentNameIsNotDPC_throwsException() {
         assertThrows(NeneException.class,
                 () -> sTestApis.devicePolicy().setDeviceOwner(sUser, NOT_DPC_COMPONENT_NAME));
     }
 
     @Test
-    public void setDeviceOwner_userAlreadyOnDevice_throwsException() {
-        UserReference user = sTestApis.users().createUser().create();
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
+    @EnsureHasSecondaryUser
+    public void setDeviceOwner_preS_userAlreadyOnDevice_throwsException() {
+        assumeFalse("After S, device owner can be set with users on the device",
+                Versions.meetsMinimumSdkVersionRequirement(Build.VERSION_CODES.S));
+
+        assertThrows(NeneException.class,
+                () -> sTestApis.devicePolicy().setDeviceOwner(sUser, DPC_COMPONENT_NAME));
+    }
+
+    @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
+    @EnsureHasSecondaryUser
+    public void setDeviceOwner_sPlus_userAlreadyOnDevice_deviceOwnerIsSet() {
+        assumeTrue("After S, device owner can be set with users on the device",
+                Versions.meetsMinimumSdkVersionRequirement(Build.VERSION_CODES.S));
+
+        DeviceOwner deviceOwner =
+                sTestApis.devicePolicy().setDeviceOwner(sUser, DPC_COMPONENT_NAME);
 
         try {
-            assertThrows(NeneException.class,
-                    () -> sTestApis.devicePolicy().setDeviceOwner(sUser, DPC_COMPONENT_NAME));
+            assertThat(sTestApis.devicePolicy().getDeviceOwner()).isNotNull();
         } finally {
-            user.remove();
+            deviceOwner.remove();
         }
     }
 
@@ -245,6 +296,8 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
     public void setDeviceOwner_nullUser_throwsException() {
         assertThrows(NullPointerException.class,
                 () -> sTestApis.devicePolicy().setDeviceOwner(
@@ -252,6 +305,8 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
     public void setDeviceOwner_nullComponentName_throwsException() {
         assertThrows(NullPointerException.class,
                 () -> sTestApis.devicePolicy().setDeviceOwner(
@@ -259,6 +314,8 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
     public void setDeviceOwner_userDoesNotExist_throwsException() {
         assertThrows(NeneException.class,
                 () -> sTestApis.devicePolicy().setDeviceOwner(
@@ -266,6 +323,8 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
     public void getDeviceOwner_returnsDeviceOwner() {
         DeviceOwner deviceOwner =
                 sTestApis.devicePolicy().setDeviceOwner(sUser, DPC_COMPONENT_NAME);
@@ -278,13 +337,14 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
     public void getDeviceOwner_noDeviceOwner_returnsNull() {
-        // We must assume no device owner entering the test
-        // TODO(scottjonathan): Encode this assumption in the annotations when Harrier supports
         assertThat(sTestApis.devicePolicy().getDeviceOwner()).isNull();
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
     public void profileOwner_autoclose_removesProfileOwner() {
         try (ProfileOwner profileOwner =
                      sTestApis.devicePolicy().setProfileOwner(sUser, DPC_COMPONENT_NAME)) {
@@ -295,6 +355,8 @@
     }
 
     @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
     public void deviceOwner_autoclose_removesDeviceOwner() {
         try (DeviceOwner deviceOwner =
                      sTestApis.devicePolicy().setDeviceOwner(sUser, DPC_COMPONENT_NAME)) {
@@ -303,4 +365,48 @@
 
         assertThat(sTestApis.devicePolicy().getDeviceOwner()).isNull();
     }
+
+    @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
+    public void setDeviceOwner_recentlyUnsetProfileOwner_sets() {
+        sTestApis.devicePolicy().setProfileOwner(sUser, DPC_COMPONENT_NAME).remove();
+
+        sTestApis.devicePolicy().setDeviceOwner(sUser, DPC_COMPONENT_NAME);
+
+        assertThat(sTestApis.devicePolicy().getDeviceOwner()).isNotNull();
+    }
+
+    @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
+    public void setDeviceOwner_recentlyUnsetDeviceOwner_sets() {
+        sTestApis.devicePolicy().setDeviceOwner(sUser, DPC_COMPONENT_NAME).remove();
+
+        sTestApis.devicePolicy().setDeviceOwner(sUser, DPC_COMPONENT_NAME);
+
+        assertThat(sTestApis.devicePolicy().getDeviceOwner()).isNotNull();
+    }
+
+    @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
+    public void setProfileOwner_recentlyUnsetProfileOwner_sets() {
+        sTestApis.devicePolicy().setProfileOwner(sUser, DPC_COMPONENT_NAME).remove();
+
+        sTestApis.devicePolicy().setProfileOwner(sUser, DPC_COMPONENT_NAME);
+
+        assertThat(sTestApis.devicePolicy().getProfileOwner(sUser)).isNotNull();
+    }
+
+    @Test
+    @EnsureHasNoDeviceOwner
+    @EnsureHasNoProfileOwner
+    public void setProfileOwner_recentlyUnsetDeviceOwner_sets() {
+        sTestApis.devicePolicy().setDeviceOwner(sUser, DPC_COMPONENT_NAME).remove();
+
+        sTestApis.devicePolicy().setProfileOwner(sUser, DPC_COMPONENT_NAME);
+
+        assertThat(sTestApis.devicePolicy().getProfileOwner(sUser)).isNotNull();
+    }
 }