Merge "Integrate admin controls for content protection" into main

src/com/android/settings/security/ContentProtectionPreferenceUtils.java

@@ -15,15 +15,23 @@
  */
 package com.android.settings.security;
 
+import static android.view.contentprotection.flags.Flags.manageDevicePolicyEnabled;
+
 import static com.android.internal.R.string.config_defaultContentProtectionService;
 
+import android.app.admin.DevicePolicyManager;
 import android.content.ComponentName;
 import android.content.Context;
+import android.content.pm.PackageManager;
+import android.os.UserHandle;
+import android.os.UserManager;
 import android.provider.DeviceConfig;
 import android.view.contentcapture.ContentCaptureManager;
 
-import androidx.annotation.Nullable;
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import com.android.settings.Utils;
 
 /** Util class for content protection preference. */
 public class ContentProtectionPreferenceUtils {
@@ -60,4 +68,49 @@ public class ContentProtectionPreferenceUtils {
                 ContentCaptureManager.DEVICE_CONFIG_PROPERTY_ENABLE_CONTENT_PROTECTION_RECEIVER,
                 ContentCaptureManager.DEFAULT_ENABLE_CONTENT_PROTECTION_RECEIVER);
     }
+
+    /** Returns the managed profile or null if none exists. */
+    @Nullable
+    public static UserHandle getManagedProfile(@NonNull Context context) {
+        UserManager userManager = context.getSystemService(UserManager.class);
+        if (userManager == null) {
+            return null;
+        }
+        return Utils.getManagedProfile(userManager);
+    }
+
+    /** Returns the current content protection policy. */
+    @DevicePolicyManager.ContentProtectionPolicy
+    public static int getContentProtectionPolicy(
+            @NonNull Context context, @Nullable UserHandle managedProfile) {
+        if (!manageDevicePolicyEnabled()) {
+            return DevicePolicyManager.CONTENT_PROTECTION_DISABLED;
+        }
+        Context policyContext = createContentProtectionPolicyContext(context, managedProfile);
+        return getContentProtectionPolicyWithGivenContext(policyContext);
+    }
+
+    @NonNull
+    private static Context createContentProtectionPolicyContext(
+            @NonNull Context context, @Nullable UserHandle managedProfile) {
+        if (managedProfile == null) {
+            return context;
+        }
+        try {
+            return context.createPackageContextAsUser(
+                    context.getPackageName(), /* flags= */ 0, managedProfile);
+        } catch (PackageManager.NameNotFoundException ex) {
+            throw new IllegalStateException(ex);
+        }
+    }
+
+    @DevicePolicyManager.ContentProtectionPolicy
+    private static int getContentProtectionPolicyWithGivenContext(@NonNull Context context) {
+        DevicePolicyManager devicePolicyManager =
+                context.getSystemService(DevicePolicyManager.class);
+        if (devicePolicyManager == null) {
+            return DevicePolicyManager.CONTENT_PROTECTION_DISABLED;
+        }
+        return devicePolicyManager.getContentProtectionPolicy(/* admin= */ null);
+    }
 }

src/com/android/settings/security/ContentProtectionTogglePreferenceController.java

@@ -15,12 +15,17 @@
  */
 package com.android.settings.security;
 
+import static android.view.contentprotection.flags.Flags.manageDevicePolicyEnabled;
+
+import android.app.admin.DevicePolicyManager;
 import android.content.ContentResolver;
 import android.content.Context;
+import android.os.UserHandle;
 import android.provider.Settings;
 import android.widget.CompoundButton;
 import android.widget.CompoundButton.OnCheckedChangeListener;
 
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import androidx.preference.Preference;
@@ -40,12 +45,22 @@ public class ContentProtectionTogglePreferenceController extends TogglePreferenc
     static final String KEY_CONTENT_PROTECTION_PREFERENCE = "content_protection_user_consent";
 
     @Nullable private SettingsMainSwitchPreference mSwitchBar;
+
     @Nullable private RestrictedLockUtils.EnforcedAdmin mEnforcedAdmin;
-    private final ContentResolver mContentResolver;
+
+    @NonNull private final ContentResolver mContentResolver;
+
+    @DevicePolicyManager.ContentProtectionPolicy
+    private int mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_DISABLED;
 
     public ContentProtectionTogglePreferenceController(Context context, String preferenceKey) {
         super(context, preferenceKey);
         mContentResolver = context.getContentResolver();
+
+        if (manageDevicePolicyEnabled()) {
+            mEnforcedAdmin = getEnforcedAdmin();
+            mContentProtectionPolicy = getContentProtectionPolicy(getManagedProfile());
+        }
     }
 
     @Override
@@ -56,14 +71,30 @@ public class ContentProtectionTogglePreferenceController extends TogglePreferenc
     @Override
     public boolean isChecked() {
         if (mEnforcedAdmin != null) {
-            // If fully managed device, it should always unchecked
+            if (!manageDevicePolicyEnabled()) {
-            return false;
+                // If fully managed device, it should always unchecked
+                return false;
+            }
+
+            if (mContentProtectionPolicy == DevicePolicyManager.CONTENT_PROTECTION_DISABLED) {
+                return false;
+            }
+            if (mContentProtectionPolicy == DevicePolicyManager.CONTENT_PROTECTION_ENABLED) {
+                return true;
+            }
         }
         return Settings.Global.getInt(mContentResolver, KEY_CONTENT_PROTECTION_PREFERENCE, 0) >= 0;
     }
 
     @Override
     public boolean setChecked(boolean isChecked) {
+        if (manageDevicePolicyEnabled()) {
+            if (mEnforcedAdmin != null
+                    && mContentProtectionPolicy
+                            != DevicePolicyManager.CONTENT_PROTECTION_NOT_CONTROLLED_BY_POLICY) {
+                return false;
+            }
+        }
         Settings.Global.putInt(
                 mContentResolver, KEY_CONTENT_PROTECTION_PREFERENCE, isChecked ? 1 : -1);
         return true;
@@ -80,16 +111,20 @@ public class ContentProtectionTogglePreferenceController extends TogglePreferenc
         }
     }
 
