Add Always on VPN to Privacy Settings page

This CL adds information about always-on VPNs to the Enterprise
Privacy Settings page.

Test: make RunSettingsRoboTests
Bug: 32692748

Change-Id: I2b59e2ec4c55308b323aaa478cd9c847fe0b4b55
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 73c7e3b..6caaf16 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -8003,6 +8003,12 @@
     <string name="enterprise_privacy_security_logs">Your most recent security log</string>
     <!-- Label indicating that the date at which the admin last took a particular action was "never" (i.e. the admin never took the action so far). -->
     <string name="enterprise_privacy_never">Never</string>
+    <!-- Label explaining that an always-on VPN was set by the admin for the entire device. [CHAR LIMIT=NONE] -->
+    <string name="enterprise_privacy_always_on_vpn_device">Always-on VPN turned on</string>
+    <!-- Label explaining that an always-on VPN was set by the admin in the personal profile. [CHAR LIMIT=NONE] -->
+    <string name="enterprise_privacy_always_on_vpn_personal">Always-on VPN turned on in your personal profile</string>
+    <!-- Label explaining that an always-on VPN was set by the admin in the work profile. [CHAR LIMIT=NONE] -->
+    <string name="enterprise_privacy_always_on_vpn_work">Always-on VPN turned on in your work profile</string>
 
     <!-- Preference label for the Photos & Videos storage section. [CHAR LIMIT=50] -->
     <string name="storage_photos_videos">Photos &amp; Videos</string>
diff --git a/res/xml/enterprise_privacy_settings.xml b/res/xml/enterprise_privacy_settings.xml
index e1761e2..856706e 100644
--- a/res/xml/enterprise_privacy_settings.xml
+++ b/res/xml/enterprise_privacy_settings.xml
@@ -59,6 +59,15 @@
     </PreferenceCategory>
 
     <PreferenceCategory android:title="@string/enterprise_privacy_exposure_changes_category">
+        <com.android.settings.DividerPreference
+                android:key="always_on_vpn_primary_user"
+                settings:allowDividerBelow="true"
+                settings:multiLine="true"/>
+        <com.android.settings.DividerPreference
+                android:key="always_on_vpn_managed_profile"
+                android:title="@string/enterprise_privacy_always_on_vpn_work"
+                settings:allowDividerBelow="true"
+                settings:multiLine="true"/>
     </PreferenceCategory>
 
     <PreferenceCategory android:title="@string/enterprise_privacy_device_access_category">
