From 6b9e20d32d6cc9df338aefa908cf29bfdc2905cc Mon Sep 17 00:00:00 2001
From: Wa Gao <wagao@google.com>
Date: Wed, 8 Nov 2023 06:09:27 +0000
Subject: [PATCH] Add the dialog when clicking on the diabled switch.

Bug: 302189945
Change-Id: I0bf7774b80a59c70a2ea1725e0adba522854adde
---
 ...content_protection_preference_fragment.xml |   5 +-
 .../ContentProtectionPreferenceFragment.java  |  14 --
 ...ContentProtectionWorkSwitchController.java |  86 +++++++++
 ...ntentProtectionPreferenceFragmentTest.java |  75 --------
 ...entProtectionWorkSwitchControllerTest.java | 163 ++++++++++++++++++
 5 files changed, 252 insertions(+), 91 deletions(-)
 create mode 100644 src/com/android/settings/security/ContentProtectionWorkSwitchController.java
 create mode 100644 tests/robotests/src/com/android/settings/security/ContentProtectionWorkSwitchControllerTest.java

diff --git a/res/layout/content_protection_preference_fragment.xml b/res/layout/content_protection_preference_fragment.xml
index a412d8618a4..4c7352e38c4 100644
--- a/res/layout/content_protection_preference_fragment.xml
+++ b/res/layout/content_protection_preference_fragment.xml
@@ -36,10 +36,11 @@
         settings:lottie_rawRes="@drawable/content_protection_preference_illustration"
         settings:searchable="false" />
 
-    <SwitchPreference
+    <com.android.settingslib.RestrictedSwitchPreference
         android:key="content_protection_preference_user_consent_work_profile_switch"
         android:title="@string/content_protection_preference_user_consent_work_profile_switch_title"
-        settings:isPreferenceVisible="false" />
+        settings:restrictedSwitchSummary="@string/summary_placeholder"
+        settings:controller="com.android.settings.security.ContentProtectionWorkSwitchController"/>
 
     <com.android.settingslib.widget.FooterPreference
         android:key="content_protection_preference_subpage_footer"
