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 6fb65d5..d25cfb6 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 @@
         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 @@
         // 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),