diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index deb9ee0a41c9a..f2127e634424c 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -3231,7 +3231,6 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
                 // If admin is a device or profile owner tidy that up first.
                 if (isDeviceOwner(adminReceiver, userHandle)) {
                     clearDeviceOwnerLocked(getDeviceOwnerAdminLocked(), userHandle);
-                    clearDeviceOwnerUserRestrictionLocked(UserHandle.of(userHandle));
                 }
                 if (isProfileOwner(adminReceiver, userHandle)) {
                     final ActiveAdmin admin = getActiveAdminUncheckedLocked(adminReceiver,
@@ -3247,10 +3246,8 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
         }
     }
 
-    // It's temporary solution to clear DISALLOW_ADD_USER after CTS
-    // STOPSHIP(b/31952368) when the restriction is moved from system to the device owner,
-    // it can be removed.
     private void clearDeviceOwnerUserRestrictionLocked(UserHandle userHandle) {
+        // ManagedProvisioning/DPC sets DISALLOW_ADD_USER. Clear to recover to the original state
         if (mUserManager.hasUserRestriction(UserManager.DISALLOW_ADD_USER, userHandle)) {
             mUserManager.setUserRestriction(UserManager.DISALLOW_ADD_USER, false, userHandle);
         }
@@ -6593,6 +6590,7 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
         mOwners.writeDeviceOwner();
         updateDeviceOwnerLocked();
 
+        clearDeviceOwnerUserRestrictionLocked(UserHandle.of(userId));
         mInjector.securityLogSetLoggingEnabledProperty(false);
         mSecurityLogMonitor.stop();
         setNetworkLoggingActiveInternal(false);
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 6fb65d59e31e7..d25cfb6f7e197 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
@@ -908,6 +908,8 @@ public class DevicePolicyManagerTest extends DpmTestBase {
         assertEquals(admin1, dpm.getDeviceOwnerComponentOnAnyUser());
 
         dpm.addUserRestriction(admin1, UserManager.DISALLOW_ADD_USER);
+        when(mContext.userManager.hasUserRestriction(eq(UserManager.DISALLOW_ADD_USER),
+                MockUtils.checkUserHandle(UserHandle.USER_SYSTEM))).thenReturn(true);
 
         assertTrue(dpm.isAdminActive(admin1));
         assertFalse(dpm.isRemovingAdmin(admin1, UserHandle.USER_SYSTEM));
@@ -937,6 +939,10 @@ public class DevicePolicyManagerTest extends DpmTestBase {
         // Now DO shouldn't be set.
         assertNull(dpm.getDeviceOwnerComponentOnAnyUser());
 
+        verify(mContext.userManager).setUserRestriction(eq(UserManager.DISALLOW_ADD_USER),
+                eq(false),
+                MockUtils.checkUserHandle(UserHandle.USER_SYSTEM));
+
         verify(mContext.userManagerInternal).setDevicePolicyUserRestrictions(
                 eq(UserHandle.USER_SYSTEM),
                 eq(null),