move locktask from DeviceOwner to DeviceAndProfileOwner

Change-Id: Ifb445002531600bd87b0fe60a90d549e1c7069ef
Merged-In: If0a9aa992712c50410a926eb22fcc90dd2ea81c3
Bug: 133249781
Test: atest "com.android.cts.devicepolicy.MixedProfileOwnerTest#testLockTask"
Test: atest "com.android.cts.devicepolicy.MixedProfileOwnerTest#testLockTaskAfterReboot"
Test: atest "com.android.cts.devicepolicy.MixedProfileOwnerTest#testLockTaskAfterReboot_tryOpeningSettings"
Test: atest "com.android.cts.devicepolicy.MixedDeviceOwnerTest#testLockTask_unaffiliatedUser"
Test: atest "com.android.cts.devicepolicy.MixedDeviceOwnerTest#testLockTask_affiliatedSecondaryUser"
diff --git a/hostsidetests/devicepolicy/app/DeviceAndProfileOwner/latest/AndroidManifest.xml b/hostsidetests/devicepolicy/app/DeviceAndProfileOwner/latest/AndroidManifest.xml
index c8385d7..fedd10e 100644
--- a/hostsidetests/devicepolicy/app/DeviceAndProfileOwner/latest/AndroidManifest.xml
+++ b/hostsidetests/devicepolicy/app/DeviceAndProfileOwner/latest/AndroidManifest.xml
@@ -85,6 +85,20 @@
         <activity
             android:name="com.android.cts.deviceandprofileowner.KeyManagementActivity"
             android:theme="@android:style/Theme.Translucent.NoTitleBar" />
+
+        <activity
+            android:name="com.android.cts.deviceandprofileowner.LockTaskUtilityActivity"/>
+        <activity
+            android:name="com.android.cts.deviceandprofileowner.LockTaskUtilityActivityIfWhitelisted"
+            android:launchMode="singleInstance"
+            android:lockTaskMode="if_whitelisted">
+            <intent-filter>
+                <action android:name="android.intent.action.MAIN"/>
+                <category android:name="android.intent.category.HOME"/>
+                <category android:name="android.intent.category.DEFAULT"/>
+            </intent-filter>
+        </activity>
+
     </application>
 
     <instrumentation