diff --git a/src/com/android/settings/security/ContentProtectionPreferenceFragment.java b/src/com/android/settings/security/ContentProtectionPreferenceFragment.java
index 476d93eea4a..5c72b4d7abb 100644
--- a/src/com/android/settings/security/ContentProtectionPreferenceFragment.java
+++ b/src/com/android/settings/security/ContentProtectionPreferenceFragment.java
@@ -34,17 +34,11 @@ import com.android.settingslib.search.SearchIndexable;
 public class ContentProtectionPreferenceFragment extends DashboardFragment {
     private static final String TAG = "ContentProtectionPreferenceFragment";
 
-    @VisibleForTesting
-    static final String KEY_WORK_PROFILE_SWITCH =
-            "content_protection_preference_user_consent_work_profile_switch";
-
     // Required by @SearchIndexable to make the fragment and preferences to be indexed.
     // Do not rename.
     public static final BaseSearchIndexProvider SEARCH_INDEX_DATA_PROVIDER =
             new BaseSearchIndexProvider(R.layout.content_protection_preference_fragment);
 
-    private SwitchPreference mWorkProfileSwitch;
-
     @Override
     public void onAttach(Context context) {
         super.onAttach(context);
@@ -53,14 +47,6 @@ public class ContentProtectionPreferenceFragment extends DashboardFragment {
     @Override
     public void onActivityCreated(Bundle savedInstanceState) {
         super.onActivityCreated(savedInstanceState);
-
-        mWorkProfileSwitch = getPreferenceScreen().findPreference(KEY_WORK_PROFILE_SWITCH);
-        // If any work profile on the device, display the disable toggle unchecked
-        if (Utils.getManagedProfile(getContext().getSystemService(UserManager.class)) != null) {
-            mWorkProfileSwitch.setVisible(true);
-            mWorkProfileSwitch.setEnabled(false);
-            mWorkProfileSwitch.setChecked(false);
-        }
     }
 
     @Override
diff --git a/src/com/android/settings/security/ContentProtectionWorkSwitchController.java b/src/com/android/settings/security/ContentProtectionWorkSwitchController.java
new file mode 100644
index 00000000000..0404dcde651
--- /dev/null
+++ b/src/com/android/settings/security/ContentProtectionWorkSwitchController.java
@@ -0,0 +1,86 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.settings.security;
+
+import android.content.Context;
+import android.os.UserHandle;
+import android.os.UserManager;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
+import androidx.preference.PreferenceScreen;
+
+import com.android.settings.R;
+import com.android.settings.Utils;
+import com.android.settings.core.TogglePreferenceController;
+import com.android.settingslib.RestrictedLockUtils;
+import com.android.settingslib.RestrictedSwitchPreference;
+
+/** Preference controller for content protection work profile switch bar. */
+public class ContentProtectionWorkSwitchController extends TogglePreferenceController {
+
+    public ContentProtectionWorkSwitchController(
+            @NonNull Context context, @NonNull String preferenceKey) {
+        super(context, preferenceKey);
+    }
+
+    @Override
+    public int getAvailabilityStatus() {
+        return getManagedProfile() != null ? AVAILABLE : CONDITIONALLY_UNAVAILABLE;
+    }
+
+    // The switch is always set to unchecked until Android V by design
+    @Override
+    public boolean isChecked() {
+        return false;
+    }
+
+    // The switch is disabled until Android V by design
+    @Override
+    public boolean setChecked(boolean isChecked) {
+        return false;
+    }
+
+    @Override
+    public void displayPreference(PreferenceScreen screen) {
+        super.displayPreference(screen);
+
+        RestrictedSwitchPreference switchPreference = screen.findPreference(getPreferenceKey());
+        UserHandle managedProfile = getManagedProfile();
+        if (managedProfile != null) {
+            switchPreference.setDisabledByAdmin(getEnforcedAdmin(managedProfile));
+        }
+    }
+
+    @Override
+    public int getSliceHighlightMenuRes() {
+        return R.string.menu_key_security;
+    }
+
+    @VisibleForTesting
+    @Nullable
+    protected UserHandle getManagedProfile() {
+        return Utils.getManagedProfile(mContext.getSystemService(UserManager.class));
+    }
+
+    @VisibleForTesting
+    @Nullable
+    protected RestrictedLockUtils.EnforcedAdmin getEnforcedAdmin(
+            @NonNull UserHandle managedProfile) {
+        return RestrictedLockUtils.getProfileOrDeviceOwner(mContext, managedProfile);
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/security/ContentProtectionPreferenceFragmentTest.java b/tests/robotests/src/com/android/settings/security/ContentProtectionPreferenceFragmentTest.java
index c9b1c64bdd0..2bd91a6a980 100644
--- a/tests/robotests/src/com/android/settings/security/ContentProtectionPreferenceFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/security/ContentProtectionPreferenceFragmentTest.java
@@ -18,25 +18,18 @@ package com.android.settings.security;
 
 import static android.app.settings.SettingsEnums.CONTENT_PROTECTION_PREFERENCE;
 
-import static com.android.settings.security.ContentProtectionPreferenceFragment.KEY_WORK_PROFILE_SWITCH;
-
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.spy;
 
-import android.content.ComponentName;
 import android.content.Context;
-import android.content.pm.UserInfo;
-import android.os.UserHandle;
-import android.os.UserManager;
 import android.provider.SearchIndexableResource;
 
 import androidx.preference.PreferenceScreen;
 import androidx.preference.SwitchPreference;
 
 import com.android.settings.R;
-import com.android.settings.Utils;
 import com.android.settings.testutils.XmlTestUtils;
 import com.android.settings.testutils.shadow.ShadowDashboardFragment;
 import com.android.settings.testutils.shadow.ShadowUtils;
@@ -44,13 +37,11 @@ import com.android.settings.testutils.shadow.ShadowUtils;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
 
-import java.util.Arrays;
 import java.util.List;
 
 @RunWith(RobolectricTestRunner.class)
@@ -60,14 +51,9 @@ import java.util.List;
             ShadowUtils.class,
         })
 public class ContentProtectionPreferenceFragmentTest {
-    private static final int TEST_PRIMARY_USER_ID = 10;
-    private static final int TEST_MANAGED_PROFILE_ID = 11;
-
     private ContentProtectionPreferenceFragment mFragment;
-    @Mock private UserManager mMockUserManager;
     private Context mContext;
     private PreferenceScreen mScreen;
-    private SwitchPreference mWorkProfileSwitch;
 
     @Before
     public void setUp() throws Exception {
@@ -79,66 +65,6 @@ public class ContentProtectionPreferenceFragmentTest {
 
         doReturn(mContext).when(mFragment).getContext();
         doReturn(mScreen).when(mFragment).getPreferenceScreen();
-
-        mWorkProfileSwitch = new SwitchPreference(mContext);
-        mWorkProfileSwitch.setVisible(false);
-        doReturn(mWorkProfileSwitch).when(mScreen).findPreference(KEY_WORK_PROFILE_SWITCH);
-
-        doReturn(mMockUserManager).when(mContext).getSystemService(UserManager.class);
-        doReturn(TEST_PRIMARY_USER_ID).when(mMockUserManager).getUserHandle();
-        UserInfo primaryUser =
-                new UserInfo(
-                        TEST_PRIMARY_USER_ID,
-                        null,
-                        UserInfo.FLAG_INITIALIZED | UserInfo.FLAG_PRIMARY);
-        doReturn(primaryUser).when(mMockUserManager).getUserInfo(TEST_PRIMARY_USER_ID);
-        UserInfo managedProfile =
-                new UserInfo(
-                        TEST_MANAGED_PROFILE_ID,
-                        null,
-                        UserInfo.FLAG_INITIALIZED | UserInfo.FLAG_MANAGED_PROFILE);
-        doReturn(managedProfile).when(mMockUserManager).getUserInfo(TEST_MANAGED_PROFILE_ID);
-    }
-
-    @Test
-    public void onActivityCreated_workProfileDisplayWorkSwitch() {
-        UserHandle[] userHandles =
-                new UserHandle[] {
-                    new UserHandle(TEST_PRIMARY_USER_ID), new UserHandle(TEST_MANAGED_PROFILE_ID)
-                };
-        doReturn(Arrays.asList(userHandles)).when(mMockUserManager).getUserProfiles();
-
-        assertThat(Utils.getManagedProfile(mMockUserManager).getIdentifier())
-                .isEqualTo(TEST_MANAGED_PROFILE_ID);
-
-        mFragment.onActivityCreated(null);
-
-        assertThat(mWorkProfileSwitch.isVisible()).isTrue();
-        assertThat(mWorkProfileSwitch.isChecked()).isFalse();
-        assertThat(mWorkProfileSwitch.isEnabled()).isFalse();
-    }
-
-    @Test
-    public void onActivityCreated_fullyManagedMode_bottomSwitchInvisible() {
-        final ComponentName componentName =
-                ComponentName.unflattenFromString("com.android.test/.DeviceAdminReceiver");
-        ShadowUtils.setDeviceOwnerComponent(componentName);
-
-        mFragment.onActivityCreated(null);
-
-        assertThat(mWorkProfileSwitch.isVisible()).isFalse();
-    }
-
-    @Test
-    public void onActivityCreated_personalProfileHideWorkSwitch() {
-        UserHandle[] userHandles = new UserHandle[] {new UserHandle(TEST_PRIMARY_USER_ID)};
-        doReturn(Arrays.asList(userHandles)).when(mMockUserManager).getUserProfiles();
-
-        assertThat(Utils.getManagedProfile(mMockUserManager)).isNull();
-
-        mFragment.onActivityCreated(null);
-
-        assertThat(mWorkProfileSwitch.isVisible()).isFalse();
     }
 
     @Test
@@ -175,4 +101,3 @@ public class ContentProtectionPreferenceFragmentTest {
         assertThat(indexRes.get(0).xmlResId).isEqualTo(mFragment.getPreferenceScreenResId());
     }
 }
-
diff --git a/tests/robotests/src/com/android/settings/security/ContentProtectionWorkSwitchControllerTest.java b/tests/robotests/src/com/android/settings/security/ContentProtectionWorkSwitchControllerTest.java
new file mode 100644
index 00000000000..8d35e4dcaa2
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/security/ContentProtectionWorkSwitchControllerTest.java
@@ -0,0 +1,163 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.security;
+
+import static com.android.settings.core.BasePreferenceController.AVAILABLE;
+import static com.android.settings.core.BasePreferenceController.CONDITIONALLY_UNAVAILABLE;
+
+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.Context;
+import android.os.UserHandle;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+import androidx.preference.PreferenceScreen;
+import androidx.test.core.app.ApplicationProvider;
+
+import com.android.settingslib.RestrictedLockUtils;
+import com.android.settingslib.RestrictedSwitchPreference;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+
+@RunWith(RobolectricTestRunner.class)
+public class ContentProtectionWorkSwitchControllerTest {
+    private static final UserHandle TEST_USER_HANDLE = UserHandle.of(10);
+
+    private final Context mContext = ApplicationProvider.getApplicationContext();
+
+    @Mock private PreferenceScreen mMockPreferenceScreen;
+    private ContentProtectionWorkSwitchController mController;
+    private UserHandle mManagedProfileUserHandle;
+    private RestrictedLockUtils.EnforcedAdmin mEnforcedAdmin;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mController = new TestContentProtectionWorkSwitchController();
+    }
+
+    @Test
+    public void isAvailable_managedProfile_available() {
+        mManagedProfileUserHandle = TEST_USER_HANDLE;
+
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(AVAILABLE);
+        assertThat(mController.isAvailable()).isTrue();
+    }
+
+    @Test
+    public void isAvailable_noManagedProfile_notAvailable() {
+        mManagedProfileUserHandle = null;
+
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(CONDITIONALLY_UNAVAILABLE);
+        assertThat(mController.isAvailable()).isFalse();
+    }
+
+    @Test
+    public void isChecked_noManagedProfile_alwaysOff() {
+        mManagedProfileUserHandle = null;
+
+        assertThat(mController.isChecked()).isFalse();
+    }
+
+    @Test
+    public void isChecked_managedProfile_alwaysOff() {
+        mManagedProfileUserHandle = TEST_USER_HANDLE;
+
+        assertThat(mController.isChecked()).isFalse();
+    }
+
+    @Test
+    public void setChecked_alwaysFalse() {
+        assertThat(mController.setChecked(true)).isFalse();
+        assertThat(mController.setChecked(false)).isFalse();
+    }
+
+    @Test
+    public void displayPreference_managedProfile_disabled() {
+        mManagedProfileUserHandle = TEST_USER_HANDLE;
+        mEnforcedAdmin = new RestrictedLockUtils.EnforcedAdmin();
+        RestrictedSwitchPreference mockSwitchPreference = mock(RestrictedSwitchPreference.class);
+        when(mMockPreferenceScreen.findPreference(any())).thenReturn(mockSwitchPreference);
+        when(mockSwitchPreference.getKey()).thenReturn(mController.getPreferenceKey());
+
+        mController.displayPreference(mMockPreferenceScreen);
+
+        assertThat(mController.isAvailable()).isTrue();
+        verify(mockSwitchPreference).setDisabledByAdmin(mEnforcedAdmin);
+    }
+
+    @Test
+    public void displayPreference_noManagedProfile_notDisabled() {
+        mManagedProfileUserHandle = null;
+        mEnforcedAdmin = new RestrictedLockUtils.EnforcedAdmin();
+        RestrictedSwitchPreference mockSwitchPreference = mock(RestrictedSwitchPreference.class);
+        when(mMockPreferenceScreen.findPreference(any())).thenReturn(mockSwitchPreference);
+        when(mockSwitchPreference.getKey()).thenReturn(mController.getPreferenceKey());
+
+        mController.displayPreference(mMockPreferenceScreen);
+
+        assertThat(mController.isAvailable()).isFalse();
+        verify(mockSwitchPreference, never()).setDisabledByAdmin(any());
+    }
+
+    @Test
+    public void displayPreference_noEnforcedAdmin_notDisabled() {
+        mManagedProfileUserHandle = null;
+        mEnforcedAdmin = null;
+        RestrictedSwitchPreference mockSwitchPreference = mock(RestrictedSwitchPreference.class);
+        when(mMockPreferenceScreen.findPreference(any())).thenReturn(mockSwitchPreference);
+        when(mockSwitchPreference.getKey()).thenReturn(mController.getPreferenceKey());
+
+        mController.displayPreference(mMockPreferenceScreen);
+
+        assertThat(mController.isAvailable()).isFalse();
+        verify(mockSwitchPreference, never()).setDisabledByAdmin(any());
+    }
+
+    private class TestContentProtectionWorkSwitchController
+            extends ContentProtectionWorkSwitchController {
+
+        TestContentProtectionWorkSwitchController() {
+            super(ContentProtectionWorkSwitchControllerTest.this.mContext, "key");
+        }
+
+        @Override
+        @Nullable
+        protected UserHandle getManagedProfile() {
+            return mManagedProfileUserHandle;
+        }
+
+        @Override
+        @Nullable
+        protected RestrictedLockUtils.EnforcedAdmin getEnforcedAdmin(
+                @NonNull UserHandle managedProfile) {
+            return mEnforcedAdmin;
+        }
+    }
+}