Add dialog for top switch on fully managed devices.

Bug: 311233745
Change-Id: Id4043675041d1cf9c0ed307f3a0c2dc65768641b
diff --git a/src/com/android/settings/security/ContentProtectionTogglePreferenceController.java b/src/com/android/settings/security/ContentProtectionTogglePreferenceController.java
index c874a5e..101364b 100644
--- a/src/com/android/settings/security/ContentProtectionTogglePreferenceController.java
+++ b/src/com/android/settings/security/ContentProtectionTogglePreferenceController.java
@@ -15,20 +15,22 @@
  */
 package com.android.settings.security;
 
-
 import android.content.ContentResolver;
 import android.content.Context;
 import android.provider.Settings;
 import android.widget.CompoundButton;
 import android.widget.CompoundButton.OnCheckedChangeListener;
 
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
+import androidx.preference.Preference;
 import androidx.preference.PreferenceScreen;
 
 import com.android.settings.R;
-import com.android.settings.Utils;
 import com.android.settings.core.TogglePreferenceController;
 import com.android.settings.widget.SettingsMainSwitchPreference;
+import com.android.settingslib.RestrictedLockUtils;
+import com.android.settingslib.RestrictedLockUtilsInternal;
 
 /** Preference controller for content protection toggle switch bar. */
 public class ContentProtectionTogglePreferenceController extends TogglePreferenceController
@@ -37,9 +39,9 @@
     @VisibleForTesting
     static final String KEY_CONTENT_PROTECTION_PREFERENCE = "content_protection_user_consent";
 
-    private SettingsMainSwitchPreference mSwitchBar;
+    @Nullable private SettingsMainSwitchPreference mSwitchBar;
+    @Nullable private RestrictedLockUtils.EnforcedAdmin mEnforcedAdmin;
     private final ContentResolver mContentResolver;