diff --git a/hostsidetests/devicepolicy/app/DeviceAndProfileOwner/src/com/android/cts/deviceandprofileowner/AffiliationTest.java b/hostsidetests/devicepolicy/app/DeviceAndProfileOwner/src/com/android/cts/deviceandprofileowner/AffiliationTest.java
new file mode 100644
index 0000000..5b8f4fb
--- /dev/null
+++ b/hostsidetests/devicepolicy/app/DeviceAndProfileOwner/src/com/android/cts/deviceandprofileowner/AffiliationTest.java
@@ -0,0 +1,114 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.cts.deviceandprofileowner;
+
+import static com.android.cts.deviceandprofileowner.BaseDeviceAdminTest.ADMIN_RECEIVER_COMPONENT;
+
+import static junit.framework.Assert.assertEquals;
+import static junit.framework.Assert.assertFalse;
+import static junit.framework.Assert.assertTrue;
+import static junit.framework.Assert.fail;
+
+import static org.junit.Assert.assertArrayEquals;
+
+import android.app.admin.DevicePolicyManager;
+import android.content.Context;
+
+import androidx.test.InstrumentationRegistry;
+import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.Collections;
+import java.util.Set;
+
+@RunWith(AndroidJUnit4.class)
+public class AffiliationTest {
+
+    private DevicePolicyManager mDevicePolicyManager;
+
+    @Before
+    public void setUp() {
+        Context context = InstrumentationRegistry.getContext();
+        mDevicePolicyManager = (DevicePolicyManager)
+                context.getSystemService(Context.DEVICE_POLICY_SERVICE);
+    }
+
+    @Test
+    public void testSetAffiliationId1() {
+        setAffiliationIds(Collections.singleton("id.number.1"));
+    }
+
+    @Test
+    public void testSetAffiliationId2() {
+        setAffiliationIds(Collections.singleton("id.number.2"));
+    }
+
+    @Test
+    public void testLockTaskMethodsThrowExceptionIfUnaffiliated() {
+        checkLockTaskMethodsThrow();
+    }
+
+    /** Assumes that the calling user is already affiliated before calling this method */
+    @Test
+    public void testSetLockTaskPackagesClearedIfUserBecomesUnaffiliated() {
+        final String[] packages = {"package1", "package2"};
+        mDevicePolicyManager.setLockTaskPackages(ADMIN_RECEIVER_COMPONENT, packages);
+        assertArrayEquals(packages,
+                mDevicePolicyManager.getLockTaskPackages(ADMIN_RECEIVER_COMPONENT));
+        assertTrue(mDevicePolicyManager.isLockTaskPermitted("package1"));
+        assertFalse(mDevicePolicyManager.isLockTaskPermitted("package3"));
+
+        final Set<String> previousAffiliationIds =
+                mDevicePolicyManager.getAffiliationIds(ADMIN_RECEIVER_COMPONENT);
+        try {
+            // Clearing affiliation ids for this user. Lock task methods unavailable.
+            setAffiliationIds(Collections.emptySet());
+            checkLockTaskMethodsThrow();
+            assertFalse(mDevicePolicyManager.isLockTaskPermitted("package1"));
+
+            // Affiliating the user again. Previously set packages have been cleared.
+            setAffiliationIds(previousAffiliationIds);
+            assertEquals(0,
+                    mDevicePolicyManager.getLockTaskPackages(ADMIN_RECEIVER_COMPONENT).length);
+            assertFalse(mDevicePolicyManager.isLockTaskPermitted("package1"));
+        } finally {
+            mDevicePolicyManager.setAffiliationIds(ADMIN_RECEIVER_COMPONENT,
+                    previousAffiliationIds);
+        }
+    }
+
+    private void setAffiliationIds(Set<String> ids) {
+        mDevicePolicyManager.setAffiliationIds(ADMIN_RECEIVER_COMPONENT, ids);
+        assertEquals(ids, mDevicePolicyManager.getAffiliationIds(ADMIN_RECEIVER_COMPONENT));
+    }
+
+    private void checkLockTaskMethodsThrow() {
+        try {
+            mDevicePolicyManager.setLockTaskPackages(ADMIN_RECEIVER_COMPONENT, new String[0]);
+            fail("setLockTaskPackages did not throw expected SecurityException");
+        } catch (SecurityException expected) {
+        }
+        try {
+            mDevicePolicyManager.getLockTaskPackages(ADMIN_RECEIVER_COMPONENT);
+            fail("getLockTaskPackages did not throw expected SecurityException");
+        } catch (SecurityException expected) {
+        }
+    }
+}
diff --git a/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/LockTaskHostDrivenTest.java b/hostsidetests/devicepolicy/app/DeviceAndProfileOwner/src/com/android/cts/deviceandprofileowner/LockTaskHostDrivenTest.java
similarity index 90%
rename from hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/LockTaskHostDrivenTest.java
rename to hostsidetests/devicepolicy/app/DeviceAndProfileOwner/src/com/android/cts/deviceandprofileowner/LockTaskHostDrivenTest.java
index 99b97bd..c2e2121 100644
--- a/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/LockTaskHostDrivenTest.java
+++ b/hostsidetests/devicepolicy/app/DeviceAndProfileOwner/src/com/android/cts/deviceandprofileowner/LockTaskHostDrivenTest.java
@@ -13,10 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.cts.deviceowner;
-
-import static junit.framework.Assert.assertEquals;
-import static junit.framework.Assert.assertTrue;
+package com.android.cts.deviceandprofileowner;
 
 import android.app.ActivityManager;
 import android.app.admin.DevicePolicyManager;
@@ -44,7 +41,7 @@
  * state after running.
  */
 @RunWith(AndroidJUnit4.class)