-    /**
+    // Workaround for SettingsMainSwitchPreference.setDisabledByAdmin without user restriction.
-     * 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 (!manageDevicePolicyEnabled()) {
-        if (mSwitchBar != null && mEnforcedAdmin != null) {
+            // Assign the value to mEnforcedAdmin since it's needed in isChecked()
+            mEnforcedAdmin = getEnforcedAdmin();
+            mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_DISABLED;
+        }
+        if (mSwitchBar != null
+                && mEnforcedAdmin != null
+                && mContentProtectionPolicy
+                        != DevicePolicyManager.CONTENT_PROTECTION_NOT_CONTROLLED_BY_POLICY) {
             mSwitchBar.setDisabledByAdmin(mEnforcedAdmin);
         }
     }
@@ -107,7 +142,20 @@ public class ContentProtectionTogglePreferenceController extends TogglePreferenc
     }
 
     @VisibleForTesting
+    @Nullable
+    protected UserHandle getManagedProfile() {
+        return ContentProtectionPreferenceUtils.getManagedProfile(mContext);
+    }
+
+    @VisibleForTesting
+    @Nullable
     protected RestrictedLockUtils.EnforcedAdmin getEnforcedAdmin() {
         return RestrictedLockUtilsInternal.getDeviceOwner(mContext);
     }
+
+    @VisibleForTesting
+    @DevicePolicyManager.ContentProtectionPolicy
+    protected int getContentProtectionPolicy(@Nullable UserHandle userHandle) {
+        return ContentProtectionPreferenceUtils.getContentProtectionPolicy(mContext, userHandle);
+    }
 }

src/com/android/settings/security/ContentProtectionWorkSwitchController.java

@@ -15,9 +15,11 @@
  */
 package com.android.settings.security;
 
+import static android.view.contentprotection.flags.Flags.manageDevicePolicyEnabled;
+
+import android.app.admin.DevicePolicyManager;
 import android.content.Context;
 import android.os.UserHandle;
