diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index dd44aa0a53cb5..2b5a06dac0a7e 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -3192,7 +3192,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,
@@ -3208,10 +3207,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);
         }
@@ -6573,6 +6570,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 23a1bb4462f40..d0e5159bea5f8 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
@@ -918,6 +918,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));
@@ -947,6 +949,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),