Add the controller for the top switch.

Bug: 302189945
Change-Id: Ie43ac181e643a8b215ca830afc6ef670d91a5762
diff --git a/src/com/android/settings/security/ContentProtectionPreferenceFragment.java b/src/com/android/settings/security/ContentProtectionPreferenceFragment.java
index a58f6e5..476d93e 100644
--- a/src/com/android/settings/security/ContentProtectionPreferenceFragment.java
+++ b/src/com/android/settings/security/ContentProtectionPreferenceFragment.java
@@ -16,26 +16,35 @@
 
 package com.android.settings.security;
 
-import android.content.Context;
 import android.app.settings.SettingsEnums;
+import android.content.Context;
+import android.os.Bundle;
+import android.os.UserManager;
 
-import com.android.settings.dashboard.DashboardFragment;
+import androidx.annotation.VisibleForTesting;
+import androidx.preference.SwitchPreference;
+
 import com.android.settings.R;
-import com.android.settings.SettingsPreferenceFragment;
+import com.android.settings.Utils;
+import com.android.settings.dashboard.DashboardFragment;
 import com.android.settings.search.BaseSearchIndexProvider;
 import com.android.settingslib.search.SearchIndexable;
 
-import android.os.Bundle;
-
 @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);
@@ -44,7 +53,14 @@
     @Override
     public void onActivityCreated(Bundle savedInstanceState) {
         super.onActivityCreated(savedInstanceState);
-	// TODO(b/304681048): Update the toggles' behavior according to user's profile
+
+        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/ContentProtectionTogglePreferenceController.java b/src/com/android/settings/security/ContentProtectionTogglePreferenceController.java
new file mode 100644
index 0000000..686b25b
--- /dev/null
+++ b/src/com/android/settings/security/ContentProtectionTogglePreferenceController.java
@@ -0,0 +1,94 @@
+/*
+ * 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.ContentResolver;
+import android.content.Context;
+import android.provider.Settings;
+import android.widget.Switch;
+
+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.settings.widget.SettingsMainSwitchPreference;
+import com.android.settingslib.widget.OnMainSwitchChangeListener;
+
+/** Preference controller for content protection toggle switch bar. */
+public class ContentProtectionTogglePreferenceController extends TogglePreferenceController
+        implements OnMainSwitchChangeListener {
+
+    @VisibleForTesting
+    static final String KEY_CONTENT_PROTECTION_PREFERENCE = "content_protection_user_consent";
+
+    private SettingsMainSwitchPreference mSwitchBar;
+    private final ContentResolver mContentResolver;
+    private final boolean isFullyManagedDevice = Utils.getDeviceOwnerComponent(mContext) != null;
+
+    public ContentProtectionTogglePreferenceController(Context context, String preferenceKey) {
+        super(context, preferenceKey);
+        mContentResolver = context.getContentResolver();
+    }
+
+    @Override
+    public int getAvailabilityStatus() {
+        return AVAILABLE;
+    }
+
+    @Override
+    public boolean isChecked() {
+        if (isFullyManagedDevice) {
+            // If fully managed device, it should always unchecked
+            return false;
+        }
+        return Settings.Global.getInt(mContentResolver, KEY_CONTENT_PROTECTION_PREFERENCE, 0) >= 0;
+    }
+
+    @Override
+    public boolean setChecked(boolean isChecked) {
+        mSwitchBar.setChecked(isChecked);
+        Settings.Global.putInt(
+                mContentResolver, KEY_CONTENT_PROTECTION_PREFERENCE, isChecked ? 1 : -1);
+        return true;
+    }
+
+    @Override
+    public void displayPreference(PreferenceScreen screen) {
+        super.displayPreference(screen);
+
+        mSwitchBar = screen.findPreference(getPreferenceKey());
+        mSwitchBar.addOnSwitchChangeListener(this);
+        if (isFullyManagedDevice) {
+            // If fully managed device, the switch bar is greyed out
+            mSwitchBar.setEnabled(false);
+        }
+    }
+
+    @Override
+    public void onSwitchChanged(Switch switchView, boolean isChecked) {
+        if (isChecked != isChecked()) {
+            setChecked(isChecked);
+        }
+    }
+
+    @Override
+    public int getSliceHighlightMenuRes() {
+        return R.string.menu_key_security;
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/security/ContentProtectionPreferenceFragmentTest.java b/tests/robotests/src/com/android/settings/security/ContentProtectionPreferenceFragmentTest.java
index 431dd0c..c9b1c64 100644
--- a/tests/robotests/src/com/android/settings/security/ContentProtectionPreferenceFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/security/ContentProtectionPreferenceFragmentTest.java
@@ -17,16 +17,29 @@
 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;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -37,44 +50,116 @@
 import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
 
+import java.util.Arrays;
 import java.util.List;
 
 @RunWith(RobolectricTestRunner.class)
-@Config(shadows = ShadowDashboardFragment.class)
+@Config(
+        shadows = {
+            ShadowDashboardFragment.class,
+            ShadowUtils.class,
+        })
 public class ContentProtectionPreferenceFragmentTest {
+    private static final int TEST_PRIMARY_USER_ID = 10;
+    private static final int TEST_MANAGED_PROFILE_ID = 11;
 
-    @Mock
-    private ContentProtectionPreferenceFragment mMockFragment;
+    private ContentProtectionPreferenceFragment mFragment;
+    @Mock private UserManager mMockUserManager;
     private Context mContext;
+    private PreferenceScreen mScreen;
+    private SwitchPreference mWorkProfileSwitch;
 
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
-        mContext = spy(RuntimeEnvironment.application);
-        mMockFragment = spy(new ContentProtectionPreferenceFragment());
 
-        doReturn(mContext).when(mMockFragment).getContext();
+        mContext = spy(RuntimeEnvironment.application);
+        mFragment = spy(new ContentProtectionPreferenceFragment());
+        mScreen = spy(new PreferenceScreen(mContext, /* attrs= */ null));
+
+        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
     public void getMetricsCategory() {
-        assertThat(mMockFragment.getMetricsCategory()).isEqualTo(CONTENT_PROTECTION_PREFERENCE);
+        assertThat(mFragment.getMetricsCategory()).isEqualTo(CONTENT_PROTECTION_PREFERENCE);
     }
 
     @Test
-    public void getPreferenceScreenResId(){
-       assertThat(mMockFragment.getPreferenceScreenResId())
-               .isEqualTo(R.layout.content_protection_preference_fragment);
+    public void getPreferenceScreenResId() {
+        assertThat(mFragment.getPreferenceScreenResId())
+                .isEqualTo(R.layout.content_protection_preference_fragment);
     }
 
     @Test
     public void getNonIndexableKeys_existInXmlLayout() {
         final List<String> nonIndexableKeys =
-                             ContentProtectionPreferenceFragment.SEARCH_INDEX_DATA_PROVIDER
-                .getNonIndexableKeys(mContext);
+                ContentProtectionPreferenceFragment.SEARCH_INDEX_DATA_PROVIDER.getNonIndexableKeys(
+                        mContext);
         final List<String> allKeys =
-                XmlTestUtils.getKeysFromPreferenceXml(mContext,
-                        R.layout.content_protection_preference_fragment);
+                XmlTestUtils.getKeysFromPreferenceXml(
+                        mContext, R.layout.content_protection_preference_fragment);
 
         assertThat(allKeys).containsAtLeastElementsIn(nonIndexableKeys);
     }
@@ -83,10 +168,11 @@
     public void searchIndexProvider_shouldIndexResource() {
         final List<SearchIndexableResource> indexRes =
                 ContentProtectionPreferenceFragment.SEARCH_INDEX_DATA_PROVIDER
-                        .getXmlResourcesToIndex(mContext, /* enabled = */ true);
+                        .getXmlResourcesToIndex(mContext, /* enabled= */ true);
 
         assertThat(indexRes).isNotNull();
         assertThat(indexRes).isNotEmpty();
-        assertThat(indexRes.get(0).xmlResId).isEqualTo(mMockFragment.getPreferenceScreenResId());
+        assertThat(indexRes.get(0).xmlResId).isEqualTo(mFragment.getPreferenceScreenResId());
     }
 }
+
diff --git a/tests/robotests/src/com/android/settings/security/ContentProtectionTogglePreferenceControllerTest.java b/tests/robotests/src/com/android/settings/security/ContentProtectionTogglePreferenceControllerTest.java
new file mode 100644
index 0000000..b10ff22
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/security/ContentProtectionTogglePreferenceControllerTest.java
@@ -0,0 +1,146 @@
+/*
+ * 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.security.ContentProtectionTogglePreferenceController.KEY_CONTENT_PROTECTION_PREFERENCE;
+
+import static com.google.common.truth.Truth.assertThat;
+
+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;
+
+import androidx.preference.PreferenceScreen;
+import androidx.test.core.app.ApplicationProvider;
+
+import com.android.settings.testutils.shadow.ShadowUtils;
+import com.android.settings.widget.SettingsMainSwitchPreference;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.annotation.Config;
+
+@RunWith(RobolectricTestRunner.class)
+@Config(
+        shadows = {
+            ShadowUtils.class,
+        })
+public class ContentProtectionTogglePreferenceControllerTest {
+
+    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+    @Rule public final MockitoRule mMockitoRule = MockitoJUnit.rule();
+
+    @Mock private PreferenceScreen mScreen;
+
+    private SettingsMainSwitchPreference mSwitchPreference;
+    private final Context mContext = ApplicationProvider.getApplicationContext();
+    private ContentProtectionTogglePreferenceController mController;
+    private int mSettingBackupValue;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mController = new ContentProtectionTogglePreferenceController(mContext, "key");
+        mSwitchPreference = new SettingsMainSwitchPreference(mContext);
+        when(mScreen.findPreference(mController.getPreferenceKey())).thenReturn(mSwitchPreference);
+        mSettingBackupValue = getContentProtectionGlobalSetting();
+        Settings.Global.putInt(mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, 0);
+    }
+
+    @After
+    public void tearDown() {
+        Settings.Global.putInt(
+                mContext.getContentResolver(),
+                KEY_CONTENT_PROTECTION_PREFERENCE,
+                mSettingBackupValue);
+    }
+
+    @Test
+    public void isAvailable_alwaysAvailable() {
+        assertThat(mController.isAvailable()).isTrue();
+    }
+
+    @Test
+    public void isChecked_settingTurnOn() {
+        Settings.Global.putInt(mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, 1);
+
+        assertThat(mController.isChecked()).isTrue();
+    }
+
+    @Test
+    public void isChecked_fullyManagedMode_settingTurnOff() {
+        final ComponentName componentName =
+                ComponentName.unflattenFromString("com.android.test/.DeviceAdminReceiver");
+        ShadowUtils.setDeviceOwnerComponent(componentName);
+        Settings.Global.putInt(mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, 1);
+
+        ContentProtectionTogglePreferenceController controller =
+                new ContentProtectionTogglePreferenceController(mContext, "key");
+
+        assertThat(controller.isChecked()).isFalse();
+    }
+
+    @Test
+    public void isChecked_settingTurnOff() {
+        Settings.Global.putInt(
+                mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, -1);
+
+        assertThat(mController.isChecked()).isFalse();
+        assertThat(getContentProtectionGlobalSetting()).isEqualTo(-1);
+    }
+
+    @Test
+    public void isChecked_settingDefaultOn() {
+        assertThat(mController.isChecked()).isTrue();
+        assertThat(getContentProtectionGlobalSetting()).isEqualTo(0);
+    }
+
+    @Test
+    public void onSwitchChanged_switchChecked_manuallyEnabled() {
+        mController.displayPreference(mScreen);
+        mController.setChecked(false);
+
+        mController.onSwitchChanged(/* switchView= */ null, /* isChecked= */ true);
+
+        assertThat(getContentProtectionGlobalSetting()).isEqualTo(1);
+    }
+
+    @Test
+    public void onSwitchChanged_switchUnchecked_manuallyDisabled() {
+        mController.displayPreference(mScreen);
+
+        mController.onSwitchChanged(/* switchView= */ null, /* isChecked= */ false);
+
+        assertThat(getContentProtectionGlobalSetting()).isEqualTo(-1);
+    }
+
+    private int getContentProtectionGlobalSetting() {
+        return Settings.Global.getInt(
+                mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, 0);
+    }
+}