diff --git a/src/com/android/settings/enterprise/AlwaysOnVpnManagedProfilePreferenceController.java b/src/com/android/settings/enterprise/AlwaysOnVpnManagedProfilePreferenceController.java
new file mode 100644
index 0000000..52625ec
--- /dev/null
+++ b/src/com/android/settings/enterprise/AlwaysOnVpnManagedProfilePreferenceController.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright (C) 2017 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.enterprise;
+
+import android.content.Context;
+import android.support.v7.preference.Preference;
+
+import com.android.settings.core.PreferenceController;
+import com.android.settings.overlay.FeatureFactory;
+
+public class AlwaysOnVpnManagedProfilePreferenceController extends PreferenceController {
+
+    private static final String KEY_ALWAYS_ON_VPN_MANAGED_PROFILE = "always_on_vpn_managed_profile";
+    private final EnterprisePrivacyFeatureProvider mFeatureProvider;
+
+    public AlwaysOnVpnManagedProfilePreferenceController(Context context) {
+        super(context);
+        mFeatureProvider = FeatureFactory.getFactory(context)
+                .getEnterprisePrivacyFeatureProvider(context);
+    }
+
+    @Override
+    public void updateState(Preference preference) {
+        preference.setVisible(mFeatureProvider.isAlwaysOnVpnSetInManagedProfile());
+    }
+
+    @Override
+    public boolean isAvailable() {
+        return true;
+    }
+
+    @Override
+    public String getPreferenceKey() {
+        return KEY_ALWAYS_ON_VPN_MANAGED_PROFILE;
+    }
+}
diff --git a/src/com/android/settings/enterprise/AlwaysOnVpnPrimaryUserPreferenceController.java b/src/com/android/settings/enterprise/AlwaysOnVpnPrimaryUserPreferenceController.java
new file mode 100644
index 0000000..c7ffeaf
--- /dev/null
+++ b/src/com/android/settings/enterprise/AlwaysOnVpnPrimaryUserPreferenceController.java
@@ -0,0 +1,51 @@
+/*
+ * Copyright (C) 2017 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.enterprise;
+
+import android.content.Context;
+import android.support.v7.preference.Preference;
+
+import com.android.settings.R;
+import com.android.settings.core.PreferenceController;
+import com.android.settings.overlay.FeatureFactory;
+
+public class AlwaysOnVpnPrimaryUserPreferenceController extends PreferenceController {
+
+    private static final String KEY_ALWAYS_ON_VPN_PRIMARY_USER = "always_on_vpn_primary_user";
+    private final EnterprisePrivacyFeatureProvider mFeatureProvider;
+
+    public AlwaysOnVpnPrimaryUserPreferenceController(Context context) {
+        super(context);
+        mFeatureProvider = FeatureFactory.getFactory(context)
+                .getEnterprisePrivacyFeatureProvider(context);
+    }
+
+    @Override
+    public void updateState(Preference preference) {
+        preference.setTitle(mFeatureProvider.isInCompMode()
+                ? R.string.enterprise_privacy_always_on_vpn_personal
+                : R.string.enterprise_privacy_always_on_vpn_device);
+        preference.setVisible(mFeatureProvider.isAlwaysOnVpnSetInPrimaryUser());
+    }
+
+    @Override
+    public boolean isAvailable() {
+        return true;
+    }
+
+    @Override
+    public String getPreferenceKey() {
+        return KEY_ALWAYS_ON_VPN_PRIMARY_USER;
+    }
+}
diff --git a/src/com/android/settings/enterprise/EnterprisePrivacyFeatureProvider.java b/src/com/android/settings/enterprise/EnterprisePrivacyFeatureProvider.java
index efc02d6..dec2d80 100644
--- a/src/com/android/settings/enterprise/EnterprisePrivacyFeatureProvider.java
+++ b/src/com/android/settings/enterprise/EnterprisePrivacyFeatureProvider.java
@@ -26,6 +26,12 @@
     boolean hasDeviceOwner();
 
     /**
+     * Returns whether the device is in COMP mode (primary user managed by a Device Owner app and
+     * work profile managed by a Profile Owner app).
+     */
+    boolean isInCompMode();
+
+    /**
      * Returns the time at which the Device Owner last retrieved security logs, or {@code null} if
      * logs were never retrieved by the Device Owner on this device.
      */
@@ -42,4 +48,14 @@
      * logs were never retrieved by the Device Owner on this device.
      */
     Date getLastNetworkLogRetrievalTime();
+
+    /**
+     * Returns whether the Device Owner in the primary user set an always-on VPN.
+     */
+    boolean isAlwaysOnVpnSetInPrimaryUser();
+
+    /**
+     * Returns whether the Profile Owner in the managed profile (if any) set an always-on VPN.
+     */
+    boolean isAlwaysOnVpnSetInManagedProfile();
 }
diff --git a/src/com/android/settings/enterprise/EnterprisePrivacyFeatureProviderImpl.java b/src/com/android/settings/enterprise/EnterprisePrivacyFeatureProviderImpl.java
index 2e8b7f6..a742cc3 100644
--- a/src/com/android/settings/enterprise/EnterprisePrivacyFeatureProviderImpl.java
+++ b/src/com/android/settings/enterprise/EnterprisePrivacyFeatureProviderImpl.java
@@ -17,20 +17,32 @@
 package com.android.settings.enterprise;
 
 import android.content.pm.PackageManager;
