Block sleep options in Settings app when DISALLOW_CONFIG_SCREEN_TIMEOUT is set.

Test: m -j RunSettingsRoboTests
Bug: 67675734
Change-Id: Icdcadd03caa2884338d4300fb41f0036b8b9eb51
diff --git a/src/com/android/settings/display/TimeoutPreferenceController.java b/src/com/android/settings/display/TimeoutPreferenceController.java
index 1d86999..84f9641 100644
--- a/src/com/android/settings/display/TimeoutPreferenceController.java
+++ b/src/com/android/settings/display/TimeoutPreferenceController.java
@@ -13,9 +13,11 @@
  */
 package com.android.settings.display;
 
-import android.app.admin.DevicePolicyManager;
+import static android.provider.Settings.System.SCREEN_OFF_TIMEOUT;
+
 import android.content.Context;
 import android.os.UserHandle;
+import android.os.UserManager;
 import android.provider.Settings;
 import android.support.v7.preference.Preference;
 import android.util.Log;
@@ -25,10 +27,9 @@
 import com.android.settings.core.PreferenceControllerMixin;
 import com.android.settings.wrapper.DevicePolicyManagerWrapper;
 import com.android.settingslib.RestrictedLockUtils;
+import com.android.settingslib.RestrictedLockUtils.EnforcedAdmin;
 import com.android.settingslib.core.AbstractPreferenceController;
 
-import static android.provider.Settings.System.SCREEN_OFF_TIMEOUT;
-
 public class TimeoutPreferenceController extends AbstractPreferenceController implements
         PreferenceControllerMixin, Preference.OnPreferenceChangeListener {
 
@@ -69,6 +70,13 @@
             timeoutListPreference.removeUnusableTimeouts(maxTimeout, admin);
         }
         updateTimeoutPreferenceDescription(timeoutListPreference, currentTimeout);
+
+        EnforcedAdmin admin = RestrictedLockUtils.checkIfRestrictionEnforced(
+                        mContext, UserManager.DISALLOW_CONFIG_SCREEN_TIMEOUT,
+                        UserHandle.myUserId());
+        if(admin != null) {
+            timeoutListPreference.removeUnusableTimeouts(0/* disable all*/, admin);
+        }
     }
 
     @Override
diff --git a/tests/robotests/src/com/android/settings/display/TimeoutPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/display/TimeoutPreferenceControllerTest.java
index 33bc4f8..d4af8b8 100644
--- a/tests/robotests/src/com/android/settings/display/TimeoutPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/display/TimeoutPreferenceControllerTest.java
@@ -20,8 +20,12 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static junit.framework.Assert.assertEquals;
+import static junit.framework.Assert.assertTrue;
+
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -34,16 +38,20 @@
 import com.android.settings.TimeoutListPreference;
 import com.android.settings.testutils.SettingsRobolectricTestRunner;
 import com.android.settings.testutils.shadow.ShadowDevicePolicyManagerWrapper;
+import com.android.settingslib.RestrictedLockUtils.EnforcedAdmin;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
 
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.List;
 
 @RunWith(SettingsRobolectricTestRunner.class)
 @Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION,
@@ -93,4 +101,34 @@
         mController.updateState(mPreference);
         verify(mPreference).removeUnusableTimeouts(timeout, null);
     }
+
+    @Test
+    public void testUpdateStateWithAdminTimeoutsAndRestriction() {
+        final int profileUserId = UserHandle.myUserId();
+        final long timeout = 100;
+        when(mUserManager.getProfiles(profileUserId)).thenReturn(Collections.emptyList());
+        ShadowDevicePolicyManagerWrapper.setMaximumTimeToLock(profileUserId, timeout);
+
+        int userId = UserHandle.myUserId();
+        List<UserManager.EnforcingUser> enforcingUsers = new ArrayList<>();
+        // Add two enforcing users so that RestrictedLockUtils.checkIfRestrictionEnforced returns
+        // non-null.
+        enforcingUsers.add(new UserManager.EnforcingUser(userId,
+                UserManager.RESTRICTION_SOURCE_DEVICE_OWNER));
+        enforcingUsers.add(new UserManager.EnforcingUser(userId,
+                UserManager.RESTRICTION_SOURCE_PROFILE_OWNER));
+        when(mUserManager.getUserRestrictionSources(
+                UserManager.DISALLOW_CONFIG_SCREEN_TIMEOUT, UserHandle.of(userId)))
+                .thenReturn(enforcingUsers);
+
+        mController.updateState(mPreference);
+
+        ArgumentCaptor<Long> longCaptor = ArgumentCaptor.forClass(Long.class);
+        ArgumentCaptor<EnforcedAdmin> adminCaptor = ArgumentCaptor.forClass(EnforcedAdmin.class);
+
+        verify(mPreference, times(2)).removeUnusableTimeouts(
+                longCaptor.capture(), adminCaptor.capture());
+        assertEquals(0, (long)longCaptor.getValue());
+        assertTrue(adminCaptor.getValue() != null);
+    }
 }
diff --git a/tests/unit/src/com/android/settings/core/UserRestrictionTest.java b/tests/unit/src/com/android/settings/core/UserRestrictionTest.java
index f37c30b..7282f74 100644
--- a/tests/unit/src/com/android/settings/core/UserRestrictionTest.java
+++ b/tests/unit/src/com/android/settings/core/UserRestrictionTest.java
@@ -65,7 +65,8 @@
             UserManager.DISALLOW_CONFIG_VPN,
             UserManager.DISALLOW_CONFIG_MOBILE_NETWORKS,
             UserManager.DISALLOW_AIRPLANE_MODE,
-            UserManager.DISALLOW_CONFIG_BRIGHTNESS
+            UserManager.DISALLOW_CONFIG_BRIGHTNESS,
+            UserManager.DISALLOW_CONFIG_SCREEN_TIMEOUT
     );
 
     @Before