-import android.os.UserManager;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -25,7 +27,6 @@ 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;
@@ -33,25 +34,49 @@ import com.android.settingslib.RestrictedSwitchPreference;
 /** Preference controller for content protection work profile switch bar. */
 public class ContentProtectionWorkSwitchController extends TogglePreferenceController {
 
+    @Nullable private UserHandle mManagedProfile;
+
+    @DevicePolicyManager.ContentProtectionPolicy
+    private int mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_DISABLED;
+
     public ContentProtectionWorkSwitchController(
             @NonNull Context context, @NonNull String preferenceKey) {
         super(context, preferenceKey);
+
+        if (manageDevicePolicyEnabled()) {
+            mManagedProfile = getManagedProfile();
+            if (mManagedProfile != null) {
+                mContentProtectionPolicy = getContentProtectionPolicy(mManagedProfile);
+            }
+        }
     }
 
     @Override
     public int getAvailabilityStatus() {
-        return getManagedProfile() != null ? AVAILABLE : CONDITIONALLY_UNAVAILABLE;
+        if (!manageDevicePolicyEnabled()) {
+            return getManagedProfile() != null ? AVAILABLE : CONDITIONALLY_UNAVAILABLE;
+        }
+        if (mManagedProfile == null) {
+            return CONDITIONALLY_UNAVAILABLE;
+        }
+        if (mContentProtectionPolicy
+                == DevicePolicyManager.CONTENT_PROTECTION_NOT_CONTROLLED_BY_POLICY) {
+            return CONDITIONALLY_UNAVAILABLE;
+        }
+        return AVAILABLE;
     }
 
-    // The switch is always set to unchecked until Android V by design
     @Override
     public boolean isChecked() {
-        return false;
+        if (!manageDevicePolicyEnabled()) {
+            return false;
+        }
+        return mContentProtectionPolicy == DevicePolicyManager.CONTENT_PROTECTION_ENABLED;
     }
 
-    // The switch is disabled until Android V by design
     @Override
     public boolean setChecked(boolean isChecked) {
+        // Controlled by the admin API
         return false;
     }
 
@@ -59,10 +84,13 @@ public class ContentProtectionWorkSwitchController extends TogglePreferenceContr
     public void displayPreference(PreferenceScreen screen) {
         super.displayPreference(screen);
 
-        RestrictedSwitchPreference switchPreference = screen.findPreference(getPreferenceKey());
+        UserHandle managedProfile =
-        UserHandle managedProfile = getManagedProfile();
+                manageDevicePolicyEnabled() ? mManagedProfile : getManagedProfile();
         if (managedProfile != null) {
-            switchPreference.setDisabledByAdmin(getEnforcedAdmin(managedProfile));
+            RestrictedSwitchPreference switchPreference = screen.findPreference(getPreferenceKey());
+            if (switchPreference != null) {
+                switchPreference.setDisabledByAdmin(getEnforcedAdmin(managedProfile));
+            }
         }
     }
 
@@ -74,13 +102,18 @@ public class ContentProtectionWorkSwitchController extends TogglePreferenceContr
     @VisibleForTesting
     @Nullable
     protected UserHandle getManagedProfile() {
-        return Utils.getManagedProfile(mContext.getSystemService(UserManager.class));
+        return ContentProtectionPreferenceUtils.getManagedProfile(mContext);
     }
 
     @VisibleForTesting
     @Nullable
-    protected RestrictedLockUtils.EnforcedAdmin getEnforcedAdmin(
+    protected RestrictedLockUtils.EnforcedAdmin getEnforcedAdmin(@NonNull UserHandle userHandle) {
-            @NonNull UserHandle managedProfile) {
+        return RestrictedLockUtils.getProfileOrDeviceOwner(mContext, userHandle);
-        return RestrictedLockUtils.getProfileOrDeviceOwner(mContext, managedProfile);
+    }
+
+    @VisibleForTesting
+    @DevicePolicyManager.ContentProtectionPolicy
+    protected int getContentProtectionPolicy(@Nullable UserHandle userHandle) {
+        return ContentProtectionPreferenceUtils.getContentProtectionPolicy(mContext, userHandle);
     }
 }

tests/robotests/src/com/android/settings/security/ContentProtectionPreferenceUtilsTest.java

@@ -16,47 +16,73 @@
 
 package com.android.settings.security;
 
+import static android.view.contentprotection.flags.Flags.FLAG_MANAGE_DEVICE_POLICY_ENABLED;
+
 import static com.android.internal.R.string.config_defaultContentProtectionService;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.junit.Assert.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.app.admin.DevicePolicyManager;
 import android.content.ComponentName;
 import android.content.Context;
+import android.content.pm.PackageManager;
+import android.content.pm.UserInfo;
+import android.os.UserHandle;
+import android.os.UserManager;
+import android.platform.test.flag.junit.SetFlagsRule;
 import android.provider.DeviceConfig;
 import android.view.contentcapture.ContentCaptureManager;
 
 import com.android.settings.testutils.shadow.ShadowDeviceConfig;
 
 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;
 
+import java.util.List;
+
 @RunWith(RobolectricTestRunner.class)
-@Config(
+@Config(shadows = {ShadowDeviceConfig.class})
-        shadows = {
-            ShadowDeviceConfig.class,
-        })
 public class ContentProtectionPreferenceUtilsTest {
+
     private static final String PACKAGE_NAME = "com.test.package";
 
     private static final ComponentName COMPONENT_NAME =
             new ComponentName(PACKAGE_NAME, "TestClass");
 
-    private String mConfigDefaultContentProtectionService = COMPONENT_NAME.flattenToString();
+    private static final UserHandle USER_HANDLE = UserHandle.of(111);
+
+    private static final int PROCESS_USER_ID = 222;
+
+    private final String mConfigDefaultContentProtectionService = COMPONENT_NAME.flattenToString();
+
+    @Rule public final MockitoRule mMockitoRule = MockitoJUnit.rule();
+
+    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
     @Mock private Context mMockContext;
 
-    @Before
+    @Mock private Context mMockUserContext;
-    public void setUp() {
+
-        MockitoAnnotations.initMocks(this);
+    @Mock private UserManager mMockUserManager;
-    }
+
+    @Mock private DevicePolicyManager mMockDevicePolicyManager;
+
+    @Mock private UserInfo mMockUserInfo;
 
     @After
     public void tearDown() {
@@ -134,7 +160,6 @@ public class ContentProtectionPreferenceUtilsTest {
         assertThat(ContentProtectionPreferenceUtils.isAvailable(mMockContext)).isFalse();
     }
 
-
     @Test
     public void isAvailable_bothDisabled_false() {
         DeviceConfig.setProperty(
@@ -145,4 +170,113 @@ public class ContentProtectionPreferenceUtilsTest {
 
         assertThat(ContentProtectionPreferenceUtils.isAvailable(mMockContext)).isFalse();
     }
+
+    @Test
+    public void getManagedProfile_noProfiles() {
+        when(mMockContext.getSystemService(UserManager.class)).thenReturn(mMockUserManager);
+        when(mMockUserManager.getUserProfiles()).thenReturn(List.of());
+
+        UserHandle actual = ContentProtectionPreferenceUtils.getManagedProfile(mMockContext);
+
+        assertThat(actual).isNull();
+    }
+
+    @Test
+    public void getManagedProfile_notManaged() {
+        when(mMockContext.getSystemService(UserManager.class)).thenReturn(mMockUserManager);
+        when(mMockUserManager.getUserProfiles()).thenReturn(List.of(USER_HANDLE));
+        when(mMockUserManager.getProcessUserId()).thenReturn(PROCESS_USER_ID);
+        when(mMockUserManager.getUserInfo(USER_HANDLE.getIdentifier())).thenReturn(mMockUserInfo);
+
+        UserHandle actual = ContentProtectionPreferenceUtils.getManagedProfile(mMockContext);
+
+        assertThat(actual).isNull();
+        verify(mMockUserInfo).isManagedProfile();
+    }
+
+    @Test
+    public void getManagedProfile_managed() {
+        when(mMockContext.getSystemService(UserManager.class)).thenReturn(mMockUserManager);
+        when(mMockUserManager.getUserProfiles()).thenReturn(List.of(USER_HANDLE));
+        when(mMockUserManager.getProcessUserId()).thenReturn(PROCESS_USER_ID);
+        when(mMockUserManager.getUserInfo(USER_HANDLE.getIdentifier())).thenReturn(mMockUserInfo);
+        when(mMockUserInfo.isManagedProfile()).thenReturn(true);
+
+        UserHandle actual = ContentProtectionPreferenceUtils.getManagedProfile(mMockContext);
+
+        assertThat(actual).isEqualTo(USER_HANDLE);
+    }
+
+    @Test
+    public void getContentProtectionPolicy_flagDisabled_managedProfileNull() {
+        mSetFlagsRule.disableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+
+        int actual =
+                ContentProtectionPreferenceUtils.getContentProtectionPolicy(
+                        mMockContext, /* managedProfile= */ null);
+
+        assertThat(actual).isEqualTo(DevicePolicyManager.CONTENT_PROTECTION_DISABLED);
+    }
+
+    @Test
+    public void getContentProtectionPolicy_flagDisabled_managedProfileNotNull() {
+        mSetFlagsRule.disableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+
+        int actual =
+                ContentProtectionPreferenceUtils.getContentProtectionPolicy(
+                        mMockContext, USER_HANDLE);
+
+        assertThat(actual).isEqualTo(DevicePolicyManager.CONTENT_PROTECTION_DISABLED);
+    }
+
+    @Test
+    public void getContentProtectionPolicy_flagEnabled_managedProfileNull() throws Exception {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        when(mMockContext.getSystemService(DevicePolicyManager.class))
+                .thenReturn(mMockDevicePolicyManager);
+        when(mMockDevicePolicyManager.getContentProtectionPolicy(/* admin= */ null))
+                .thenReturn(DevicePolicyManager.CONTENT_PROTECTION_ENABLED);
+
+        int actual =
+                ContentProtectionPreferenceUtils.getContentProtectionPolicy(
+                        mMockContext, /* managedProfile= */ null);
+
+        assertThat(actual).isEqualTo(DevicePolicyManager.CONTENT_PROTECTION_ENABLED);
+        verify(mMockContext, never()).createPackageContextAsUser(anyString(), anyInt(), any());
+    }
+
+    @Test
+    public void getContentProtectionPolicy_flagEnabled_managedProfileNotNull() throws Exception {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        when(mMockContext.getPackageName()).thenReturn(PACKAGE_NAME);
+        when(mMockContext.createPackageContextAsUser(PACKAGE_NAME, /* flags= */ 0, USER_HANDLE))
+                .thenReturn(mMockUserContext);
+        when(mMockUserContext.getSystemService(DevicePolicyManager.class))
+                .thenReturn(mMockDevicePolicyManager);
+        when(mMockDevicePolicyManager.getContentProtectionPolicy(/* admin= */ null))
+                .thenReturn(DevicePolicyManager.CONTENT_PROTECTION_ENABLED);
+
+        int actual =
+                ContentProtectionPreferenceUtils.getContentProtectionPolicy(
+                        mMockContext, USER_HANDLE);
+
+        assertThat(actual).isEqualTo(DevicePolicyManager.CONTENT_PROTECTION_ENABLED);
+    }
+
+    @Test
+    public void getContentProtectionPolicy_flagEnabled_managedProfileNotNull_nameNotFound()
+            throws Exception {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        when(mMockContext.getPackageName()).thenReturn(PACKAGE_NAME);
+        when(mMockContext.createPackageContextAsUser(PACKAGE_NAME, /* flags= */ 0, USER_HANDLE))
+                .thenThrow(new PackageManager.NameNotFoundException());
+
+        assertThrows(
+                IllegalStateException.class,
+                () ->
+                        ContentProtectionPreferenceUtils.getContentProtectionPolicy(
+                                mMockContext, USER_HANDLE));
+
+        verify(mMockContext, never()).getSystemService(DevicePolicyManager.class);
+    }
 }

tests/robotests/src/com/android/settings/security/ContentProtectionTogglePreferenceControllerTest.java

@@ -15,20 +15,25 @@
  */
 package com.android.settings.security;
 
+import static android.view.contentprotection.flags.Flags.FLAG_MANAGE_DEVICE_POLICY_ENABLED;
+
+import static com.android.settings.core.BasePreferenceController.AVAILABLE;
 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.app.admin.DevicePolicyManager;
 import android.content.Context;
+import android.os.UserHandle;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.provider.Settings;
 
+import androidx.annotation.Nullable;
 import androidx.preference.PreferenceScreen;
 import androidx.test.core.app.ApplicationProvider;
 
@@ -42,37 +47,40 @@ 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(
+@Config(shadows = {ShadowUtils.class})
-        shadows = {
-            ShadowUtils.class,
-        })
 public class ContentProtectionTogglePreferenceControllerTest {
 
-    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
     @Rule public final MockitoRule mMockitoRule = MockitoJUnit.rule();
 
-    @Mock private PreferenceScreen mMockScreen;
+    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
-    private RestrictedLockUtils.EnforcedAdmin mAdmin;
-    private SettingsMainSwitchPreference mSwitchPreference;
     private final Context mContext = ApplicationProvider.getApplicationContext();
-    private ContentProtectionTogglePreferenceController mController;
+
+    @Mock private PreferenceScreen mMockPreferenceScreen;
+
+    @Mock private SettingsMainSwitchPreference mMockSwitchPreference;
+
+    @Nullable private RestrictedLockUtils.EnforcedAdmin mEnforcedAdmin;
+
+    @DevicePolicyManager.ContentProtectionPolicy
+    private int mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_DISABLED;
+
+    private TestContentProtectionTogglePreferenceController mController;
+
     private int mSettingBackupValue;
 
     @Before
     public void setUp() {
-        MockitoAnnotations.initMocks(this);
         mController = new TestContentProtectionTogglePreferenceController();
-        mSwitchPreference = new SettingsMainSwitchPreference(mContext);
+        SettingsMainSwitchPreference switchPreference = new SettingsMainSwitchPreference(mContext);
-        when(mMockScreen.findPreference(mController.getPreferenceKey()))
+        when(mMockPreferenceScreen.findPreference(mController.getPreferenceKey()))
-                .thenReturn(mSwitchPreference);
+                .thenReturn(switchPreference);
         mSettingBackupValue = getContentProtectionGlobalSetting();
         Settings.Global.putInt(mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, 0);
     }
@@ -87,89 +95,223 @@ public class ContentProtectionTogglePreferenceControllerTest {
     }
 
     @Test
-    public void isAvailable_alwaysAvailable() {
+    public void constructor_flagDisabled_doesNotFetchData() {
+        mSetFlagsRule.disableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mController = new TestContentProtectionTogglePreferenceController();
+
+        assertThat(mController.mCounterGetManagedProfile).isEqualTo(0);
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(0);
+        assertThat(mController.mCounterGetContentProtectionPolicy).isEqualTo(0);
+    }
+
+    @Test
+    public void constructor_flagEnabled_fetchesData() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mController = new TestContentProtectionTogglePreferenceController();
+
+        assertThat(mController.mCounterGetManagedProfile).isEqualTo(1);
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(1);
+        assertThat(mController.mCounterGetContentProtectionPolicy).isEqualTo(1);
+    }
+
+    @Test
+    public void getAvailabilityStatus_available() {
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(AVAILABLE);
         assertThat(mController.isAvailable()).isTrue();
     }
 
     @Test
-    public void displayPreference() {
+    public void isChecked_noEnforcedAdmin_readsSettingsTrue() {
-        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);
 
         assertThat(mController.isChecked()).isTrue();
     }
 
     @Test
-    public void isChecked_fullyManagedMode_settingTurnOff() {
+    public void isChecked_noEnforcedAdmin_readsSettingsFalse() {
-        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());
-
-        mController = new TestContentProtectionTogglePreferenceController();
-        mController.displayPreference(mMockScreen);
-        mController.updateState(mockSwitchPreference);
-
-        assertThat(mController.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() {
+    public void isChecked_noEnforcedAdmin_readsSettingsDefaultTrue() {
         assertThat(mController.isChecked()).isTrue();
-        assertThat(getContentProtectionGlobalSetting()).isEqualTo(0);
+    }
+
+    @Test
+    public void isChecked_enforcedAdmin_flagDisabled_false() {
+        mSetFlagsRule.disableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mEnforcedAdmin = new RestrictedLockUtils.EnforcedAdmin();
+        Settings.Global.putInt(mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, 1);
+        setupForUpdateState();
+        mController.updateState(mMockSwitchPreference);
+
+        assertThat(mController.isChecked()).isFalse();
+    }
+
+    @Test
+    public void isChecked_enforcedAdmin_flagEnabled_policyDisabled_false() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mEnforcedAdmin = new RestrictedLockUtils.EnforcedAdmin();
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_DISABLED;
+        Settings.Global.putInt(mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, 1);
+        mController = new TestContentProtectionTogglePreferenceController();
+
+        assertThat(mController.isChecked()).isFalse();
+    }
+
+    @Test
+    public void isChecked_enforcedAdmin_flagEnabled_policyEnabled_true() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mEnforcedAdmin = new RestrictedLockUtils.EnforcedAdmin();
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_ENABLED;
+        Settings.Global.putInt(
+                mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, -1);
+        mController = new TestContentProtectionTogglePreferenceController();
+
+        assertThat(mController.isChecked()).isTrue();
+    }
+
+    @Test
+    public void isChecked_enforcedAdmin_flagEnabled_policyNotControlled_readsSettingsTrue() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mEnforcedAdmin = new RestrictedLockUtils.EnforcedAdmin();
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_NOT_CONTROLLED_BY_POLICY;
+        Settings.Global.putInt(mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, 1);
+        mController = new TestContentProtectionTogglePreferenceController();
+
+        assertThat(mController.isChecked()).isTrue();
+    }
+
+    @Test
+    public void isChecked_enforcedAdmin_flagEnabled_policyNotControlled_readsSettingsFalse() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mEnforcedAdmin = new RestrictedLockUtils.EnforcedAdmin();
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_NOT_CONTROLLED_BY_POLICY;
+        Settings.Global.putInt(
+                mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, -1);
+        mController = new TestContentProtectionTogglePreferenceController();
+
+        assertThat(mController.isChecked()).isFalse();
+    }
+
+    @Test
+    public void isChecked_enforcedAdmin_flagEnabled_policyNotControlled_readsSettingsDefaultTrue() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mEnforcedAdmin = new RestrictedLockUtils.EnforcedAdmin();
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_NOT_CONTROLLED_BY_POLICY;
+        mController = new TestContentProtectionTogglePreferenceController();
+
+        assertThat(mController.isChecked()).isTrue();
+    }
+
+    @Test
+    public void displayPreference() {
+        setupForDisplayPreference();
+
+        mController.displayPreference(mMockPreferenceScreen);
+
+        verify(mMockSwitchPreference).addOnSwitchChangeListener(mController);
+    }
+
+    @Test
+    public void updateState_flagDisabled_noEnforcedAdmin() {
+        mSetFlagsRule.disableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        setupForUpdateState();
+
+        mController.updateState(mMockSwitchPreference);
+
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(1);
+        verify(mMockSwitchPreference, never()).setDisabledByAdmin(any());
+    }
+
+    @Test
+    public void updateState_flagDisabled_enforcedAdmin() {
+        mSetFlagsRule.disableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mEnforcedAdmin = new RestrictedLockUtils.EnforcedAdmin();
+        setupForUpdateState();
+
+        mController.updateState(mMockSwitchPreference);
+
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(1);
+        verify(mMockSwitchPreference).setDisabledByAdmin(mEnforcedAdmin);
+    }
+
+    @Test
+    public void updateState_flagEnabled_noEnforcedAdmin_policyDisabled() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_DISABLED;
+        setupForUpdateState();
+
+        mController.updateState(mMockSwitchPreference);
+
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(1);
+        verify(mMockSwitchPreference, never()).setDisabledByAdmin(any());
+    }
+
+    @Test
+    public void updateState_flagEnabled_noEnforcedAdmin_policyEnabled() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_ENABLED;
+        setupForUpdateState();
+
+        mController.updateState(mMockSwitchPreference);
+
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(1);
+        verify(mMockSwitchPreference, never()).setDisabledByAdmin(any());
+    }
+
+    @Test
+    public void updateState_flagEnabled_noEnforcedAdmin_policyNotControlled() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_NOT_CONTROLLED_BY_POLICY;
+        setupForUpdateState();
+
+        mController.updateState(mMockSwitchPreference);
+
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(1);
+        verify(mMockSwitchPreference, never()).setDisabledByAdmin(any());
+    }
+
+    @Test
+    public void updateState_flagEnabled_enforcedAdmin_policyDisabled() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mEnforcedAdmin = new RestrictedLockUtils.EnforcedAdmin();
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_DISABLED;
+        setupForUpdateState();
+
+        mController.updateState(mMockSwitchPreference);
+
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(1);
+        verify(mMockSwitchPreference).setDisabledByAdmin(mEnforcedAdmin);
+    }
+
+    @Test
+    public void updateState_flagEnabled_enforcedAdmin_policyEnabled() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mEnforcedAdmin = new RestrictedLockUtils.EnforcedAdmin();
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_ENABLED;
+        setupForUpdateState();
+
+        mController.updateState(mMockSwitchPreference);
+
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(1);
+        verify(mMockSwitchPreference).setDisabledByAdmin(mEnforcedAdmin);
+    }
+
+    @Test
+    public void updateState_flagEnabled_enforcedAdmin_policyNotControlled() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mEnforcedAdmin = new RestrictedLockUtils.EnforcedAdmin();
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_NOT_CONTROLLED_BY_POLICY;
+        setupForUpdateState();
+
+        mController.updateState(mMockSwitchPreference);
+
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(1);
+        verify(mMockSwitchPreference, never()).setDisabledByAdmin(any());
     }
 
     @Test
