diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 8702578026082..cb78ad8345410 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -9290,22 +9290,39 @@ public class DevicePolicyManagerService extends BaseIDevicePolicyManager {
                     return poComponent;
                 }
             }
-            final String supervisor = mContext.getResources().getString(
-                    com.android.internal.R.string.config_defaultSupervisionProfileOwnerComponent);
-            if (supervisor == null) {
-                return null;
+
+            // Check profile owner first as that is what most likely is set.
+            if (isSupervisionComponent(poComponent)) {
+                return poComponent;
             }
-            final ComponentName supervisorComponent = ComponentName.unflattenFromString(supervisor);
-            if (supervisorComponent == null) {
-                return null;
+
+            if (isSupervisionComponent(doComponent)) {
+                return doComponent;
             }
-            if (supervisorComponent.equals(doComponent) || supervisorComponent.equals(
-                    poComponent)) {
-                return supervisorComponent;
-            } else {
-                return null;
+
+            return null;
+        }
+    }
+
+    private boolean isSupervisionComponent(@Nullable ComponentName who) {
+        if (who == null) {
+            return false;
+        }
+
+        final String configComponent = mContext.getResources().getString(
+                com.android.internal.R.string.config_defaultSupervisionProfileOwnerComponent);
+        if (configComponent != null) {
+            final ComponentName componentName = ComponentName.unflattenFromString(configComponent);
+            if (who.equals(componentName)) {
+                return true;
             }
         }
+
+        // Check the system supervision role.
+        final String configPackage = mContext.getResources().getString(
+                com.android.internal.R.string.config_systemSupervision);
+
+        return who.getPackageName().equals(configPackage);
     }
 
     @Override
@@ -9491,22 +9508,7 @@ public class DevicePolicyManagerService extends BaseIDevicePolicyManager {
                     "Cannot set the profile owner on a user which is already set-up");
 
             if (!mIsWatch) {
-                final String supervisionRolePackage = mContext.getResources().getString(
-                        com.android.internal.R.string.config_systemSupervision);
-                // Only the default supervision profile owner or supervision role holder
-                // can be set as profile owner after SUW
-                final String supervisor = mContext.getResources().getString(
-                        com.android.internal.R.string
-                                .config_defaultSupervisionProfileOwnerComponent);
-                if (supervisor == null && supervisionRolePackage == null) {
-                    throw new IllegalStateException("Unable to set profile owner post-setup, no"
-                            + "default supervisor profile owner defined");
-                }
-
-                final ComponentName supervisorComponent = ComponentName.unflattenFromString(
-                        supervisor);
-                if (!owner.equals(supervisorComponent)
-                        && !owner.getPackageName().equals(supervisionRolePackage)) {
+                if (!isSupervisionComponent(owner)) {
                     throw new IllegalStateException("Unable to set non-default profile owner"
                             + " post-setup " + owner);
                 }
@@ -12100,7 +12102,7 @@ public class DevicePolicyManagerService extends BaseIDevicePolicyManager {
         synchronized (getLockObject()) {
             // Allow testOnly admins to bypass supervision config requirement.
             Preconditions.checkCallAuthorization(isAdminTestOnlyLocked(who, caller.getUserId())
-                    || isDefaultSupervisor(caller), "Admin %s is not the "
+                    || isSupervisionComponent(caller.getComponentName()), "Admin %s is not the "
                     + "default supervision component", caller.getComponentName());
             DevicePolicyData policy = getUserData(caller.getUserId());
             policy.mSecondaryLockscreenEnabled = enabled;
@@ -12119,16 +12121,6 @@ public class DevicePolicyManagerService extends BaseIDevicePolicyManager {
         return isProfileOwner(caller) && isManagedProfile(caller.getUserId());
     }
 
-    private boolean isDefaultSupervisor(CallerIdentity caller) {
-        final String supervisor = mContext.getResources().getString(
-                com.android.internal.R.string.config_defaultSupervisionProfileOwnerComponent);
-        if (supervisor == null) {
-            return false;
-        }
-        final ComponentName supervisorComponent = ComponentName.unflattenFromString(supervisor);
-        return caller.getComponentName().equals(supervisorComponent);
-    }
-
     @Override
     public void setPreferentialNetworkServiceConfigs(
             List<PreferentialNetworkServiceConfig> preferentialNetworkServiceConfigs) {
@@ -13012,16 +13004,7 @@ public class DevicePolicyManagerService extends BaseIDevicePolicyManager {
                     return false;
                 }
 
-                final String supervisionString = mContext.getResources().getString(
-                        com.android.internal.R.string
-                                .config_defaultSupervisionProfileOwnerComponent);
-                if (supervisionString == null) {
-                    return false;
-                }
-
-                final ComponentName supervisorComponent = ComponentName.unflattenFromString(
-                        supervisionString);
-                return admin.info.getComponent().equals(supervisorComponent);
+                return isSupervisionComponent(admin.info.getComponent());
             }
         }
 
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 ec6b67450a1a9..8014d2502f48a 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
@@ -3322,19 +3322,48 @@ public class DevicePolicyManagerTest extends DpmTestBase {
     }
 
     @Test
-    public void testIsActiveSupervisionApp() throws Exception {
-        when(mServiceContext.resources
-                .getString(R.string.config_defaultSupervisionProfileOwnerComponent))
-                .thenReturn(admin1.flattenToString());
+    public void testSupervisionConfig() throws Exception {
+        final int uid = UserHandle.getUid(15, 19436);
+        addManagedProfile(admin1, uid, admin1);
+        mContext.binder.callingUid = uid;
 
-        final int PROFILE_USER = 15;
-        final int PROFILE_ADMIN = UserHandle.getUid(PROFILE_USER, 19436);
-        addManagedProfile(admin1, PROFILE_ADMIN, admin1);
-        mContext.binder.callingUid = PROFILE_ADMIN;
+        verifySupervisionConfig(uid, null, null);
+        verifySupervisionConfig(uid, "", null);
+        verifySupervisionConfig(uid, null, "");
+        verifySupervisionConfig(uid, "", "");
 
+        verifySupervisionConfig(uid, admin1.flattenToString(), null);
+        verifySupervisionConfig(uid, admin1.flattenToString(), "");
+
+        verifySupervisionConfig(uid, null, admin1.getPackageName());
+        verifySupervisionConfig(uid, "", admin1.getPackageName());
+    }
+
+    private void verifySupervisionConfig(
+            int uid , String configComponentName, String configPackageName) {
+        final boolean isAdmin = admin1.flattenToString().equals(configComponentName)
+                || admin1.getPackageName().equals(configPackageName);
+
+        final UserHandle user = UserHandle.getUserHandleForUid(uid);
         final DevicePolicyManagerInternal dpmi =
                 LocalServices.getService(DevicePolicyManagerInternal.class);
-        assertThat(dpmi.isActiveSupervisionApp(PROFILE_ADMIN)).isTrue();
+
+        when(mServiceContext.resources
+                .getString(R.string.config_defaultSupervisionProfileOwnerComponent))
+                .thenReturn(configComponentName);
+
+        when(mServiceContext.resources
+                .getString(R.string.config_systemSupervision))
+                .thenReturn(configPackageName);
+
+        if (isAdmin) {
+            assertThat(dpmi.isActiveSupervisionApp(uid)).isTrue();
+            assertThat(dpm.getProfileOwnerOrDeviceOwnerSupervisionComponent(user))
+                        .isEqualTo(admin1);
+        } else {
+            assertThat(dpmi.isActiveSupervisionApp(uid)).isFalse();
+            assertThat(dpm.getProfileOwnerOrDeviceOwnerSupervisionComponent(user)).isNull();
+        }
     }
 
     // Test if lock timeout on managed profile is handled correctly depending on whether profile