+import android.content.pm.UserInfo;
+import android.os.UserHandle;
+import android.os.UserManager;
 
 import com.android.settings.applications.PackageManagerWrapper;
+import com.android.settings.vpn2.ConnectivityManagerWrapper;
+import com.android.settings.vpn2.VpnUtils;
 
 import java.util.Date;
+import java.util.List;
 
 public class EnterprisePrivacyFeatureProviderImpl implements EnterprisePrivacyFeatureProvider {
 
     private final DevicePolicyManagerWrapper mDpm;
     private final PackageManagerWrapper mPm;
+    private final UserManager mUm;
+    private final ConnectivityManagerWrapper mCm;
+
+    private static final int MY_USER_ID = UserHandle.myUserId();
 
     public EnterprisePrivacyFeatureProviderImpl(DevicePolicyManagerWrapper dpm,
-            PackageManagerWrapper pm) {
+            PackageManagerWrapper pm, UserManager um, ConnectivityManagerWrapper cm) {
         mDpm = dpm;
         mPm = pm;
+        mUm = um;
+        mCm = cm;
     }
 
     @Override
@@ -41,19 +53,47 @@
         return mDpm.getDeviceOwnerComponentOnAnyUser() != null;
     }
 
+    private int getManagedProfileUserId() {
+        for (final UserInfo userInfo : mUm.getProfiles(MY_USER_ID)) {
+            if (userInfo.isManagedProfile()) {
+                return userInfo.id;
+            }
+        }
+        return -1;
+    }
+
+    @Override
+    public boolean isInCompMode() {
+        return hasDeviceOwner() && getManagedProfileUserId() != -1;
+    }
+
     @Override
     public Date getLastSecurityLogRetrievalTime() {
         final long timestamp = mDpm.getLastSecurityLogRetrievalTime();
         return timestamp < 0 ? null : new Date(timestamp);
     }
 
+    @Override
     public Date getLastBugReportRequestTime() {
         final long timestamp = mDpm.getLastBugReportRequestTime();
         return timestamp < 0 ? null : new Date(timestamp);
     }
 
+    @Override
     public Date getLastNetworkLogRetrievalTime() {
         final long timestamp = mDpm.getLastNetworkLogRetrievalTime();
         return timestamp < 0 ? null : new Date(timestamp);
     }
+
+    @Override
+    public boolean isAlwaysOnVpnSetInPrimaryUser() {
+        return VpnUtils.isAlwaysOnVpnSet(mCm, MY_USER_ID);
+    }
+
+    @Override
+    public boolean isAlwaysOnVpnSetInManagedProfile() {
+        final int managedProfileUserId = getManagedProfileUserId();
+        return managedProfileUserId != -1 &&
+                VpnUtils.isAlwaysOnVpnSet(mCm, managedProfileUserId);
+    }
 }
diff --git a/src/com/android/settings/enterprise/EnterprisePrivacySettings.java b/src/com/android/settings/enterprise/EnterprisePrivacySettings.java
index 91d3a65..478f7fe 100644
--- a/src/com/android/settings/enterprise/EnterprisePrivacySettings.java
+++ b/src/com/android/settings/enterprise/EnterprisePrivacySettings.java
@@ -61,6 +61,8 @@
         controllers.add(new NetworkLogsPreferenceController(context));
         controllers.add(new BugReportsPreferenceController(context));
         controllers.add(new SecurityLogsPreferenceController(context));
+        controllers.add(new AlwaysOnVpnPrimaryUserPreferenceController(context));
+        controllers.add(new AlwaysOnVpnManagedProfilePreferenceController(context));
         return controllers;
     }
 