@@ -193,20 +335,49 @@ public class ContentProtectionTogglePreferenceControllerTest {
                 mContext.getContentResolver(), KEY_CONTENT_PROTECTION_PREFERENCE, 0);
     }
 
-    private void setUpFullyManagedMode() {
+    private void setupForDisplayPreference() {
-        mAdmin = new RestrictedLockUtils.EnforcedAdmin();
+        when(mMockPreferenceScreen.findPreference(any())).thenReturn(mMockSwitchPreference);
+        when(mMockSwitchPreference.getKey()).thenReturn(mController.getPreferenceKey());
+        mController = new TestContentProtectionTogglePreferenceController();
+    }
+
+    private void setupForUpdateState() {
+        setupForDisplayPreference();
+        mController.displayPreference(mMockPreferenceScreen);
     }
 
     private class TestContentProtectionTogglePreferenceController
             extends ContentProtectionTogglePreferenceController {
 
+        public int mCounterGetManagedProfile;
+
+        public int mCounterGetEnforcedAdmin;
+
+        public int mCounterGetContentProtectionPolicy;
+
         TestContentProtectionTogglePreferenceController() {
             super(ContentProtectionTogglePreferenceControllerTest.this.mContext, "key");
         }
 
         @Override
+        @Nullable
+        protected UserHandle getManagedProfile() {
+            mCounterGetManagedProfile++;
+            return null;
+        }
+
+        @Override
+        @Nullable
         protected RestrictedLockUtils.EnforcedAdmin getEnforcedAdmin() {
-            return mAdmin;
+            mCounterGetEnforcedAdmin++;
+            return mEnforcedAdmin;
+        }
+
+        @Override
+        @DevicePolicyManager.ContentProtectionPolicy
+        protected int getContentProtectionPolicy(@Nullable UserHandle userHandle) {
+            mCounterGetContentProtectionPolicy++;
+            return mContentProtectionPolicy;
         }
     }
 }

