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 ca918c6..70e4b6d 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
@@ -34,7 +34,6 @@
 import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.nullable;
 import static org.mockito.Mockito.reset;
@@ -2244,27 +2243,32 @@
         reset(getServices().settings);
 
         dpm.setMaximumTimeToLock(admin1, 0);
-        verifyScreenTimeoutCall(null, false);
+        verifyScreenTimeoutCall(null, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(false);
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
         dpm.setMaximumTimeToLock(admin1, 1);
-        verifyScreenTimeoutCall(1, true);
+        verifyScreenTimeoutCall(1L, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(true);
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
         dpm.setMaximumTimeToLock(admin2, 10);
-        verifyScreenTimeoutCall(null, false);
+        verifyScreenTimeoutCall(null, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(false);
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
         dpm.setMaximumTimeToLock(admin1, 5);
-        verifyScreenTimeoutCall(5, true);
+        verifyScreenTimeoutCall(5L, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(true);
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
         dpm.setMaximumTimeToLock(admin2, 4);
-        verifyScreenTimeoutCall(4, true);
+        verifyScreenTimeoutCall(4L, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(true);
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
@@ -2272,24 +2276,89 @@
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
-        dpm.setMaximumTimeToLock(admin2, Integer.MAX_VALUE);
-        verifyScreenTimeoutCall(Integer.MAX_VALUE, true);
-        reset(getServices().powerManagerInternal);
-        reset(getServices().settings);
-
-        dpm.setMaximumTimeToLock(admin2, Integer.MAX_VALUE + 1);
-        verifyScreenTimeoutCall(Integer.MAX_VALUE, true);
+        dpm.setMaximumTimeToLock(admin2, Long.MAX_VALUE);
+        verifyScreenTimeoutCall(Long.MAX_VALUE, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(true);
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
         dpm.setMaximumTimeToLock(admin2, 10);
-        verifyScreenTimeoutCall(10, true);
+        verifyScreenTimeoutCall(10L, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(true);
         reset(getServices().powerManagerInternal);
         reset(getServices().settings);
 
-        // There's no restriction; shold be set to MAX.
+        // There's no restriction; should be set to MAX.
         dpm.setMaximumTimeToLock(admin2, 0);
-        verifyScreenTimeoutCall(Integer.MAX_VALUE, false);
+        verifyScreenTimeoutCall(Long.MAX_VALUE, UserHandle.USER_SYSTEM);
+        verifyStayOnWhilePluggedCleared(false);
+    }
+
+    // Test if lock timeout on managed profile is handled correctly depending on whether profile
+    // uses separate challenge.
+    public void testSetMaximumTimeToLockProfile() throws Exception {
+        final int PROFILE_USER = 15;
+        final int PROFILE_ADMIN = UserHandle.getUid(PROFILE_USER, 19436);
+        addManagedProfile(admin1, PROFILE_ADMIN, admin1);
+        mContext.binder.callingUid = PROFILE_ADMIN;
+        final DevicePolicyManagerInternal dpmi =
+                LocalServices.getService(DevicePolicyManagerInternal.class);
+
+        dpm.setMaximumTimeToLock(admin1, 0);
+
+        reset(getServices().powerManagerInternal);
+        reset(getServices().settings);
+
+        // First add timeout for the profile.
+        dpm.setMaximumTimeToLock(admin1, 10);
+        verifyScreenTimeoutCall(10L, UserHandle.USER_SYSTEM);
+
+        reset(getServices().powerManagerInternal);
+        reset(getServices().settings);
+
+        // Add separate challenge
+        when(getServices().lockPatternUtils
+                .isSeparateProfileChallengeEnabled(eq(PROFILE_USER))).thenReturn(true);
+        dpmi.reportSeparateProfileChallengeChanged(PROFILE_USER);
+
+        verifyScreenTimeoutCall(10L, PROFILE_USER);
+        verifyScreenTimeoutCall(Long.MAX_VALUE, UserHandle.USER_SYSTEM);
+
+        reset(getServices().powerManagerInternal);
+        reset(getServices().settings);
+
+        // Remove the timeout.
+        dpm.setMaximumTimeToLock(admin1, 0);
+        verifyScreenTimeoutCall(Long.MAX_VALUE, PROFILE_USER);
+        verifyScreenTimeoutCall(null , UserHandle.USER_SYSTEM);
+
+        reset(getServices().powerManagerInternal);
+        reset(getServices().settings);
+
+        // Add it back.
+        dpm.setMaximumTimeToLock(admin1, 10);
+        verifyScreenTimeoutCall(10L, PROFILE_USER);
+        verifyScreenTimeoutCall(null, UserHandle.USER_SYSTEM);
+
+        reset(getServices().powerManagerInternal);
+        reset(getServices().settings);
+
+        // Remove separate challenge.
+        reset(getServices().lockPatternUtils);
+        when(getServices().lockPatternUtils
+                .isSeparateProfileChallengeEnabled(eq(PROFILE_USER))).thenReturn(false);
+        dpmi.reportSeparateProfileChallengeChanged(PROFILE_USER);
+
+        verifyScreenTimeoutCall(Long.MAX_VALUE, PROFILE_USER);
+        verifyScreenTimeoutCall(10L , UserHandle.USER_SYSTEM);
+
+        reset(getServices().powerManagerInternal);
+        reset(getServices().settings);
+
+        // Remove the timeout.
+        dpm.setMaximumTimeToLock(admin1, 0);
+        verifyScreenTimeoutCall(null, PROFILE_USER);
+        verifyScreenTimeoutCall(Long.MAX_VALUE, UserHandle.USER_SYSTEM);
     }
 
     public void testSetRequiredStrongAuthTimeout_DeviceOwner() throws Exception {
@@ -2365,15 +2434,17 @@
                 () -> dpm.setRequiredStrongAuthTimeout(admin1, -ONE_MINUTE));
     }
 
-    private void verifyScreenTimeoutCall(Integer expectedTimeout,
-            boolean shouldStayOnWhilePluggedInBeCleared) {
+    private void verifyScreenTimeoutCall(Long expectedTimeout, int userId) {
         if (expectedTimeout == null) {
             verify(getServices().powerManagerInternal, times(0))
-                    .setMaximumScreenOffTimeoutFromDeviceAdmin(anyInt());
+                    .setMaximumScreenOffTimeoutFromDeviceAdmin(eq(userId), anyLong());
         } else {
             verify(getServices().powerManagerInternal, times(1))
-                    .setMaximumScreenOffTimeoutFromDeviceAdmin(eq(expectedTimeout));
+                    .setMaximumScreenOffTimeoutFromDeviceAdmin(eq(userId), eq(expectedTimeout));
         }
+    }
+
+    private void verifyStayOnWhilePluggedCleared(boolean cleared) {
         // TODO Verify calls to settingsGlobalPutInt.  Tried but somehow mockito threw
         // UnfinishedVerificationException.
     }