diff --git a/src/com/android/settings/overlay/FeatureFactoryImpl.java b/src/com/android/settings/overlay/FeatureFactoryImpl.java
index eb5d065..060b58c 100644
--- a/src/com/android/settings/overlay/FeatureFactoryImpl.java
+++ b/src/com/android/settings/overlay/FeatureFactoryImpl.java
@@ -18,6 +18,8 @@
 
 import android.app.admin.DevicePolicyManager;
 import android.content.Context;
+import android.net.ConnectivityManager;
+import android.os.UserManager;
 import android.support.annotation.Keep;
 
 import com.android.settings.applications.ApplicationFeatureProvider;
@@ -37,6 +39,7 @@
 import com.android.settings.security.SecurityFeatureProviderImpl;
 import com.android.settings.search2.SearchFeatureProvider;
 import com.android.settings.search2.SearchFeatureProviderImpl;
+import com.android.settings.vpn2.ConnectivityManagerWrapperImpl;
 
 /**
  * {@link FeatureFactory} implementation for AOSP Settings.
@@ -101,7 +104,10 @@
             mEnterprisePrivacyFeatureProvider = new EnterprisePrivacyFeatureProviderImpl(
                     new DevicePolicyManagerWrapperImpl((DevicePolicyManager) context
                             .getSystemService(Context.DEVICE_POLICY_SERVICE)),
-                    new PackageManagerWrapperImpl(context.getPackageManager()));
+                    new PackageManagerWrapperImpl(context.getPackageManager()),
+                    UserManager.get(context),
+                    new ConnectivityManagerWrapperImpl((ConnectivityManager) context
+                            .getSystemService(Context.CONNECTIVITY_SERVICE)));
         }
         return mEnterprisePrivacyFeatureProvider;
     }
diff --git a/src/com/android/settings/vpn2/ConnectivityManagerWrapper.java b/src/com/android/settings/vpn2/ConnectivityManagerWrapper.java
new file mode 100644
index 0000000..938db50
--- /dev/null
+++ b/src/com/android/settings/vpn2/ConnectivityManagerWrapper.java
@@ -0,0 +1,33 @@
+/*
+ * Copyright (C) 2017 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.vpn2;
+
+/**
+ * This interface replicates a subset of the android.net.ConnectivityManager (CM). The interface
+ * exists so that we can use a thin wrapper around the CM in production code and a mock in tests.
+ * We cannot directly mock or shadow the CM, because some of the methods we rely on are marked as
+ * hidden and are thus invisible to Robolectric.
+ */
+public interface ConnectivityManagerWrapper {
+
+    /**
+     * Calls {@code ConnectivityManager.getAlwaysOnVpnPackageForUser()}.
+     *
+     * @see android.net.ConnectivityManager#getAlwaysOnVpnPackageForUser
+     */
+   String getAlwaysOnVpnPackageForUser(int userId);
+}
diff --git a/src/com/android/settings/vpn2/ConnectivityManagerWrapperImpl.java b/src/com/android/settings/vpn2/ConnectivityManagerWrapperImpl.java
new file mode 100644
index 0000000..ad1b4eb
--- /dev/null
+++ b/src/com/android/settings/vpn2/ConnectivityManagerWrapperImpl.java
@@ -0,0 +1,33 @@
+/*
+ * Copyright (C) 2017 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.vpn2;
+
+import android.net.ConnectivityManager;
+
+public class ConnectivityManagerWrapperImpl implements ConnectivityManagerWrapper {
+
+    private final ConnectivityManager mCm;
+
+    public ConnectivityManagerWrapperImpl(ConnectivityManager cm) {
+        mCm = cm;
+    }
+
+    @Override
+    public String getAlwaysOnVpnPackageForUser(int userId) {
+        return mCm.getAlwaysOnVpnPackageForUser(userId);
+    }
+}
diff --git a/src/com/android/settings/vpn2/VpnUtils.java b/src/com/android/settings/vpn2/VpnUtils.java
index 07e6c52..c9f971d 100644
--- a/src/com/android/settings/vpn2/VpnUtils.java
+++ b/src/com/android/settings/vpn2/VpnUtils.java
@@ -82,4 +82,8 @@
         return IConnectivityManager.Stub.asInterface(
                 ServiceManager.getService(Context.CONNECTIVITY_SERVICE));
     }
+
+    public static boolean isAlwaysOnVpnSet(ConnectivityManagerWrapper cm, final int userId) {
+        return cm.getAlwaysOnVpnPackageForUser(userId) != null;
+    }
 }
diff --git a/tests/robotests/src/com/android/settings/enterprise/AlwaysOnVpnManagedProfilePreferenceControllerTest.java b/tests/robotests/src/com/android/settings/enterprise/AlwaysOnVpnManagedProfilePreferenceControllerTest.java
new file mode 100644
index 0000000..a52e049
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/enterprise/AlwaysOnVpnManagedProfilePreferenceControllerTest.java
@@ -0,0 +1,89 @@
+/*
+ * Copyright (C) 2017 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.enterprise;
+
+import android.content.Context;
+import android.support.v7.preference.Preference;
+
+import com.android.settings.SettingsRobolectricTestRunner;
+import com.android.settings.TestConfig;
+import com.android.settings.testutils.FakeFeatureFactory;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Answers;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.annotation.Config;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.when;
+
+/**
+ * Tests for {@link AlwaysOnVpnManagedProfilePreferenceController}.
+ */
+@RunWith(SettingsRobolectricTestRunner.class)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
+public final class AlwaysOnVpnManagedProfilePreferenceControllerTest {
+
+    @Mock(answer = Answers.RETURNS_DEEP_STUBS)
+    private Context mContext;
+    private FakeFeatureFactory mFeatureFactory;
+
+    private AlwaysOnVpnManagedProfilePreferenceController mController;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        FakeFeatureFactory.setupForTest(mContext);
+        mFeatureFactory = (FakeFeatureFactory) FakeFeatureFactory.getFactory(mContext);
+        mController = new AlwaysOnVpnManagedProfilePreferenceController(mContext);
+    }
+
+    @Test
+    public void testUpdateState() {
+        final Preference preference = new Preference(mContext, null, 0, 0);
+        preference.setVisible(true);
+
+        when(mFeatureFactory.enterprisePrivacyFeatureProvider.isAlwaysOnVpnSetInManagedProfile())
+                .thenReturn(false);
+        mController.updateState(preference);
+        assertThat(preference.isVisible()).isFalse();
+
+        when(mFeatureFactory.enterprisePrivacyFeatureProvider.isAlwaysOnVpnSetInManagedProfile())
+                .thenReturn(true);
+        mController.updateState(preference);
+        assertThat(preference.isVisible()).isTrue();
+    }
+
+    @Test
+    public void testIsAvailable() {
+        assertThat(mController.isAvailable()).isTrue();
+    }
+
+    @Test
+    public void testHandlePreferenceTreeClick() {
+        assertThat(mController.handlePreferenceTreeClick(new Preference(mContext, null, 0, 0)))
+                .isFalse();
+    }
+
+    @Test
+    public void testGetPreferenceKey() {
+        assertThat(mController.getPreferenceKey()).isEqualTo("always_on_vpn_managed_profile");
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/enterprise/AlwaysOnVpnPrimaryUserPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/enterprise/AlwaysOnVpnPrimaryUserPreferenceControllerTest.java
new file mode 100644
index 0000000..d504b84
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/enterprise/AlwaysOnVpnPrimaryUserPreferenceControllerTest.java
@@ -0,0 +1,113 @@
+/*
+ * Copyright (C) 2017 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.enterprise;
+
+import android.content.Context;
+import com.android.settings.R;
+import android.support.v7.preference.Preference;
+
+import com.android.settings.SettingsRobolectricTestRunner;
+import com.android.settings.TestConfig;
+import com.android.settings.testutils.FakeFeatureFactory;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Answers;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.annotation.Config;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.when;
+
+/**
+ * Tests for {@link AlwaysOnVpnPrimaryUserPreferenceController}.
+ */
+@RunWith(SettingsRobolectricTestRunner.class)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
+public final class AlwaysOnVpnPrimaryUserPreferenceControllerTest {
+
+    private final String VPN_SET_DEVICE = "VPN set";
+    private final String VPN_SET_PERSONAL = "VPN set in personal profile";
+
+    @Mock(answer = Answers.RETURNS_DEEP_STUBS)
+    private Context mContext;
+    private FakeFeatureFactory mFeatureFactory;
+
+    private AlwaysOnVpnPrimaryUserPreferenceController mController;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        FakeFeatureFactory.setupForTest(mContext);
+        mFeatureFactory = (FakeFeatureFactory) FakeFeatureFactory.getFactory(mContext);
+        mController = new AlwaysOnVpnPrimaryUserPreferenceController(mContext);
+        when(mContext.getString(R.string.enterprise_privacy_always_on_vpn_device))
+                .thenReturn(VPN_SET_DEVICE);
+        when(mContext.getString(R.string.enterprise_privacy_always_on_vpn_personal))
+                .thenReturn(VPN_SET_PERSONAL);
+    }
+
+    @Test
+    public void testUpdateState() {
+        final Preference preference = new Preference(mContext, null, 0, 0);
+        preference.setVisible(true);
+
+        when(mFeatureFactory.enterprisePrivacyFeatureProvider.isInCompMode()).thenReturn(false);
+
+        when(mFeatureFactory.enterprisePrivacyFeatureProvider.isAlwaysOnVpnSetInPrimaryUser())
+                .thenReturn(false);
+        mController.updateState(preference);
+        assertThat(preference.isVisible()).isFalse();
+
+        when(mFeatureFactory.enterprisePrivacyFeatureProvider.isAlwaysOnVpnSetInPrimaryUser())
+                .thenReturn(true);
+        mController.updateState(preference);
+        assertThat(preference.isVisible()).isTrue();
+        assertThat(preference.getTitle()).isEqualTo(VPN_SET_DEVICE);
+
+        when(mFeatureFactory.enterprisePrivacyFeatureProvider.isInCompMode()).thenReturn(true);
+
+        when(mFeatureFactory.enterprisePrivacyFeatureProvider.isAlwaysOnVpnSetInPrimaryUser())
+                .thenReturn(false);
+        mController.updateState(preference);
+        assertThat(preference.isVisible()).isFalse();
+
+        when(mFeatureFactory.enterprisePrivacyFeatureProvider.isAlwaysOnVpnSetInPrimaryUser())
+                .thenReturn(true);
+        mController.updateState(preference);
+        assertThat(preference.isVisible()).isTrue();
+        assertThat(preference.getTitle()).isEqualTo(VPN_SET_PERSONAL);
+    }
+
+    @Test
+    public void testIsAvailable() {
+        assertThat(mController.isAvailable()).isTrue();
+    }
+
+    @Test
+    public void testHandlePreferenceTreeClick() {
+        assertThat(mController.handlePreferenceTreeClick(new Preference(mContext, null, 0, 0)))
+                .isFalse();
+    }
+
+    @Test
+    public void testGetPreferenceKey() {
+        assertThat(mController.getPreferenceKey()).isEqualTo("always_on_vpn_primary_user");
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/enterprise/EnterprisePrivacyFeatureProviderImplTest.java b/tests/robotests/src/com/android/settings/enterprise/EnterprisePrivacyFeatureProviderImplTest.java
index 9688c12..3fc6f7c 100644
--- a/tests/robotests/src/com/android/settings/enterprise/EnterprisePrivacyFeatureProviderImplTest.java
+++ b/tests/robotests/src/com/android/settings/enterprise/EnterprisePrivacyFeatureProviderImplTest.java
@@ -17,12 +17,15 @@
 package com.android.settings.enterprise;
 
 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 com.android.settings.SettingsRobolectricTestRunner;
 import com.android.settings.TestConfig;
 import com.android.settings.applications.PackageManagerWrapper;
+import com.android.settings.vpn2.ConnectivityManagerWrapper;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -31,7 +34,9 @@
 import org.mockito.MockitoAnnotations;
 import org.robolectric.annotation.Config;
 
+import java.util.ArrayList;
 import java.util.Date;
+import java.util.List;
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Mockito.when;
@@ -45,9 +50,16 @@
 
     private final ComponentName DEVICE_OWNER = new ComponentName("dummy", "component");
     private final Date TIMESTAMP = new Date(2011, 11, 11);
+    private final int MY_USER_ID = UserHandle.myUserId();
+    private final int MANAGED_PROFILE_USER_ID = MY_USER_ID + 1;
+    private final String VPN_PACKAGE_ID = "com.example.vpn";
+
+    private List<UserInfo> mProfiles = new ArrayList();
 
     private @Mock DevicePolicyManagerWrapper mDevicePolicyManager;
     private @Mock PackageManagerWrapper mPackageManager;
+    private @Mock UserManager mUserManager;
+    private @Mock ConnectivityManagerWrapper mConnectivityManger;
 
     private EnterprisePrivacyFeatureProvider mProvider;
 
@@ -57,8 +69,11 @@
 
         when(mPackageManager.hasSystemFeature(PackageManager.FEATURE_DEVICE_ADMIN))
                 .thenReturn(true);
+        when(mUserManager.getProfiles(MY_USER_ID)).thenReturn(mProfiles);
+        mProfiles.add(new UserInfo(MY_USER_ID, "", "", 0 /* flags */));
 
-        mProvider = new EnterprisePrivacyFeatureProviderImpl(mDevicePolicyManager, mPackageManager);
+        mProvider = new EnterprisePrivacyFeatureProviderImpl(mDevicePolicyManager, mPackageManager,
+                mUserManager, mConnectivityManger);
     }
 
     @Test
@@ -71,6 +86,15 @@
     }
 
     @Test