-public class LockTaskHostDrivenTest {
+public class LockTaskHostDrivenTest extends BaseDeviceAdminTest {
 
     private static final String TAG = LockTaskHostDrivenTest.class.getName();
 
@@ -96,7 +93,7 @@
     public void tearDown() {
         mContext.unregisterReceiver(mReceiver);
     }
-  
+
     @Test
     public void startLockTask() throws Exception {
         Log.d(TAG, "startLockTask on host-driven test (no cleanup)");
@@ -114,7 +111,7 @@
         mUiDevice.waitForIdle();
 
         // We need to wait until the LockTaskActivity is ready
-        // since com.android.cts.deviceowner can be killed by AMS for reason "start instr".
+        // since the package can be killed by AMS for reason "start instr".
         synchronized (mActivityResumedLock) {
             if (!mIsActivityResumed) {
                 mActivityResumedLock.wait(ACTIVITY_RESUMED_TIMEOUT_MILLIS);
@@ -140,10 +137,9 @@
     @Test
     public void clearDefaultHomeIntentReceiver() {
         mDevicePolicyManager.clearPackagePersistentPreferredActivities(
-                BasicAdminReceiver.getComponentName(mContext),
+                ADMIN_RECEIVER_COMPONENT,
                 mContext.getPackageName());
-        mDevicePolicyManager.setLockTaskPackages(BasicAdminReceiver.getComponentName(mContext),
-                new String[0]);
+        mDevicePolicyManager.setLockTaskPackages(ADMIN_RECEIVER_COMPONENT, new String[0]);
     }
 
     private void checkLockedActivityIsRunning() throws Exception {
@@ -165,13 +161,13 @@
     }
 
     private void setDefaultHomeIntentReceiver() {
-        mDevicePolicyManager.setLockTaskPackages(BasicAdminReceiver.getComponentName(mContext),
+        mDevicePolicyManager.setLockTaskPackages(ADMIN_RECEIVER_COMPONENT,
                 new String[]{mContext.getPackageName()});
         IntentFilter intentFilter = new IntentFilter(Intent.ACTION_MAIN);
         intentFilter.addCategory(Intent.CATEGORY_HOME);
         intentFilter.addCategory(Intent.CATEGORY_DEFAULT);
         mDevicePolicyManager.addPersistentPreferredActivity(
-                BasicAdminReceiver.getComponentName(mContext), intentFilter,
+                ADMIN_RECEIVER_COMPONENT, intentFilter,
                 new ComponentName(mContext.getPackageName(), LOCK_TASK_ACTIVITY));
     }
 }
diff --git a/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/LockTaskTest.java b/hostsidetests/devicepolicy/app/DeviceAndProfileOwner/src/com/android/cts/deviceandprofileowner/LockTaskTest.java
similarity index 97%
rename from hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/LockTaskTest.java
rename to hostsidetests/devicepolicy/app/DeviceAndProfileOwner/src/com/android/cts/deviceandprofileowner/LockTaskTest.java
index 6cfec10..37f1be3 100644
--- a/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/LockTaskTest.java
+++ b/hostsidetests/devicepolicy/app/DeviceAndProfileOwner/src/com/android/cts/deviceandprofileowner/LockTaskTest.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.cts.deviceowner;
+package com.android.cts.deviceandprofileowner;
 
 import static android.app.admin.DevicePolicyManager.LOCK_TASK_FEATURE_GLOBAL_ACTIONS;
 import static android.app.admin.DevicePolicyManager.LOCK_TASK_FEATURE_HOME;
@@ -23,10 +23,6 @@
 import static android.app.admin.DevicePolicyManager.LOCK_TASK_FEATURE_OVERVIEW;
 import static android.app.admin.DevicePolicyManager.LOCK_TASK_FEATURE_SYSTEM_INFO;
 
-import static junit.framework.Assert.assertEquals;
-import static junit.framework.Assert.assertFalse;
-import static junit.framework.Assert.assertTrue;
-
 import static org.junit.Assert.assertArrayEquals;
 import static org.testng.Assert.assertThrows;
 
@@ -52,19 +48,18 @@
 import java.util.concurrent.TimeUnit;
 
 @RunWith(AndroidJUnit4.class)
-public class LockTaskTest {
+public class LockTaskTest extends BaseDeviceAdminTest {
 
     private static final String TAG = "LockTaskTest";
 
     private static final String PACKAGE_NAME = LockTaskTest.class.getPackage().getName();
-    private static final ComponentName ADMIN_COMPONENT =
-            new ComponentName(PACKAGE_NAME, BasicAdminReceiver.class.getName());
+    private static final ComponentName ADMIN_COMPONENT = ADMIN_RECEIVER_COMPONENT;
     private static final String TEST_PACKAGE = "com.google.android.example.somepackage";
 
     private static final String UTILITY_ACTIVITY
-            = "com.android.cts.deviceowner.LockTaskUtilityActivity";
+            = "com.android.cts.deviceandprofileowner.LockTaskUtilityActivity";
     private static final String UTILITY_ACTIVITY_IF_WHITELISTED
-            = "com.android.cts.deviceowner.LockTaskUtilityActivityIfWhitelisted";
+            = "com.android.cts.deviceandprofileowner.LockTaskUtilityActivityIfWhitelisted";
 
     private static final String RECEIVER_ACTIVITY_PACKAGE_NAME =
             "com.android.cts.intent.receiver";
diff --git a/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/LockTaskUtilityActivity.java b/hostsidetests/devicepolicy/app/DeviceAndProfileOwner/src/com/android/cts/deviceandprofileowner/LockTaskUtilityActivity.java
similarity index 91%
rename from hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/LockTaskUtilityActivity.java
rename to hostsidetests/devicepolicy/app/DeviceAndProfileOwner/src/com/android/cts/deviceandprofileowner/LockTaskUtilityActivity.java
index aeb8abd..be3ec07 100644
--- a/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/LockTaskUtilityActivity.java
+++ b/hostsidetests/devicepolicy/app/DeviceAndProfileOwner/src/com/android/cts/deviceandprofileowner/LockTaskUtilityActivity.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.cts.deviceowner;
+package com.android.cts.deviceandprofileowner;
 
 import android.app.Activity;
 import android.content.Intent;
@@ -28,11 +28,11 @@
     public static final String START_ACTIVITY = "startActivity";
     public static final String FINISH = "finish";
 
-    public static final String CREATE_ACTION = "com.android.cts.deviceowner.LOCK_TASK_CREATE";
-    public static final String DESTROY_ACTION = "com.android.cts.deviceowner.LOCK_TASK_DESTROY";
-    public static final String PAUSE_ACTION = "com.android.cts.deviceowner.LOCK_TASK_PAUSE";
-    public static final String RESUME_ACTION = "com.android.cts.deviceowner.LOCK_TASK_RESUME";
-    public static final String INTENT_ACTION = "com.android.cts.deviceowner.LOCK_TASK_INTENT";
+    public static final String CREATE_ACTION = "com.android.cts.deviceandprofileowner.LOCK_TASK_CREATE";
+    public static final String DESTROY_ACTION = "com.android.cts.deviceandprofileowner.LOCK_TASK_DESTROY";
+    public static final String PAUSE_ACTION = "com.android.cts.deviceandprofileowner.LOCK_TASK_PAUSE";
+    public static final String RESUME_ACTION = "com.android.cts.deviceandprofileowner.LOCK_TASK_RESUME";
+    public static final String INTENT_ACTION = "com.android.cts.deviceandprofileowner.LOCK_TASK_INTENT";
 
     @Override
     protected void onNewIntent(Intent intent) {
diff --git a/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/LockTaskUtilityActivityIfWhitelisted.java b/hostsidetests/devicepolicy/app/DeviceAndProfileOwner/src/com/android/cts/deviceandprofileowner/LockTaskUtilityActivityIfWhitelisted.java
similarity index 93%
rename from hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/LockTaskUtilityActivityIfWhitelisted.java
rename to hostsidetests/devicepolicy/app/DeviceAndProfileOwner/src/com/android/cts/deviceandprofileowner/LockTaskUtilityActivityIfWhitelisted.java
index 4cf6efa..d030fdd 100644
--- a/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/LockTaskUtilityActivityIfWhitelisted.java
+++ b/hostsidetests/devicepolicy/app/DeviceAndProfileOwner/src/com/android/cts/deviceandprofileowner/LockTaskUtilityActivityIfWhitelisted.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.cts.deviceowner;
+package com.android.cts.deviceandprofileowner;
 
 public class LockTaskUtilityActivityIfWhitelisted extends LockTaskUtilityActivity {
 }
diff --git a/hostsidetests/devicepolicy/app/DeviceOwner/AndroidManifest.xml b/hostsidetests/devicepolicy/app/DeviceOwner/AndroidManifest.xml
index ebb848b..1b5bdac 100644
--- a/hostsidetests/devicepolicy/app/DeviceOwner/AndroidManifest.xml
+++ b/hostsidetests/devicepolicy/app/DeviceOwner/AndroidManifest.xml
@@ -72,19 +72,6 @@
         </service>
 
         <activity
-            android:name="com.android.cts.deviceowner.LockTaskUtilityActivity" />
-        <activity
-            android:name="com.android.cts.deviceowner.LockTaskUtilityActivityIfWhitelisted"
-            android:launchMode="singleInstance"
-            android:lockTaskMode="if_whitelisted">
-            <intent-filter>
-                <action android:name="android.intent.action.MAIN"/>
-                <category android:name="android.intent.category.HOME"/>
-                <category android:name="android.intent.category.DEFAULT"/>
-            </intent-filter>
-        </activity>
-
-        <activity
             android:name=".SetPolicyActivity"
             android:launchMode="singleTop">
             <intent-filter>
diff --git a/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/AffiliationTest.java b/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/AffiliationTest.java
index 57d200b..2424c46 100644
--- a/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/AffiliationTest.java
+++ b/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/AffiliationTest.java
@@ -17,12 +17,8 @@
 package com.android.cts.deviceowner;
 
 import static junit.framework.Assert.assertEquals;
-import static junit.framework.Assert.assertFalse;
-import static junit.framework.Assert.assertTrue;
 import static junit.framework.Assert.fail;
 
-import static org.junit.Assert.assertArrayEquals;
-
 import android.app.admin.DevicePolicyManager;
 import android.content.ComponentName;
 import android.content.Context;
@@ -81,52 +77,8 @@
         setAffiliationIds(Collections.singleton("id.number.2"));
     }
 
-    @Test
-    public void testLockTaskMethodsThrowExceptionIfUnaffiliated() {
-        checkLockTaskMethodsThrow();
-    }
-
-    /** Assumes that the calling user is already affiliated before calling this method */
-    @Test
-    public void testSetLockTaskPackagesClearedIfUserBecomesUnaffiliated() {
-        final String[] packages = {"package1", "package2"};
-        mDevicePolicyManager.setLockTaskPackages(mAdminComponent, packages);
-        assertArrayEquals(packages, mDevicePolicyManager.getLockTaskPackages(mAdminComponent));
-        assertTrue(mDevicePolicyManager.isLockTaskPermitted("package1"));
-        assertFalse(mDevicePolicyManager.isLockTaskPermitted("package3"));
-
-        final Set<String> previousAffiliationIds =
-                mDevicePolicyManager.getAffiliationIds(mAdminComponent);
-        try {
-            // Clearing affiliation ids for this user. Lock task methods unavailable.
-            setAffiliationIds(Collections.emptySet());
-            checkLockTaskMethodsThrow();
-            assertFalse(mDevicePolicyManager.isLockTaskPermitted("package1"));
-
-            // Affiliating the user again. Previously set packages have been cleared.
-            setAffiliationIds(previousAffiliationIds);
-            assertEquals(0, mDevicePolicyManager.getLockTaskPackages(mAdminComponent).length);
-            assertFalse(mDevicePolicyManager.isLockTaskPermitted("package1"));
-        } finally {
-            mDevicePolicyManager.setAffiliationIds(mAdminComponent, previousAffiliationIds);
-        }
-    }
-
     private void setAffiliationIds(Set<String> ids) {
         mDevicePolicyManager.setAffiliationIds(mAdminComponent, ids);
         assertEquals(ids, mDevicePolicyManager.getAffiliationIds(mAdminComponent));
     }
-
-    private void checkLockTaskMethodsThrow() {
-        try {
-            mDevicePolicyManager.setLockTaskPackages(mAdminComponent, new String[0]);
-            fail("setLockTaskPackages did not throw expected SecurityException");
-        } catch (SecurityException expected) {
-        }
-        try {
-            mDevicePolicyManager.getLockTaskPackages(mAdminComponent);
-            fail("getLockTaskPackages did not throw expected SecurityException");
-        } catch (SecurityException expected) {
-        }
-    }
 }
diff --git a/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/BaseDevicePolicyTest.java b/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/BaseDevicePolicyTest.java
index 3b0cbfb..0a32aba 100644
--- a/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/BaseDevicePolicyTest.java
+++ b/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/BaseDevicePolicyTest.java
@@ -185,6 +185,7 @@
         }
 
         removeOwners();
+        switchUser(USER_SYSTEM);
         removeTestUsers();
         // Unlock keyguard before test
         wakeupAndDismissKeyguard();
@@ -199,6 +200,7 @@
         getDevice().executeShellCommand("settings put global package_verifier_enable "
                 + mPackageVerifier);
         removeOwners();
+        switchUser(USER_SYSTEM);
         removeTestUsers();
         removeTestPackages();
         super.tearDown();
diff --git a/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceAndProfileOwnerTest.java b/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceAndProfileOwnerTest.java
index 56014ea..d9d9b15 100644
--- a/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceAndProfileOwnerTest.java
+++ b/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceAndProfileOwnerTest.java
@@ -50,14 +50,16 @@
 
     public static final String DEVICE_ADMIN_PKG = "com.android.cts.deviceandprofileowner";
     public static final String DEVICE_ADMIN_APK = "CtsDeviceAndProfileOwnerApp.apk";
-    public static final String ADMIN_RECEIVER_TEST_CLASS
+    protected static final String ADMIN_RECEIVER_TEST_CLASS
             = ".BaseDeviceAdminTest$BasicAdminReceiver";
+    protected static final String DEVICE_ADMIN_COMPONENT_FLATTENED =
+            DEVICE_ADMIN_PKG + "/" + ADMIN_RECEIVER_TEST_CLASS;
 
     private static final String STORAGE_ENCRYPTION_TEST_CLASS = ".StorageEncryptionTest";
     private static final String IS_PRIMARY_USER_PARAM = "isPrimaryUser";
 
-    private static final String INTENT_RECEIVER_PKG = "com.android.cts.intent.receiver";
-    private static final String INTENT_RECEIVER_APK = "CtsIntentReceiverApp.apk";
+    protected static final String INTENT_RECEIVER_PKG = "com.android.cts.intent.receiver";
+    protected static final String INTENT_RECEIVER_APK = "CtsIntentReceiverApp.apk";
 
     private static final String INTENT_SENDER_PKG = "com.android.cts.intent.sender";
     private static final String INTENT_SENDER_APK = "CtsIntentSenderApp.apk";
@@ -1056,6 +1058,77 @@
                     .build());
     }
 
+    public void testLockTask() throws Exception {
+        if (!mHasFeature) {
+            return;
+        }
+        try {
+            installAppAsUser(INTENT_RECEIVER_APK, mUserId);
+            executeDeviceTestClass(".LockTaskTest");
+            assertMetricsLogged(
+                    getDevice(),
+                    () -> executeDeviceTestMethod(".LockTaskTest", "testStartLockTask"),
+                    new DevicePolicyEventWrapper.Builder(EventId.SET_LOCKTASK_MODE_ENABLED_VALUE)
+                            .setAdminPackageName(DEVICE_ADMIN_PKG)
+                            .setBoolean(true)
+                            .setStrings(DEVICE_ADMIN_PKG)
+                            .build());
+        } catch (AssertionError ex) {
+            // STOPSHIP(b/32771855), remove this once we fixed the bug.
+            executeShellCommand("dumpsys activity activities");
+            executeShellCommand("dumpsys window -a");
+            executeShellCommand("dumpsys activity service com.android.systemui");
+            throw ex;
+        } finally {
+            getDevice().uninstallPackage(INTENT_RECEIVER_PKG);
+        }
+    }
+
+    public void testLockTaskAfterReboot() throws Exception {
+        if (!mHasFeature) {
+            return;
+        }
+
+        try {
+            // Just start kiosk mode
+            executeDeviceTestMethod(".LockTaskHostDrivenTest", "startLockTask");
+
+            // Reboot while in kiosk mode and then unlock the device
+            rebootAndWaitUntilReady();
+
+            // Check that kiosk mode is working and can't be interrupted
+            executeDeviceTestMethod(".LockTaskHostDrivenTest",
+                    "testLockTaskIsActiveAndCantBeInterrupted");
+        } finally {
+            executeDeviceTestMethod(".LockTaskHostDrivenTest",
+                    "clearDefaultHomeIntentReceiver");
+        }
+    }
+
+    public void testLockTaskAfterReboot_tryOpeningSettings() throws Exception {
+        if (!mHasFeature) {
+            return;
+        }
+
+        try {
+            // Just start kiosk mode
+            executeDeviceTestMethod(".LockTaskHostDrivenTest", "startLockTask");
+
+            // Reboot while in kiosk mode and then unlock the device
+            rebootAndWaitUntilReady();
+
+            // Try to open settings via adb
+            executeShellCommand("am start -a android.settings.SETTINGS");
+
+            // Check again
+            executeDeviceTestMethod(".LockTaskHostDrivenTest",
+                    "testLockTaskIsActiveAndCantBeInterrupted");
+        } finally {
+            executeDeviceTestMethod(".LockTaskHostDrivenTest",
+                    "clearDefaultHomeIntentReceiver");
+        }
+    }
+
     public void testSuspendPackage() throws Exception {
         if (!mHasFeature) {
             return;
diff --git a/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceOwnerTest.java b/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceOwnerTest.java
index a42799e..b88fed9 100644
--- a/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceOwnerTest.java
+++ b/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceOwnerTest.java
@@ -18,6 +18,8 @@
 
 import static com.android.cts.devicepolicy.metrics.DevicePolicyEventLogVerifier.assertMetricsLogged;
 
+import android.stats.devicepolicy.EventId;
+
 import com.android.compatibility.common.tradefed.build.CompatibilityBuildHelper;
 import com.android.cts.devicepolicy.metrics.DevicePolicyEventWrapper;
 import com.android.tradefed.device.DeviceNotAvailableException;
@@ -33,8 +35,6 @@
 import java.util.List;
 import java.util.Map;
 
-import android.stats.devicepolicy.EventId;
-
 /**
  * Set of tests for Device Owner use cases.
  */
@@ -568,113 +568,6 @@
         executeDeviceTestMethod(".AffiliationTest", "testSetAffiliationId_containsEmptyString");
     }
 
-    public void testLockTask_deviceOwnerUser() throws Exception {
-        if (!mHasFeature) {
-            return;
-        }
-        try {
-            installAppAsUser(INTENT_RECEIVER_APK, mPrimaryUserId);
-            executeDeviceOwnerTest("LockTaskTest");
-            assertMetricsLogged(getDevice(), () -> {
-                runDeviceTestsAsUser(DEVICE_OWNER_PKG, ".LockTaskTest", "testStartLockTask",
-                        mPrimaryUserId);
-            }, new DevicePolicyEventWrapper.Builder(EventId.SET_LOCKTASK_MODE_ENABLED_VALUE)
-                    .setAdminPackageName(DEVICE_OWNER_PKG)
-                    .setBoolean(true)
-                    .setStrings(DEVICE_OWNER_PKG)
-                    .build());
-        } catch (AssertionError ex) {
-            // STOPSHIP(b/32771855), remove this once we fixed the bug.
-            executeShellCommand("dumpsys activity activities");
-            executeShellCommand("dumpsys window -a");
-            executeShellCommand("dumpsys activity service com.android.systemui");
-            throw ex;
-        } finally {
-            getDevice().uninstallPackage(INTENT_RECEIVER_PKG);
-        }
-    }
-
-    public void testLockTaskAfterReboot_deviceOwnerUser() throws Exception {
-        if (!mHasFeature) {
-            return;
-        }
-
-        try {
-            // Just start kiosk mode
-            runDeviceTestsAsUser(DEVICE_OWNER_PKG, ".LockTaskHostDrivenTest", "startLockTask",
-                    mPrimaryUserId);
-
-            // Reboot while in kiosk mode and then unlock the device
-            rebootAndWaitUntilReady();
-
-            // Check that kiosk mode is working and can't be interrupted
-            runDeviceTestsAsUser(DEVICE_OWNER_PKG, ".LockTaskHostDrivenTest",
-                    "testLockTaskIsActiveAndCantBeInterrupted", mPrimaryUserId);
-        } finally {
-            runDeviceTestsAsUser(DEVICE_OWNER_PKG, ".LockTaskHostDrivenTest",
-                    "clearDefaultHomeIntentReceiver", mPrimaryUserId);
-        }
-    }
-
-    public void testLockTaskAfterReboot_tryOpeningSettings_deviceOwnerUser() throws Exception {
-        if (!mHasFeature) {
-            return;
-        }
-
-        try {
-            // Just start kiosk mode
-            runDeviceTestsAsUser(DEVICE_OWNER_PKG, ".LockTaskHostDrivenTest", "startLockTask",
-                    mPrimaryUserId);
-
-            // Reboot while in kiosk mode and then unlock the device
-            rebootAndWaitUntilReady();
-
-            // Try to open settings via adb
-            executeShellCommand("am start -a android.settings.SETTINGS");
-
-            // Check again
-            runDeviceTestsAsUser(DEVICE_OWNER_PKG, ".LockTaskHostDrivenTest",
-                    "testLockTaskIsActiveAndCantBeInterrupted", mPrimaryUserId);
-        } finally {
-            runDeviceTestsAsUser(DEVICE_OWNER_PKG, ".LockTaskHostDrivenTest",
-                    "clearDefaultHomeIntentReceiver", mPrimaryUserId);
-        }
-    }
-
-    public void testLockTask_unaffiliatedUser() throws Exception {
-        if (!mHasFeature || !canCreateAdditionalUsers(1)) {
-            return;
-        }
-
-        final int userId = createUser();
-        installAppAsUser(DEVICE_OWNER_APK, userId);
-        setProfileOwnerOrFail(DEVICE_OWNER_COMPONENT, userId);
-
-        runDeviceTestsAsUser(
-                DEVICE_OWNER_PKG,
-                ".AffiliationTest",
-                "testLockTaskMethodsThrowExceptionIfUnaffiliated",
-                userId);
-
-        runDeviceTestsAsUser(
-                DEVICE_OWNER_PKG, ".AffiliationTest", "testSetAffiliationId1", mPrimaryUserId);
-        runDeviceTestsAsUser(
-                DEVICE_OWNER_PKG, ".AffiliationTest", "testSetAffiliationId1", userId);
-        runDeviceTestsAsUser(
-                DEVICE_OWNER_PKG,
-                ".AffiliationTest",
-                "testSetLockTaskPackagesClearedIfUserBecomesUnaffiliated",
-                userId);
-    }
-
-    public void testLockTask_affiliatedSecondaryUser() throws Exception {
-        if (!mHasFeature || !canCreateAdditionalUsers(1)) {
-            return;
-        }
-        final int userId = createAffiliatedSecondaryUser();
-        executeAffiliatedProfileOwnerTest("LockTaskTest", userId);
-    }
-
     public void testSystemUpdatePolicy() throws Exception {
         if (!mHasFeature) {
             return;
@@ -1135,15 +1028,6 @@
         runDeviceTestsAsUser(DEVICE_OWNER_PKG, testClass, mPrimaryUserId);
     }
 
-    private void executeAffiliatedProfileOwnerTest(String testClassName, int userId)
-            throws Exception {
-        if (!mHasFeature) {
-            return;
-        }
-        String testClass = DEVICE_OWNER_PKG + "." + testClassName;
-        runDeviceTestsAsUser(DEVICE_OWNER_PKG, testClass, userId);
-    }
-
     private void executeDeviceTestMethod(String className, String testName) throws Exception {
         if (!mHasFeature) {
             return;
@@ -1162,7 +1046,7 @@
         waitForBroadcastIdle();
         wakeupAndDismissKeyguard();
 
-        // Setting the same affiliation ids on both users and running the lock task tests.
+        // Setting the same affiliation ids on both users
         runDeviceTestsAsUser(
                 DEVICE_OWNER_PKG, ".AffiliationTest", "testSetAffiliationId1", mPrimaryUserId);
         runDeviceTestsAsUser(
diff --git a/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/MixedDeviceOwnerTest.java b/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/MixedDeviceOwnerTest.java
index 1dc783e4..6f3c07c 100644
--- a/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/MixedDeviceOwnerTest.java
+++ b/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/MixedDeviceOwnerTest.java
@@ -16,16 +16,11 @@
 
 package com.android.cts.devicepolicy;
 
-import static com.android.cts.devicepolicy.metrics.DevicePolicyEventLogVerifier.assertMetricsLogged;
-
 import android.stats.devicepolicy.EventId;
 
 import com.android.cts.devicepolicy.metrics.DevicePolicyEventWrapper;
 
-import java.io.File;
 import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -46,10 +41,9 @@
             mUserId = mPrimaryUserId;
 
             installAppAsUser(DEVICE_ADMIN_APK, mUserId);
-            if (!setDeviceOwner(
-                    DEVICE_ADMIN_PKG + "/" + ADMIN_RECEIVER_TEST_CLASS, mUserId,
-                    /*expectFailure*/ false)) {
-                removeAdmin(DEVICE_ADMIN_PKG + "/" + ADMIN_RECEIVER_TEST_CLASS, mUserId);
+            if (!setDeviceOwner(DEVICE_ADMIN_COMPONENT_FLATTENED, mUserId, /*expectFailure*/
+                    false)) {
+                removeAdmin(DEVICE_ADMIN_COMPONENT_FLATTENED, mUserId);
                 getDevice().uninstallPackage(DEVICE_ADMIN_PKG);
                 fail("Failed to set device owner");
             }
@@ -60,11 +54,39 @@
     protected void tearDown() throws Exception {
         if (mHasFeature) {
             assertTrue("Failed to remove device owner",
-                    removeAdmin(DEVICE_ADMIN_PKG + "/" + ADMIN_RECEIVER_TEST_CLASS, mUserId));
+                    removeAdmin(DEVICE_ADMIN_COMPONENT_FLATTENED, mUserId));
         }
         super.tearDown();
     }
 
+    public void testLockTask_unaffiliatedUser() throws Exception {
+        if (!mHasFeature || !canCreateAdditionalUsers(1)) {
+            return;
+        }
+
+        final int userId = createSecondaryUserAsProfileOwner();
+        runDeviceTestsAsUser(
+                DEVICE_ADMIN_PKG, ".AffiliationTest",
+                "testLockTaskMethodsThrowExceptionIfUnaffiliated", userId);
+
+        setUserAsAffiliatedUserToPrimary(userId);
+        runDeviceTestsAsUser(
+                DEVICE_ADMIN_PKG,
+                ".AffiliationTest",
+                "testSetLockTaskPackagesClearedIfUserBecomesUnaffiliated",
+                userId);
+    }
+
+    public void testLockTask_affiliatedSecondaryUser() throws Exception {
+        if (!mHasFeature || !canCreateAdditionalUsers(1)) {
+            return;
+        }
+        final int userId = createSecondaryUserAsProfileOwner();
+        switchToUser(userId);
+        setUserAsAffiliatedUserToPrimary(userId);
+        runDeviceTestsAsUser(DEVICE_ADMIN_PKG, ".LockTaskTest", userId);
+    }
+
     public void testDelegatedCertInstallerDeviceIdAttestation() throws Exception {
         if (!mHasFeature) {
             return;
@@ -105,4 +127,26 @@
         result.add(DELEGATION_NETWORK_LOGGING);
         return result;
     }
+
+    private int createSecondaryUserAsProfileOwner() throws Exception {
+        final int userId = createUser();
+        installAppAsUser(INTENT_RECEIVER_APK, userId);
+        installAppAsUser(DEVICE_ADMIN_APK, userId);
+        setProfileOwnerOrFail(DEVICE_ADMIN_COMPONENT_FLATTENED, userId);
+        return userId;
+    }
+
+    private void switchToUser(int userId) throws Exception {
+        switchUser(userId);
+        waitForBroadcastIdle();
+        wakeupAndDismissKeyguard();
+    }
+
+    private void setUserAsAffiliatedUserToPrimary(int userId) throws Exception {
+        // Setting the same affiliation ids on both users
+        runDeviceTestsAsUser(
+                DEVICE_ADMIN_PKG, ".AffiliationTest", "testSetAffiliationId1", mPrimaryUserId);
+        runDeviceTestsAsUser(
+                DEVICE_ADMIN_PKG, ".AffiliationTest", "testSetAffiliationId1", userId);
+    }
 }