-    private final boolean isFullyManagedDevice = Utils.getDeviceOwnerComponent(mContext) != null;
 
     public ContentProtectionTogglePreferenceController(Context context, String preferenceKey) {
         super(context, preferenceKey);
@@ -53,7 +55,7 @@
 
     @Override
     public boolean isChecked() {
-        if (isFullyManagedDevice) {
+        if (mEnforcedAdmin != null) {
             // If fully managed device, it should always unchecked
             return false;
         }
@@ -70,12 +72,25 @@
     @Override
     public void displayPreference(PreferenceScreen screen) {
         super.displayPreference(screen);
+        final Preference preference = screen.findPreference(getPreferenceKey());
 
-        mSwitchBar = screen.findPreference(getPreferenceKey());
-        mSwitchBar.addOnSwitchChangeListener(this);
-        if (isFullyManagedDevice) {
-            // If fully managed device, the switch bar is greyed out
-            mSwitchBar.setEnabled(false);
+        if (preference instanceof SettingsMainSwitchPreference) {
+            mSwitchBar = (SettingsMainSwitchPreference) preference;
+            mSwitchBar.addOnSwitchChangeListener(this);
+        }
+    }
+
+    /**
+     * Temporary workaround for SettingsMainSwitchPreference.setDisabledByAdmin without user
+     * restriction.
+     */
+    @Override
+    public void updateState(Preference preference) {
+        super.updateState(preference);
+        // Assign the value to mEnforcedAdmin since it's needed in isChecked()
+        mEnforcedAdmin = getEnforcedAdmin();
+        if (mSwitchBar != null && mEnforcedAdmin != null) {
+            mSwitchBar.setDisabledByAdmin(mEnforcedAdmin);
         }
     }
 
@@ -90,4 +105,9 @@
     public int getSliceHighlightMenuRes() {
         return R.string.menu_key_security;
     }
+
+    @VisibleForTesting
+    protected RestrictedLockUtils.EnforcedAdmin getEnforcedAdmin() {
+        return RestrictedLockUtilsInternal.getDeviceOwner(mContext);
+    }
 }
diff --git a/src/com/android/settings/widget/SettingsMainSwitchPreference.java b/src/com/android/settings/widget/SettingsMainSwitchPreference.java
index fc43e58..9f6d787 100644
--- a/src/com/android/settings/widget/SettingsMainSwitchPreference.java
+++ b/src/com/android/settings/widget/SettingsMainSwitchPreference.java
@@ -78,7 +78,7 @@
         holder.setDividerAllowedAbove(false);
         holder.setDividerAllowedBelow(false);
 
-        if (mRestrictedHelper != null) {
+        if (mEnforcedAdmin == null && mRestrictedHelper != null) {
             mEnforcedAdmin = mRestrictedHelper.checkRestrictionEnforced();
         }
         mMainSwitchBar = (SettingsMainSwitchBar) holder.findViewById(R.id.main_switch_bar);
diff --git a/tests/robotests/src/com/android/settings/security/ContentProtectionTogglePreferenceControllerTest.java b/tests/robotests/src/com/android/settings/security/ContentProtectionTogglePreferenceControllerTest.java
index 63b6c3e..50e9a5c 100644
--- a/tests/robotests/src/com/android/settings/security/ContentProtectionTogglePreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/security/ContentProtectionTogglePreferenceControllerTest.java
@@ -13,16 +13,18 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 package com.android.settings.security;
 
 import static com.android.settings.security.ContentProtectionTogglePreferenceController.KEY_CONTENT_PROTECTION_PREFERENCE;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import android.content.ComponentName;
 import android.content.Context;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.provider.Settings;
@@ -32,6 +34,7 @@
 
 import com.android.settings.testutils.shadow.ShadowUtils;
 import com.android.settings.widget.SettingsMainSwitchPreference;
+import com.android.settingslib.RestrictedLockUtils;
 
 import org.junit.After;
 import org.junit.Before;
@@ -55,8 +58,9 @@
     @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
     @Rule public final MockitoRule mMockitoRule = MockitoJUnit.rule();
 
-    @Mock private PreferenceScreen mScreen;
+    @Mock private PreferenceScreen mMockScreen;
 
+    private RestrictedLockUtils.EnforcedAdmin mAdmin;
     private SettingsMainSwitchPreference mSwitchPreference;
     private final Context mContext = ApplicationProvider.getApplicationContext();
     private ContentProtectionTogglePreferenceController mController;
@@ -65,9 +69,10 @@
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mController = new ContentProtectionTogglePreferenceController(mContext, "key");
+        mController = new TestContentProtectionTogglePreferenceController();
         mSwitchPreference = new SettingsMainSwitchPreference(mContext);
-        when(mScreen.findPreference(mController.getPreferenceKey())).thenReturn(mSwitchPreference);
+        when(mMockScreen.findPreference(mController.getPreferenceKey()))
+                .thenReturn(mSwitchPreference);
         mSettingBackupValue = getContentProtectionGlobalSetting();
         Settings.Global.putInt(mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, 0);
     }
@@ -78,6 +83,7 @@
                 mContext.getContentResolver(),
                 KEY_CONTENT_PROTECTION_PREFERENCE,
                 mSettingBackupValue);
+        ShadowUtils.reset();
     }
 
     @Test
@@ -86,6 +92,49 @@
     }
 
     @Test