+    public void testIsInCompMode() {
+        when(mDevicePolicyManager.getDeviceOwnerComponentOnAnyUser()).thenReturn(DEVICE_OWNER);
+        assertThat(mProvider.isInCompMode()).isFalse();
+
+        mProfiles.add(new UserInfo(MANAGED_PROFILE_USER_ID, "", "", UserInfo.FLAG_MANAGED_PROFILE));
+        assertThat(mProvider.isInCompMode()).isTrue();
+    }
+
+    @Test
     public void testGetLastSecurityLogRetrievalTime() {
         when(mDevicePolicyManager.getLastSecurityLogRetrievalTime()).thenReturn(-1L);
         assertThat(mProvider.getLastSecurityLogRetrievalTime()).isNull();
@@ -97,4 +121,29 @@
         when(mDevicePolicyManager.getLastNetworkLogRetrievalTime()).thenReturn(TIMESTAMP.getTime());
         assertThat(mProvider.getLastNetworkLogRetrievalTime()).isEqualTo(TIMESTAMP);
     }
+
+    @Test
+    public void testIsAlwaysOnVpnSetInPrimaryUser() {
+        when(mConnectivityManger.getAlwaysOnVpnPackageForUser(MY_USER_ID)).thenReturn(null);
+        assertThat(mProvider.isAlwaysOnVpnSetInPrimaryUser()).isFalse();
+
+        when(mConnectivityManger.getAlwaysOnVpnPackageForUser(MY_USER_ID))
+                .thenReturn(VPN_PACKAGE_ID);
+        assertThat(mProvider.isAlwaysOnVpnSetInPrimaryUser()).isTrue();
+    }
+
+    @Test
+    public void testIsAlwaysOnVpnSetInManagedProfileProfile() {
+        assertThat(mProvider.isAlwaysOnVpnSetInManagedProfile()).isFalse();
+
+        mProfiles.add(new UserInfo(MANAGED_PROFILE_USER_ID, "", "", UserInfo.FLAG_MANAGED_PROFILE));
+
+        when(mConnectivityManger.getAlwaysOnVpnPackageForUser(MANAGED_PROFILE_USER_ID))
+                .thenReturn(null);
+        assertThat(mProvider.isAlwaysOnVpnSetInManagedProfile()).isFalse();
+
+        when(mConnectivityManger.getAlwaysOnVpnPackageForUser(MANAGED_PROFILE_USER_ID))
+                .thenReturn(VPN_PACKAGE_ID);
+        assertThat(mProvider.isAlwaysOnVpnSetInManagedProfile()).isTrue();
+    }
 }