tests/robotests/src/com/android/settings/security/ContentProtectionWorkSwitchControllerTest.java

@@ -16,19 +16,22 @@
 
 package com.android.settings.security;
 
+import static android.view.contentprotection.flags.Flags.FLAG_MANAGE_DEVICE_POLICY_ENABLED;
+
 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.app.admin.DevicePolicyManager;
 import android.content.Context;
 import android.os.UserHandle;
+import android.platform.test.flag.junit.SetFlagsRule;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -39,55 +42,169 @@ import com.android.settingslib.RestrictedLockUtils;
 import com.android.settingslib.RestrictedSwitchPreference;
 
 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;
 
 @RunWith(RobolectricTestRunner.class)
 public class ContentProtectionWorkSwitchControllerTest {
+
     private static final UserHandle TEST_USER_HANDLE = UserHandle.of(10);
 
+    @Rule public final MockitoRule mMockitoRule = MockitoJUnit.rule();
+
+    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
     private final Context mContext = ApplicationProvider.getApplicationContext();
 
     @Mock private PreferenceScreen mMockPreferenceScreen;
-    private ContentProtectionWorkSwitchController mController;
+
-    private UserHandle mManagedProfileUserHandle;
+    @Mock private RestrictedSwitchPreference mMockSwitchPreference;
-    private RestrictedLockUtils.EnforcedAdmin mEnforcedAdmin;
+
+    @Nullable private UserHandle mManagedProfileUserHandle;
+
+    @Nullable private RestrictedLockUtils.EnforcedAdmin mEnforcedAdmin;
+
+    @DevicePolicyManager.ContentProtectionPolicy
+    private int mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_DISABLED;
+
+    private TestContentProtectionWorkSwitchController mController;
 
     @Before
     public void setUp() {
-        MockitoAnnotations.initMocks(this);
         mController = new TestContentProtectionWorkSwitchController();
     }
 
     @Test
-    public void isAvailable_managedProfile_available() {
+    public void constructor_flagDisabled_doesNotFetchData() {
+        mSetFlagsRule.disableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mController = new TestContentProtectionWorkSwitchController();
+
+        assertThat(mController.mCounterGetManagedProfile).isEqualTo(0);
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(0);
+        assertThat(mController.mCounterGetContentProtectionPolicy).isEqualTo(0);
+    }
+
+    @Test
+    public void constructor_flagEnabled_fetchesManagedProfile() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mController = new TestContentProtectionWorkSwitchController();
+
+        assertThat(mController.mCounterGetManagedProfile).isEqualTo(1);
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(0);
+        assertThat(mController.mCounterGetContentProtectionPolicy).isEqualTo(0);
+    }
+
+    @Test
+    public void constructor_flagEnabled_withManagedProfile_fetchesPolicy() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
         mManagedProfileUserHandle = TEST_USER_HANDLE;
+        mController = new TestContentProtectionWorkSwitchController();
+
+        assertThat(mController.mCounterGetManagedProfile).isEqualTo(1);
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(0);
+        assertThat(mController.mCounterGetContentProtectionPolicy).isEqualTo(1);
+    }
+
+    @Test
+    public void getAvailabilityStatus_flagDisabled_managedProfile_available() {
+        mSetFlagsRule.disableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mManagedProfileUserHandle = TEST_USER_HANDLE;
+        mController = new TestContentProtectionWorkSwitchController();
 
         assertThat(mController.getAvailabilityStatus()).isEqualTo(AVAILABLE);
         assertThat(mController.isAvailable()).isTrue();
     }
 
     @Test