+    public void displayPreference() {
+        setUpFullyManagedMode();
+        SettingsMainSwitchPreference mockSwitchPreference =
+                mock(SettingsMainSwitchPreference.class);
+        when(mMockScreen.findPreference(any())).thenReturn(mockSwitchPreference);
+        when(mockSwitchPreference.getKey()).thenReturn(mController.getPreferenceKey());
+
+        mController = new TestContentProtectionTogglePreferenceController();
+        mController.displayPreference(mMockScreen);
+
+        assertThat(mockSwitchPreference).isNotNull();
+    }
+
+    @Test
+    public void updateState_notFullyManagedMode_enabled() {
+        SettingsMainSwitchPreference mockSwitchPreference =
+                mock(SettingsMainSwitchPreference.class);
+        when(mMockScreen.findPreference(any())).thenReturn(mockSwitchPreference);
+        when(mockSwitchPreference.getKey()).thenReturn(mController.getPreferenceKey());
+
+        mController = new TestContentProtectionTogglePreferenceController();
+        mController.displayPreference(mMockScreen);
+        mController.updateState(mockSwitchPreference);
+
+        verify(mockSwitchPreference, never()).setDisabledByAdmin(any());
+    }
+
+    @Test
+    public void updateState_fullyManagedMode_disabled() {
+        setUpFullyManagedMode();
+        SettingsMainSwitchPreference mockSwitchPreference =
+                mock(SettingsMainSwitchPreference.class);
+        when(mMockScreen.findPreference(any())).thenReturn(mockSwitchPreference);
+        when(mockSwitchPreference.getKey()).thenReturn(mController.getPreferenceKey());
+
+        mController = new TestContentProtectionTogglePreferenceController();
+        mController.displayPreference(mMockScreen);
+        mController.updateState(mockSwitchPreference);
+
+        verify(mockSwitchPreference).setDisabledByAdmin(mAdmin);
+    }
+
+    @Test
     public void isChecked_settingTurnOn() {
         Settings.Global.putInt(mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, 1);
 
@@ -94,15 +143,18 @@
 
     @Test
     public void isChecked_fullyManagedMode_settingTurnOff() {
-        final ComponentName componentName =
-                ComponentName.unflattenFromString("com.android.test/.DeviceAdminReceiver");
-        ShadowUtils.setDeviceOwnerComponent(componentName);
+        setUpFullyManagedMode();
         Settings.Global.putInt(mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, 1);
+        SettingsMainSwitchPreference mockSwitchPreference =
+                mock(SettingsMainSwitchPreference.class);
+        when(mMockScreen.findPreference(any())).thenReturn(mockSwitchPreference);
+        when(mockSwitchPreference.getKey()).thenReturn(mController.getPreferenceKey());
 
-        ContentProtectionTogglePreferenceController controller =
-                new ContentProtectionTogglePreferenceController(mContext, "key");
+        mController = new TestContentProtectionTogglePreferenceController();
+        mController.displayPreference(mMockScreen);
+        mController.updateState(mockSwitchPreference);
 
-        assertThat(controller.isChecked()).isFalse();
+        assertThat(mController.isChecked()).isFalse();
     }
 
     @Test
@@ -122,7 +174,6 @@
 
     @Test
     public void onSwitchChanged_switchChecked_manuallyEnabled() {
-        mController.displayPreference(mScreen);
         mController.setChecked(false);
 
         mController.onCheckedChanged(/* switchView= */ null, /* isChecked= */ true);
@@ -132,8 +183,6 @@
 
     @Test
     public void onSwitchChanged_switchUnchecked_manuallyDisabled() {
-        mController.displayPreference(mScreen);
-
         mController.onCheckedChanged(/* switchView= */ null, /* isChecked= */ false);
 
         assertThat(getContentProtectionGlobalSetting()).isEqualTo(-1);
@@ -143,4 +192,21 @@
         return Settings.Global.getInt(
                 mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, 0);
     }
+
+    private void setUpFullyManagedMode() {
+        mAdmin = new RestrictedLockUtils.EnforcedAdmin();
+    }
+
+    private class TestContentProtectionTogglePreferenceController
+            extends ContentProtectionTogglePreferenceController {
+
+        TestContentProtectionTogglePreferenceController() {
+            super(ContentProtectionTogglePreferenceControllerTest.this.mContext, "key");
+        }
+
+        @Override
+        protected RestrictedLockUtils.EnforcedAdmin getEnforcedAdmin() {
+            return mAdmin;
+        }
+    }
 }