diff --git a/tests/robotests/src/com/android/settings/enterprise/EnterprisePrivacySettingsTest.java b/tests/robotests/src/com/android/settings/enterprise/EnterprisePrivacySettingsTest.java
index f6e18c6..4d86e61 100644
--- a/tests/robotests/src/com/android/settings/enterprise/EnterprisePrivacySettingsTest.java
+++ b/tests/robotests/src/com/android/settings/enterprise/EnterprisePrivacySettingsTest.java
@@ -73,10 +73,14 @@
         final List<PreferenceController> controllers = mSettings.getPreferenceControllers(
                 ShadowApplication.getInstance().getApplicationContext());
         assertThat(controllers).isNotNull();
-        assertThat(controllers.size()).isEqualTo(4);
+        assertThat(controllers.size()).isEqualTo(6);
         assertThat(controllers.get(0)).isInstanceOf(InstalledPackagesPreferenceController.class);
         assertThat(controllers.get(1)).isInstanceOf(NetworkLogsPreferenceController.class);
         assertThat(controllers.get(2)).isInstanceOf(BugReportsPreferenceController.class);
         assertThat(controllers.get(3)).isInstanceOf(SecurityLogsPreferenceController.class);
+        assertThat(controllers.get(4)).isInstanceOf(
+                AlwaysOnVpnPrimaryUserPreferenceController.class);
+        assertThat(controllers.get(5)).isInstanceOf(
+                AlwaysOnVpnManagedProfilePreferenceController.class);
     }
 }
