Merge commit '087d29e0a2380ccb90e847e5daa52f50235731f0' into nyc-mr1-dev-plus-aosp

* commit '087d29e0a2380ccb90e847e5daa52f50235731f0':
  Captive portals: login activity probes like NetworkMonitor
  Logging improvements in CaptivePortalLoginActivity
  Captive portal systel log improvements
  Fix flaky IpPrefixTest.
  Removing bogus file added by mistake
  frameworks-test: adding missing @SmallTest
  Netd events: record connect() success/errno
  Add missing dependency.
  Show notification for always-on app VPN

core/java/android/net/ConnectivityManager.java

@@ -229,6 +229,13 @@ public class ConnectivityManager {
      */
     public static final String EXTRA_CAPTIVE_PORTAL_URL = "android.net.extra.CAPTIVE_PORTAL_URL";
 
+    /**
+     * Key for passing a user agent string to the captive portal login activity.
+     * {@hide}
+     */
+    public static final String EXTRA_CAPTIVE_PORTAL_USER_AGENT =
+            "android.net.extra.CAPTIVE_PORTAL_USER_AGENT";
+
     /**
      * Broadcast action to indicate the change of data activity status
      * (idle or active) on a network in a recent period.

core/tests/coretests/src/android/net/IpPrefixTest.java

@@ -18,14 +18,14 @@ package android.net;
 
 import android.net.IpPrefix;
 import android.os.Parcel;
-import static android.test.MoreAsserts.assertNotEqual;
 import android.test.suitebuilder.annotation.SmallTest;
-
-import static org.junit.Assert.assertArrayEquals;
 import java.net.InetAddress;
 import java.util.Random;
 import junit.framework.TestCase;
 
+import static android.test.MoreAsserts.assertNotEqual;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
 
 public class IpPrefixTest extends TestCase {
 
@@ -242,25 +242,42 @@ public class IpPrefixTest extends TestCase {
 
     @SmallTest
     public void testHashCode() {
-        IpPrefix p;
+        IpPrefix p = new IpPrefix(new byte[4], 0);
-        int oldCode = -1;
         Random random = new Random();
         for (int i = 0; i < 100; i++) {
+            final IpPrefix oldP = p;
             if (random.nextBoolean()) {
                 // IPv4.
                 byte[] b = new byte[4];
                 random.nextBytes(b);
                 p = new IpPrefix(b, random.nextInt(33));
-                assertNotEqual(oldCode, p.hashCode());
-                oldCode = p.hashCode();
             } else {
                 // IPv6.
                 byte[] b = new byte[16];
                 random.nextBytes(b);
                 p = new IpPrefix(b, random.nextInt(129));
-                assertNotEqual(oldCode, p.hashCode());
-                oldCode = p.hashCode();
             }
+            if (p.equals(oldP)) {
+              assertEquals(p.hashCode(), oldP.hashCode());
+            }
+            if (p.hashCode() != oldP.hashCode()) {
+              assertNotEqual(p, oldP);
+            }
+        }
+    }
+
+    @SmallTest
+    public void testHashCodeIsNotConstant() {
+        IpPrefix[] prefixes = {
+            new IpPrefix("2001:db8:f00::ace:d00d/127"),
+            new IpPrefix("192.0.2.0/23"),
+            new IpPrefix("::/0"),
+            new IpPrefix("0.0.0.0/0"),
+        };
+        for (int i = 0; i < prefixes.length; i++) {
+          for (int j = i + 1; j < prefixes.length; j++) {
+            assertNotEqual(prefixes[i].hashCode(), prefixes[j].hashCode());
+          }
         }
     }
 

services/core/java/com/android/server/connectivity/NetworkNotificationManager.java

@@ -114,7 +114,7 @@ public class NetworkNotificationManager {
         }
 
         if (DBG) {
-            Slog.d(TAG, "showNotification " + notifyType
+            Slog.d(TAG, "showNotification id=" + id + " " + notifyType
                     + " transportType=" + getTransportName(transportType)
                     + " extraInfo=" + extraInfo + " highPriority=" + highPriority);
         }
@@ -187,7 +187,7 @@ public class NetworkNotificationManager {
         try {
             mNotificationManager.notifyAsUser(NOTIFICATION_ID, id, notification, UserHandle.ALL);
         } catch (NullPointerException npe) {
-            Slog.d(TAG, "setNotificationVisible: visible notificationManager npe=" + npe);
+            Slog.d(TAG, "setNotificationVisible: visible notificationManager error", npe);
         }
     }
 
@@ -198,7 +198,7 @@ public class NetworkNotificationManager {
         try {
             mNotificationManager.cancelAsUser(NOTIFICATION_ID, id, UserHandle.ALL);
         } catch (NullPointerException npe) {
-            Slog.d(TAG, "setNotificationVisible: cancel notificationManager npe=" + npe);
+            Slog.d(TAG, "setNotificationVisible: cancel notificationManager error", npe);
         }
     }
 

tests/net/java/com/android/server/ConnectivityServiceTest.java

@@ -70,7 +70,6 @@ import android.provider.Settings;
 import android.test.AndroidTestCase;
 import android.test.FlakyTest;
 import android.test.mock.MockContentResolver;
-import android.test.suitebuilder.annotation.LargeTest;
 import android.test.suitebuilder.annotation.SmallTest;
 import android.util.Log;
 import android.util.LogPrinter;
@@ -196,6 +195,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
     }
 
     // Tests that IdleableHandlerThread works as expected.
+    @SmallTest
     public void testIdleableHandlerThread() {
         final int attempts = 50;  // Causes the test to take about 200ms on bullhead-eng.
 
@@ -220,6 +220,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         }
     }
 
+    @SmallTest
     @FlakyTest(tolerance = 3)
     public void testNotWaitingForIdleCausesRaceConditions() {
         // Bring up a network that we can use to send messages to ConnectivityService.
@@ -840,7 +841,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         return cv;
     }
 
-    @LargeTest
+    @SmallTest
     public void testLingering() throws Exception {
         verifyNoNetwork();
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
@@ -880,7 +881,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         verifyNoNetwork();
     }
 
-    @LargeTest
+    @SmallTest
     public void testValidatedCellularOutscoresUnvalidatedWiFi() throws Exception {
         // Test bringing up unvalidated WiFi
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
@@ -915,7 +916,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         verifyNoNetwork();
     }
 
-    @LargeTest
+    @SmallTest
     public void testUnvalidatedWifiOutscoresUnvalidatedCellular() throws Exception {
         // Test bringing up unvalidated cellular.
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
@@ -941,7 +942,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         verifyNoNetwork();
     }
 
-    @LargeTest
+    @SmallTest
     public void testUnlingeringDoesNotValidate() throws Exception {
         // Test bringing up unvalidated WiFi.
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
@@ -969,7 +970,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
                 NET_CAPABILITY_VALIDATED));
     }
 
-    @LargeTest
+    @SmallTest
     public void testCellularOutscoresWeakWifi() throws Exception {
         // Test bringing up validated cellular.
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
@@ -995,7 +996,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         verifyActiveNetwork(TRANSPORT_WIFI);
     }
 
-    @LargeTest
+    @SmallTest
     public void testReapingNetwork() throws Exception {
         // Test bringing up WiFi without NET_CAPABILITY_INTERNET.
         // Expect it to be torn down immediately because it satisfies no requests.
@@ -1028,7 +1029,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         waitFor(cv);
     }
 
-    @LargeTest
+    @SmallTest
     public void testCellularFallback() throws Exception {
         // Test bringing up validated cellular.
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
@@ -1066,7 +1067,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         verifyActiveNetwork(TRANSPORT_WIFI);
     }
 
-    @LargeTest
+    @SmallTest
     public void testWiFiFallback() throws Exception {
         // Test bringing up unvalidated WiFi.
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
@@ -1208,7 +1209,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         }
     }
 
-    @LargeTest
+    @SmallTest
     public void testStateChangeNetworkCallbacks() throws Exception {
         final TestNetworkCallback genericNetworkCallback = new TestNetworkCallback();
         final TestNetworkCallback wifiNetworkCallback = new TestNetworkCallback();
@@ -1591,7 +1592,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         handlerThread.quit();
     }
 
-    @LargeTest
+    @SmallTest
     public void testNetworkFactoryRequests() throws Exception {
         tryNetworkFactoryRequests(NET_CAPABILITY_MMS);
         tryNetworkFactoryRequests(NET_CAPABILITY_SUPL);
@@ -1611,7 +1612,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         // Skipping VALIDATED and CAPTIVE_PORTAL as they're disallowed.
     }
 
-    @LargeTest
+    @SmallTest
     public void testNoMutableNetworkRequests() throws Exception {
         PendingIntent pendingIntent = PendingIntent.getBroadcast(mContext, 0, new Intent("a"), 0);
         NetworkRequest.Builder builder = new NetworkRequest.Builder();
@@ -1636,7 +1637,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         } catch (IllegalArgumentException expected) {}
     }
 
-    @LargeTest
+    @SmallTest
     public void testMMSonWiFi() throws Exception {
         // Test bringing up cellular without MMS NetworkRequest gets reaped
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
@@ -1671,7 +1672,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         verifyActiveNetwork(TRANSPORT_WIFI);
     }
 
-    @LargeTest
+    @SmallTest
     public void testMMSonCell() throws Exception {
         // Test bringing up cellular without MMS
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
@@ -1697,7 +1698,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         verifyActiveNetwork(TRANSPORT_CELLULAR);
     }
 
-    @LargeTest
+    @SmallTest
     public void testCaptivePortal() {
         final TestNetworkCallback captivePortalCallback = new TestNetworkCallback();
         final NetworkRequest captivePortalRequest = new NetworkRequest.Builder()
@@ -1746,7 +1747,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         validatedCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
     }
 
-    @LargeTest
+    @SmallTest
     public void testAvoidOrIgnoreCaptivePortals() {
         final TestNetworkCallback captivePortalCallback = new TestNetworkCallback();
         final NetworkRequest captivePortalRequest = new NetworkRequest.Builder()
@@ -1817,7 +1818,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
                 execptionCalled);
     }
 
-    @LargeTest
+    @SmallTest
     public void testRegisterDefaultNetworkCallback() throws Exception {
         final TestNetworkCallback defaultNetworkCallback = new TestNetworkCallback();
         mCm.registerDefaultNetworkCallback(defaultNetworkCallback);
@@ -1878,7 +1879,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         }
     }
 
-    @LargeTest
+    @SmallTest
     public void testRequestCallbackUpdates() throws Exception {
         // File a network request for mobile.
         final TestNetworkCallback cellNetworkCallback = new TestRequestUpdateCallback();
@@ -2519,6 +2520,7 @@ public class ConnectivityServiceTest extends AndroidTestCase {
         return mWiFiNetworkAgent.getNetwork();
     }
 
+    @SmallTest
     public void testPacketKeepalives() throws Exception {
         InetAddress myIPv4 = InetAddress.getByName("192.0.2.129");
         InetAddress notMyIPv4 = InetAddress.getByName("192.0.2.35");

tests/net/java/com/android/server/connectivity/LingerMonitorTest.java

@@ -24,6 +24,7 @@ import android.net.Network;
 import android.net.NetworkCapabilities;
 import android.net.NetworkInfo;
 import android.net.NetworkMisc;
+import android.test.suitebuilder.annotation.SmallTest;
 import android.text.format.DateUtils;
 import com.android.internal.R;
 import com.android.server.ConnectivityService;
@@ -70,6 +71,7 @@ public class LingerMonitorTest extends TestCase {
         mMonitor = new TestableLingerMonitor(mCtx, mNotifier, HIGH_DAILY_LIMIT, HIGH_RATE_LIMIT);
     }
 
+    @SmallTest
     public void testTransitions() {
         setNotificationSwitch(transition(WIFI, CELLULAR));
         NetworkAgentInfo nai1 = wifiNai(100);
@@ -79,6 +81,7 @@ public class LingerMonitorTest extends TestCase {
         assertFalse(mMonitor.isNotificationEnabled(nai2, nai1));
     }
 
+    @SmallTest
     public void testNotificationOnLinger() {
         setNotificationSwitch(transition(WIFI, CELLULAR));
         setNotificationType(LingerMonitor.NOTIFY_TYPE_NOTIFICATION);
@@ -89,6 +92,7 @@ public class LingerMonitorTest extends TestCase {
         verifyNotification(from, to);
     }
 
+    @SmallTest
     public void testToastOnLinger() {
         setNotificationSwitch(transition(WIFI, CELLULAR));
         setNotificationType(LingerMonitor.NOTIFY_TYPE_TOAST);
@@ -99,6 +103,7 @@ public class LingerMonitorTest extends TestCase {
         verifyToast(from, to);
     }
 
+    @SmallTest
     public void testNotificationClearedAfterDisconnect() {
         setNotificationSwitch(transition(WIFI, CELLULAR));
         setNotificationType(LingerMonitor.NOTIFY_TYPE_NOTIFICATION);
@@ -112,6 +117,7 @@ public class LingerMonitorTest extends TestCase {
         verify(mNotifier, times(1)).clearNotification(100);
     }
 
+    @SmallTest
     public void testNotificationClearedAfterSwitchingBack() {
         setNotificationSwitch(transition(WIFI, CELLULAR));
         setNotificationType(LingerMonitor.NOTIFY_TYPE_NOTIFICATION);
@@ -125,6 +131,7 @@ public class LingerMonitorTest extends TestCase {
         verify(mNotifier, times(1)).clearNotification(100);
     }
 
+    @SmallTest
     public void testUniqueToast() {
         setNotificationSwitch(transition(WIFI, CELLULAR));
         setNotificationType(LingerMonitor.NOTIFY_TYPE_TOAST);
@@ -142,6 +149,7 @@ public class LingerMonitorTest extends TestCase {
         verifyNoNotifications();
     }
 
+    @SmallTest
     public void testMultipleNotifications() {
         setNotificationSwitch(transition(WIFI, CELLULAR));
         setNotificationType(LingerMonitor.NOTIFY_TYPE_NOTIFICATION);
@@ -160,6 +168,7 @@ public class LingerMonitorTest extends TestCase {
         verifyNotification(wifi2, cell);
     }
 
+    @SmallTest
     public void testRateLimiting() throws InterruptedException {
         mMonitor = new TestableLingerMonitor(mCtx, mNotifier, HIGH_DAILY_LIMIT, LOW_RATE_LIMIT);
 
@@ -185,6 +194,7 @@ public class LingerMonitorTest extends TestCase {
         verifyNoNotifications();
     }
 
+    @SmallTest
     public void testDailyLimiting() throws InterruptedException {
         mMonitor = new TestableLingerMonitor(mCtx, mNotifier, LOW_DAILY_LIMIT, HIGH_RATE_LIMIT);
 
@@ -211,6 +221,7 @@ public class LingerMonitorTest extends TestCase {
         verifyNoNotifications();
     }
 
+    @SmallTest
     public void testUniqueNotification() {
         setNotificationSwitch(transition(WIFI, CELLULAR));
         setNotificationType(LingerMonitor.NOTIFY_TYPE_NOTIFICATION);
@@ -227,6 +238,7 @@ public class LingerMonitorTest extends TestCase {
         verifyNotification(from, to);
     }
 
+    @SmallTest
     public void testIgnoreNeverValidatedNetworks() {
         setNotificationType(LingerMonitor.NOTIFY_TYPE_TOAST);
         setNotificationSwitch(transition(WIFI, CELLULAR));
@@ -238,6 +250,7 @@ public class LingerMonitorTest extends TestCase {
         verifyNoNotifications();
     }
 
+    @SmallTest
     public void testIgnoreCurrentlyValidatedNetworks() {
         setNotificationType(LingerMonitor.NOTIFY_TYPE_TOAST);
         setNotificationSwitch(transition(WIFI, CELLULAR));
@@ -249,6 +262,7 @@ public class LingerMonitorTest extends TestCase {
         verifyNoNotifications();
     }
 
+    @SmallTest
     public void testNoNotificationType() {
         setNotificationType(LingerMonitor.NOTIFY_TYPE_TOAST);
         setNotificationSwitch();
@@ -259,6 +273,7 @@ public class LingerMonitorTest extends TestCase {
         verifyNoNotifications();
     }
 
+    @SmallTest
     public void testNoTransitionToNotify() {
         setNotificationType(LingerMonitor.NOTIFY_TYPE_NONE);
         setNotificationSwitch(transition(WIFI, CELLULAR));
@@ -269,6 +284,7 @@ public class LingerMonitorTest extends TestCase {
         verifyNoNotifications();
     }
 
+    @SmallTest
     public void testDifferentTransitionToNotify() {
         setNotificationType(LingerMonitor.NOTIFY_TYPE_TOAST);
         setNotificationSwitch(transition(CELLULAR, WIFI));

tests/net/java/com/android/server/connectivity/NetdEventListenerServiceTest.java

@@ -22,6 +22,8 @@ import android.net.Network;
 import android.net.metrics.DnsEvent;
 import android.net.metrics.INetdEventListener;
 import android.net.metrics.IpConnectivityLog;
+import android.os.RemoteException;
+import android.test.suitebuilder.annotation.SmallTest;
 
 import junit.framework.TestCase;
 import org.junit.Before;
@@ -82,6 +84,7 @@ public class NetdEventListenerServiceTest extends TestCase {
         verify(mCm, times(1)).registerNetworkCallback(any(), mCallbackCaptor.capture());
     }
 
+    @SmallTest
     public void testOneBatch() throws Exception {
         log(105, LATENCIES);
         log(106, Arrays.copyOf(LATENCIES, BATCH_SIZE - 1)); // one lookup short of a batch event
@@ -96,6 +99,7 @@ public class NetdEventListenerServiceTest extends TestCase {
             new DnsEvent(106, EVENT_TYPES, RETURN_CODES, LATENCIES));
     }
 
+    @SmallTest
     public void testSeveralBatches() throws Exception {
         log(105, LATENCIES);
         log(106, LATENCIES);
@@ -109,6 +113,7 @@ public class NetdEventListenerServiceTest extends TestCase {
             new DnsEvent(107, EVENT_TYPES, RETURN_CODES, LATENCIES));
     }
 
+    @SmallTest
     public void testBatchAndNetworkLost() throws Exception {
         byte[] eventTypes = Arrays.copyOf(EVENT_TYPES, 20);
         byte[] returnCodes = Arrays.copyOf(RETURN_CODES, 20);
@@ -125,6 +130,7 @@ public class NetdEventListenerServiceTest extends TestCase {
             new DnsEvent(105, EVENT_TYPES, RETURN_CODES, LATENCIES));
     }
 
+    @SmallTest
     public void testConcurrentBatchesAndDumps() throws Exception {
         final long stop = System.currentTimeMillis() + 100;
         final PrintWriter pw = new PrintWriter(new FileOutputStream("/dev/null"));
@@ -146,6 +152,7 @@ public class NetdEventListenerServiceTest extends TestCase {
             new DnsEvent(107, EVENT_TYPES, RETURN_CODES, LATENCIES));
     }
 
+    @SmallTest
     public void testConcurrentBatchesAndNetworkLoss() throws Exception {
         logAsync(105, LATENCIES);
         Thread.sleep(10L);
@@ -157,9 +164,13 @@ public class NetdEventListenerServiceTest extends TestCase {
     }
 
     void log(int netId, int[] latencies) {
-        for (int l : latencies) {
+        try {
-            mNetdEventListenerService.onDnsEvent(netId, EVENT_TYPE, RETURN_CODE, l, null, null, 0,
+            for (int l : latencies) {
-                    0);
+                mNetdEventListenerService.onDnsEvent(netId, EVENT_TYPE, RETURN_CODE, l, null, null,
+                        0, 0);
+            }
+        } catch (RemoteException re) {
+            throw re.rethrowFromSystemServer();
         }
     }
 

tests/net/java/com/android/server/connectivity/VpnTest.java

@@ -25,9 +25,11 @@ import static org.mockito.Mockito.*;
 
 import android.annotation.UserIdInt;
 import android.app.AppOpsManager;
+import android.app.NotificationManager;
 import android.content.Context;
 import android.content.pm.PackageManager;
 import android.content.pm.UserInfo;
+import android.net.NetworkInfo.DetailedState;
 import android.net.UidRange;
 import android.os.INetworkManagementService;
 import android.os.Looper;
@@ -43,6 +45,8 @@ import java.util.Arrays;
 import java.util.Map;
 import java.util.Set;
 
+import org.mockito.ArgumentCaptor;
+import org.mockito.InOrder;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -88,14 +92,18 @@ public class VpnTest extends AndroidTestCase {
     @Mock private PackageManager mPackageManager;
     @Mock private INetworkManagementService mNetService;
     @Mock private AppOpsManager mAppOps;
+    @Mock private NotificationManager mNotificationManager;
 
     @Override
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
         when(mContext.getPackageManager()).thenReturn(mPackageManager);
         setMockedPackages(mPackages);
+        when(mContext.getPackageName()).thenReturn(Vpn.class.getPackage().getName());
         when(mContext.getSystemService(eq(Context.USER_SERVICE))).thenReturn(mUserManager);
         when(mContext.getSystemService(eq(Context.APP_OPS_SERVICE))).thenReturn(mAppOps);
+        when(mContext.getSystemService(eq(Context.NOTIFICATION_SERVICE)))
+                .thenReturn(mNotificationManager);
         doNothing().when(mNetService).registerObserver(any());
     }
 
@@ -103,7 +111,7 @@ public class VpnTest extends AndroidTestCase {
     public void testRestrictedProfilesAreAddedToVpn() {
         setMockedUsers(primaryUser, secondaryUser, restrictedProfileA, restrictedProfileB);
 
-        final Vpn vpn = new MockVpn(primaryUser.id);
+        final Vpn vpn = spyVpn(primaryUser.id);
         final Set<UidRange> ranges = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
                 null, null);
 
@@ -117,7 +125,7 @@ public class VpnTest extends AndroidTestCase {
     public void testManagedProfilesAreNotAddedToVpn() {
         setMockedUsers(primaryUser, managedProfileA);
 
-        final Vpn vpn = new MockVpn(primaryUser.id);
+        final Vpn vpn = spyVpn(primaryUser.id);
         final Set<UidRange> ranges = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
                 null, null);
 
@@ -130,7 +138,7 @@ public class VpnTest extends AndroidTestCase {
     public void testAddUserToVpnOnlyAddsOneUser() {
         setMockedUsers(primaryUser, restrictedProfileA, managedProfileA);
 
-        final Vpn vpn = new MockVpn(primaryUser.id);
+        final Vpn vpn = spyVpn(primaryUser.id);
         final Set<UidRange> ranges = new ArraySet<>();
         vpn.addUserToRanges(ranges, primaryUser.id, null, null);
 
@@ -141,7 +149,7 @@ public class VpnTest extends AndroidTestCase {
 
     @SmallTest
     public void testUidWhiteAndBlacklist() throws Exception {
-        final Vpn vpn = new MockVpn(primaryUser.id);
+        final Vpn vpn = spyVpn(primaryUser.id);
         final UidRange user = UidRange.createForUser(primaryUser.id);
         final String[] packages = {PKGS[0], PKGS[1], PKGS[2]};
 
@@ -166,15 +174,15 @@ public class VpnTest extends AndroidTestCase {
 
     @SmallTest
     public void testLockdownChangingPackage() throws Exception {
-        final MockVpn vpn = new MockVpn(primaryUser.id);
+        final Vpn vpn = spyVpn(primaryUser.id);
         final UidRange user = UidRange.createForUser(primaryUser.id);
 
         // Default state.
-        vpn.assertUnblocked(user.start + PKG_UIDS[0], user.start + PKG_UIDS[1], user.start + PKG_UIDS[2], user.start + PKG_UIDS[3]);
+        assertUnblocked(vpn, user.start + PKG_UIDS[0], user.start + PKG_UIDS[1], user.start + PKG_UIDS[2], user.start + PKG_UIDS[3]);
 
         // Set always-on without lockdown.
         assertTrue(vpn.setAlwaysOnPackage(PKGS[1], false));
-        vpn.assertUnblocked(user.start + PKG_UIDS[0], user.start + PKG_UIDS[1], user.start + PKG_UIDS[2], user.start + PKG_UIDS[3]);
+        assertUnblocked(vpn, user.start + PKG_UIDS[0], user.start + PKG_UIDS[1], user.start + PKG_UIDS[2], user.start + PKG_UIDS[3]);
 
         // Set always-on with lockdown.
         assertTrue(vpn.setAlwaysOnPackage(PKGS[1], true));
@@ -182,8 +190,8 @@ public class VpnTest extends AndroidTestCase {
             new UidRange(user.start, user.start + PKG_UIDS[1] - 1),
             new UidRange(user.start + PKG_UIDS[1] + 1, user.stop)
         }));
-        vpn.assertBlocked(user.start + PKG_UIDS[0], user.start + PKG_UIDS[2], user.start + PKG_UIDS[3]);
+        assertBlocked(vpn, user.start + PKG_UIDS[0], user.start + PKG_UIDS[2], user.start + PKG_UIDS[3]);
-        vpn.assertUnblocked(user.start + PKG_UIDS[1]);
+        assertUnblocked(vpn, user.start + PKG_UIDS[1]);
 
         // Switch to another app.
         assertTrue(vpn.setAlwaysOnPackage(PKGS[3], true));
@@ -195,13 +203,13 @@ public class VpnTest extends AndroidTestCase {
             new UidRange(user.start, user.start + PKG_UIDS[3] - 1),
             new UidRange(user.start + PKG_UIDS[3] + 1, user.stop)
         }));
-        vpn.assertBlocked(user.start + PKG_UIDS[0], user.start + PKG_UIDS[1], user.start + PKG_UIDS[2]);
+        assertBlocked(vpn, user.start + PKG_UIDS[0], user.start + PKG_UIDS[1], user.start + PKG_UIDS[2]);
-        vpn.assertUnblocked(user.start + PKG_UIDS[3]);
+        assertUnblocked(vpn, user.start + PKG_UIDS[3]);
     }
 
     @SmallTest
     public void testLockdownAddingAProfile() throws Exception {
-        final MockVpn vpn = new MockVpn(primaryUser.id);
+        final Vpn vpn = spyVpn(primaryUser.id);
         setMockedUsers(primaryUser);
 
         // Make a copy of the restricted profile, as we're going to mark it deleted halfway through.
@@ -220,7 +228,7 @@ public class VpnTest extends AndroidTestCase {
         }));
 
         // Verify restricted user isn't affected at first.
-        vpn.assertUnblocked(profile.start + PKG_UIDS[0]);
+        assertUnblocked(vpn, profile.start + PKG_UIDS[0]);
 
         // Add the restricted user.
         setMockedUsers(primaryUser, tempProfile);
@@ -239,24 +247,53 @@ public class VpnTest extends AndroidTestCase {
         }));
     }
 
+    @SmallTest
+    public void testNotificationShownForAlwaysOnApp() {
+        final Vpn vpn = spyVpn(primaryUser.id);
+        final InOrder order = inOrder(vpn);
+        setMockedUsers(primaryUser);
+
+        // Don't show a notification for regular disconnected states.
+        vpn.updateState(DetailedState.DISCONNECTED, TAG);
+        order.verify(vpn).updateAlwaysOnNotificationInternal(false);
+
+        // Start showing a notification for disconnected once always-on.
+        vpn.setAlwaysOnPackage(PKGS[0], false);
+        order.verify(vpn).updateAlwaysOnNotificationInternal(true);
+
+        // Stop showing the notification once connected.
+        vpn.updateState(DetailedState.CONNECTED, TAG);
+        order.verify(vpn).updateAlwaysOnNotificationInternal(false);
+
+        // Show the notification if we disconnect again.
+        vpn.updateState(DetailedState.DISCONNECTED, TAG);
+        order.verify(vpn).updateAlwaysOnNotificationInternal(true);
+
+        // Notification should be cleared after unsetting always-on package.
+        vpn.setAlwaysOnPackage(null, false);
+        order.verify(vpn).updateAlwaysOnNotificationInternal(false);
+    }
+
     /**
-     * A subclass of {@link Vpn} with some of the fields pre-mocked.
+     * Mock some methods of vpn object.
      */
-    private class MockVpn extends Vpn {
+    private Vpn spyVpn(@UserIdInt int userId) {
-        public MockVpn(@UserIdInt int userId) {
+        final Vpn vpn = spy(new Vpn(Looper.myLooper(), mContext, mNetService, userId));
-            super(Looper.myLooper(), mContext, mNetService, userId);
-        }
 
-        public void assertBlocked(int... uids) {
+        // Block calls to the NotificationManager or PendingIntent#getActivity.
-            for (int uid : uids) {
+        doNothing().when(vpn).updateAlwaysOnNotificationInternal(anyBoolean());
-                assertTrue("Uid " + uid + " should be blocked", isBlockingUid(uid));
+        return vpn;
-            }
+    }
-        }
 
-        public void assertUnblocked(int... uids) {
+    private static void assertBlocked(Vpn vpn, int... uids) {
-            for (int uid : uids) {
+        for (int uid : uids) {
-                assertFalse("Uid " + uid + " should not be blocked", isBlockingUid(uid));
+            assertTrue("Uid " + uid + " should be blocked", vpn.isBlockingUid(uid));
-            }
+        }
+    }
+
+    private static void assertUnblocked(Vpn vpn, int... uids) {
+        for (int uid : uids) {
+            assertFalse("Uid " + uid + " should not be blocked", vpn.isBlockingUid(uid));
         }
     }