-    public void isAvailable_noManagedProfile_notAvailable() {
+    public void getAvailabilityStatus_flagDisabled_noManagedProfile_unavailable() {
-        mManagedProfileUserHandle = null;
+        mSetFlagsRule.disableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mController = new TestContentProtectionWorkSwitchController();
 
         assertThat(mController.getAvailabilityStatus()).isEqualTo(CONDITIONALLY_UNAVAILABLE);
         assertThat(mController.isAvailable()).isFalse();
     }
 
     @Test
-    public void isChecked_noManagedProfile_alwaysOff() {
+    public void getAvailabilityStatus_flagEnabled_managedProfile_policyDisabled_available() {
-        mManagedProfileUserHandle = null;
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mManagedProfileUserHandle = TEST_USER_HANDLE;
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_DISABLED;
+        mController = new TestContentProtectionWorkSwitchController();
+
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(AVAILABLE);
+        assertThat(mController.isAvailable()).isTrue();
+    }
+
+    @Test
+    public void getAvailabilityStatus_flagEnabled_managedProfile_policyEnabled_available() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mManagedProfileUserHandle = TEST_USER_HANDLE;
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_ENABLED;
+        mController = new TestContentProtectionWorkSwitchController();
+
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(AVAILABLE);
+        assertThat(mController.isAvailable()).isTrue();
+    }
+
+    @Test
+    public void getAvailabilityStatus_flagEnabled_managedProfile_policyNotControlled_unavailable() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mManagedProfileUserHandle = TEST_USER_HANDLE;
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_NOT_CONTROLLED_BY_POLICY;
+        mController = new TestContentProtectionWorkSwitchController();
+
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(CONDITIONALLY_UNAVAILABLE);
+        assertThat(mController.isAvailable()).isFalse();
+    }
+
+    @Test
+    public void getAvailabilityStatus_flagEnabled_noManagedProfile_unavailable() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mController = new TestContentProtectionWorkSwitchController();
+
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(CONDITIONALLY_UNAVAILABLE);
+        assertThat(mController.isAvailable()).isFalse();
+    }
+
+    @Test
+    public void isChecked_flagDisabled_false() {
+        mSetFlagsRule.disableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mController = new TestContentProtectionWorkSwitchController();
 
         assertThat(mController.isChecked()).isFalse();
     }
 
     @Test