diff --git a/tests/robotests/src/com/android/settings/vpn2/VpnUtilsTest.java b/tests/robotests/src/com/android/settings/vpn2/VpnUtilsTest.java
new file mode 100644
index 0000000..582f9fb
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/vpn2/VpnUtilsTest.java
@@ -0,0 +1,45 @@
+/*
+ * Copyright (C) 2017 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.vpn2;
+
+import com.android.settings.SettingsRobolectricTestRunner;
+import com.android.settings.TestConfig;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.robolectric.annotation.Config;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/**
+ * Tests for {@link VpnUtils}.
+ */
+@RunWith(SettingsRobolectricTestRunner.class)
+@Config(manifest = TestConfig.MANIFEST_PATH, sdk = TestConfig.SDK_VERSION)
+public final class VpnUtilsTest {
+    @Test
+    public void testIsAlwaysOnVpnSet() {
+        final ConnectivityManagerWrapper cm = mock(ConnectivityManagerWrapper.class);
+        when(cm.getAlwaysOnVpnPackageForUser(0)).thenReturn("com.example.vpn");
+        assertThat(VpnUtils.isAlwaysOnVpnSet(cm, 0)).isTrue();
+
+        when(cm.getAlwaysOnVpnPackageForUser(0)).thenReturn(null);
+        assertThat(VpnUtils.isAlwaysOnVpnSet(cm, 0)).isFalse();
+    }
+}