-    public void isChecked_managedProfile_alwaysOff() {
+    public void isChecked_flagEnabled_policyEnabled_true() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
         mManagedProfileUserHandle = TEST_USER_HANDLE;
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_ENABLED;
+        mController = new TestContentProtectionWorkSwitchController();
+
+        assertThat(mController.isChecked()).isTrue();
+    }
+
+    @Test
+    public void isChecked_flagEnabled_policyDisabled_false() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mManagedProfileUserHandle = TEST_USER_HANDLE;
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_DISABLED;
+        mController = new TestContentProtectionWorkSwitchController();
+
+        assertThat(mController.isChecked()).isFalse();
+    }
+
+    @Test
+    public void isChecked_flagEnabled_policyNotControlled_false() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mManagedProfileUserHandle = TEST_USER_HANDLE;
+        mContentProtectionPolicy = DevicePolicyManager.CONTENT_PROTECTION_NOT_CONTROLLED_BY_POLICY;
+        mController = new TestContentProtectionWorkSwitchController();
 
         assertThat(mController.isChecked()).isFalse();
     }
@@ -99,50 +216,72 @@ public class ContentProtectionWorkSwitchControllerTest {
     }
 
     @Test
-    public void displayPreference_managedProfile_disabled() {
+    public void displayPreference_flagDisabled_managedProfile_disabledByAdmin() {
+        mSetFlagsRule.disableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
         mManagedProfileUserHandle = TEST_USER_HANDLE;
         mEnforcedAdmin = new RestrictedLockUtils.EnforcedAdmin();
-        RestrictedSwitchPreference mockSwitchPreference = mock(RestrictedSwitchPreference.class);
+        setupForDisplayPreference();
-        when(mMockPreferenceScreen.findPreference(any())).thenReturn(mockSwitchPreference);
-        when(mockSwitchPreference.getKey()).thenReturn(mController.getPreferenceKey());
 
         mController.displayPreference(mMockPreferenceScreen);
 
-        assertThat(mController.isAvailable()).isTrue();
+        verify(mMockSwitchPreference).setDisabledByAdmin(mEnforcedAdmin);
-        verify(mockSwitchPreference).setDisabledByAdmin(mEnforcedAdmin);
+        assertThat(mController.mCounterGetManagedProfile).isEqualTo(3);
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(1);
     }
 
     @Test
-    public void displayPreference_noManagedProfile_notDisabled() {
+    public void displayPreference_flagDisabled_noManagedProfile_notDisabledByAdmin() {
-        mManagedProfileUserHandle = null;
+        mSetFlagsRule.disableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        setupForDisplayPreference();
+
+        mController.displayPreference(mMockPreferenceScreen);
+
+        verify(mMockSwitchPreference, never()).setDisabledByAdmin(any());
+        assertThat(mController.mCounterGetManagedProfile).isEqualTo(3);
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(0);
+    }
+
+    @Test
+    public void displayPreference_flagEnabled_managedProfile_disabledByAdmin() {
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
+        mManagedProfileUserHandle = TEST_USER_HANDLE;
         mEnforcedAdmin = new RestrictedLockUtils.EnforcedAdmin();
-        RestrictedSwitchPreference mockSwitchPreference = mock(RestrictedSwitchPreference.class);
+        setupForDisplayPreference();
-        when(mMockPreferenceScreen.findPreference(any())).thenReturn(mockSwitchPreference);
-        when(mockSwitchPreference.getKey()).thenReturn(mController.getPreferenceKey());
 
         mController.displayPreference(mMockPreferenceScreen);
 
-        assertThat(mController.isAvailable()).isFalse();
+        verify(mMockSwitchPreference).setDisabledByAdmin(mEnforcedAdmin);
-        verify(mockSwitchPreference, never()).setDisabledByAdmin(any());
+        assertThat(mController.mCounterGetManagedProfile).isEqualTo(1);
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(1);
     }
 
     @Test
-    public void displayPreference_noEnforcedAdmin_notDisabled() {
+    public void displayPreference_flagEnabled_noManagedProfile_notDisabledByAdmin() {
-        mManagedProfileUserHandle = null;
+        mSetFlagsRule.enableFlags(FLAG_MANAGE_DEVICE_POLICY_ENABLED);
-        mEnforcedAdmin = null;
+        setupForDisplayPreference();
-        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(mMockSwitchPreference, never()).setDisabledByAdmin(any());
-        verify(mockSwitchPreference, never()).setDisabledByAdmin(any());
+        assertThat(mController.mCounterGetManagedProfile).isEqualTo(1);
+        assertThat(mController.mCounterGetEnforcedAdmin).isEqualTo(0);
+    }
+
+    private void setupForDisplayPreference() {
+        when(mMockPreferenceScreen.findPreference(any())).thenReturn(mMockSwitchPreference);
+        when(mMockSwitchPreference.getKey()).thenReturn(mController.getPreferenceKey());
+        mController = new TestContentProtectionWorkSwitchController();
     }
 
     private class TestContentProtectionWorkSwitchController
             extends ContentProtectionWorkSwitchController {
 
+        public int mCounterGetManagedProfile;
+
+        public int mCounterGetEnforcedAdmin;
+
+        public int mCounterGetContentProtectionPolicy;
+
         TestContentProtectionWorkSwitchController() {
             super(ContentProtectionWorkSwitchControllerTest.this.mContext, "key");
         }
@@ -150,14 +289,23 @@ public class ContentProtectionWorkSwitchControllerTest {
         @Override
         @Nullable
         protected UserHandle getManagedProfile() {
+            mCounterGetManagedProfile++;
             return mManagedProfileUserHandle;
         }
 
         @Override
         @Nullable
         protected RestrictedLockUtils.EnforcedAdmin getEnforcedAdmin(
-                @NonNull UserHandle managedProfile) {
+                @NonNull UserHandle userHandle) {
+            mCounterGetEnforcedAdmin++;
             return mEnforcedAdmin;
         }
+
+        @Override
+        @DevicePolicyManager.ContentProtectionPolicy
+        protected int getContentProtectionPolicy(@Nullable UserHandle userHandle) {
+            mCounterGetContentProtectionPolicy++;
+            return mContentProtectionPolicy;
+        }
     }
 }