diff --git a/core/java/android/net/ConnectivityManager.java b/core/java/android/net/ConnectivityManager.java
index 40f47b0..4ccfea2 100644
--- a/core/java/android/net/ConnectivityManager.java
+++ b/core/java/android/net/ConnectivityManager.java
@@ -238,6 +238,14 @@
     public static final String EXTRA_CAPTIVE_PORTAL_URL = "android.net.extra.CAPTIVE_PORTAL_URL";
 
     /**
+     * Key for passing a {@link android.net.captiveportal.CaptivePortalProbeSpec} to the captive
+     * portal login activity.
+     * {@hide}
+     */
+    public static final String EXTRA_CAPTIVE_PORTAL_PROBE_SPEC =
+            "android.net.extra.CAPTIVE_PORTAL_PROBE_SPEC";
+
+    /**
      * Key for passing a user agent string to the captive portal login activity.
      * {@hide}
      */
diff --git a/core/java/android/net/Network.java b/core/java/android/net/Network.java
index f14847f..8f95607 100644
--- a/core/java/android/net/Network.java
+++ b/core/java/android/net/Network.java
@@ -79,6 +79,21 @@
     private static final long httpKeepAliveDurationMs =
             Long.parseLong(System.getProperty("http.keepAliveDuration", "300000"));  // 5 minutes.
 
+    // A boolean to control how getAllByName()/getByName() behaves in the face
+    // of Private DNS.
+    //
+    // When true, these calls will request that DNS resolution bypass any
+    // Private DNS that might otherwise apply. Use of this feature is restricted
+    // and permission checks are made by netd (attempts to bypass Private DNS
+    // without appropriate permission are silently turned into vanilla DNS
+    // requests). This only affects DNS queries made using this network object.
+    //
+    // It it not parceled to receivers because (a) it can be set or cleared at
+    // anytime and (b) receivers should be explicit about attempts to bypass
+    // Private DNS so that the intent of the code is easily determined and
+    // code search audits are possible.
+    private boolean mPrivateDnsBypass = false;
+
     /**
      * @hide
      */
@@ -102,7 +117,7 @@
      * @throws UnknownHostException if the address lookup fails.
      */
     public InetAddress[] getAllByName(String host) throws UnknownHostException {
-        return InetAddress.getAllByNameOnNet(host, netId);
+        return InetAddress.getAllByNameOnNet(host, getNetIdForResolv());
     }
 
     /**
@@ -116,7 +131,32 @@
      *             if the address lookup fails.
      */
     public InetAddress getByName(String host) throws UnknownHostException {
-        return InetAddress.getByNameOnNet(host, netId);
+        return InetAddress.getByNameOnNet(host, getNetIdForResolv());
+    }
+
+    /**
+     * Specify whether or not Private DNS should be bypassed when attempting
+     * to use {@link getAllByName()}/{@link getByName()} methods on the given
+     * instance for hostname resolution.
+     *
+     * @hide
+     */
+    public void setPrivateDnsBypass(boolean bypass) {
+        mPrivateDnsBypass = bypass;
+    }
+
+    /**
+     * Returns a netid marked with the Private DNS bypass flag.
+     *
+     * This flag must be kept in sync with the NETID_USE_LOCAL_NAMESERVERS flag
+     * in system/netd/include/NetdClient.h.
+     *
+     * @hide
+     */
+    public int getNetIdForResolv() {
+        return mPrivateDnsBypass
+                ? (int) (0x80000000L | (long) netId)  // Non-portable DNS resolution flag.
+                : netId;
     }
 
     /**
diff --git a/core/java/android/net/NetworkCapabilities.java b/core/java/android/net/NetworkCapabilities.java
index 19f0c90..83553df 100644
--- a/core/java/android/net/NetworkCapabilities.java
+++ b/core/java/android/net/NetworkCapabilities.java
@@ -17,6 +17,8 @@
 package android.net;
 
 import android.annotation.IntDef;
+import android.annotation.SystemApi;
+import android.annotation.TestApi;
 import android.net.ConnectivityManager.NetworkCallback;
 import android.os.Parcel;
 import android.os.Parcelable;
@@ -60,15 +62,7 @@
 
     public NetworkCapabilities(NetworkCapabilities nc) {
         if (nc != null) {
-            mNetworkCapabilities = nc.mNetworkCapabilities;
-            mTransportTypes = nc.mTransportTypes;
-            mLinkUpBandwidthKbps = nc.mLinkUpBandwidthKbps;
-            mLinkDownBandwidthKbps = nc.mLinkDownBandwidthKbps;
-            mNetworkSpecifier = nc.mNetworkSpecifier;
-            mSignalStrength = nc.mSignalStrength;
-            mUids = nc.mUids;
-            mEstablishingVpnAppUid = nc.mEstablishingVpnAppUid;
-            mUnwantedNetworkCapabilities = nc.mUnwantedNetworkCapabilities;
+            set(nc);
         }
     }
 
@@ -84,6 +78,24 @@
         mSignalStrength = SIGNAL_STRENGTH_UNSPECIFIED;
         mUids = null;
         mEstablishingVpnAppUid = INVALID_UID;
+        mSSID = null;
+    }
+
+    /**
+     * Set all contents of this object to the contents of a NetworkCapabilities.
+     * @hide
+     */
+    public void set(NetworkCapabilities nc) {
+        mNetworkCapabilities = nc.mNetworkCapabilities;
+        mTransportTypes = nc.mTransportTypes;
+        mLinkUpBandwidthKbps = nc.mLinkUpBandwidthKbps;
+        mLinkDownBandwidthKbps = nc.mLinkDownBandwidthKbps;
+        mNetworkSpecifier = nc.mNetworkSpecifier;
+        mSignalStrength = nc.mSignalStrength;
+        setUids(nc.mUids); // Will make the defensive copy
+        mEstablishingVpnAppUid = nc.mEstablishingVpnAppUid;
+        mUnwantedNetworkCapabilities = nc.mUnwantedNetworkCapabilities;
+        mSSID = nc.mSSID;
     }
 
     /**
@@ -275,6 +287,7 @@
      * this network can be used by system apps to upload telemetry data.
      * @hide
      */
+    @SystemApi
     public static final int NET_CAPABILITY_OEM_PAID = 22;
 
     private static final int MIN_NET_CAPABILITY = NET_CAPABILITY_MMS;
@@ -423,6 +436,7 @@
      * @return an array of capability values for this instance.
      * @hide
      */
+    @TestApi
     public @NetCapability int[] getCapabilities() {
         return BitUtils.unpackBits(mNetworkCapabilities);
     }
@@ -687,6 +701,7 @@
      * @return an array of transport type values for this instance.
      * @hide
      */
+    @TestApi
     public @Transport int[] getTransportTypes() {
         return BitUtils.unpackBits(mTransportTypes);
     }
@@ -920,7 +935,7 @@
     /**
      * Sets the signal strength. This is a signed integer, with higher values indicating a stronger
      * signal. The exact units are bearer-dependent. For example, Wi-Fi uses the same RSSI units
-     * reported by WifiManager.
+     * reported by wifi code.
      * <p>
      * Note that when used to register a network callback, this specifies the minimum acceptable
      * signal strength. When received as the state of an existing network it specifies the current
@@ -1052,7 +1067,7 @@
     }
 
     /**
-     * Tests if the set of UIDs that this network applies to is the same of the passed set of UIDs.
+     * Tests if the set of UIDs that this network applies to is the same as the passed network.
      * <p>
      * This test only checks whether equal range objects are in both sets. It will
      * return false if the ranges are not exactly the same, even if the covered UIDs
@@ -1142,6 +1157,62 @@
         mUids.addAll(nc.mUids);
     }
 
+
+    /**
+     * The SSID of the network, or null if not applicable or unknown.
+     * <p>
+     * This is filled in by wifi code.
+     * @hide
+     */
+    private String mSSID;
+
+    /**
+     * Sets the SSID of this network.
+     * @hide
+     */
+    public NetworkCapabilities setSSID(String ssid) {
+        mSSID = ssid;
+        return this;
+    }
+
+    /**
+     * Gets the SSID of this network, or null if none or unknown.
+     * @hide
+     */
+    public String getSSID() {
+        return mSSID;
+    }
+
+    /**
+     * Tests if the SSID of this network is the same as the SSID of the passed network.
+     * @hide
+     */
+    public boolean equalsSSID(NetworkCapabilities nc) {
+        return Objects.equals(mSSID, nc.mSSID);
+    }
+
+    /**
+     * Check if the SSID requirements of this object are matched by the passed object.
+     * @hide
+     */
+    public boolean satisfiedBySSID(NetworkCapabilities nc) {
+        return mSSID == null || mSSID.equals(nc.mSSID);
+    }
+
+    /**
+     * Combine SSIDs of the capabilities.
+     * <p>
+     * This is only legal if either the SSID of this object is null, or both SSIDs are
+     * equal.
+     * @hide
+     */
+    private void combineSSIDs(NetworkCapabilities nc) {
+        if (mSSID != null && !mSSID.equals(nc.mSSID)) {
+            throw new IllegalStateException("Can't combine two SSIDs");
+        }
+        setSSID(nc.mSSID);
+    }
+
     /**
      * Combine a set of Capabilities to this one.  Useful for coming up with the complete set.
      * <p>
@@ -1157,6 +1228,7 @@
         combineSpecifiers(nc);
         combineSignalStrength(nc);
         combineUids(nc);
+        combineSSIDs(nc);
     }
 
     /**
@@ -1175,7 +1247,8 @@
                 && (onlyImmutable || satisfiedByLinkBandwidths(nc))
                 && satisfiedBySpecifier(nc)
                 && (onlyImmutable || satisfiedBySignalStrength(nc))
-                && (onlyImmutable || satisfiedByUids(nc)));
+                && (onlyImmutable || satisfiedByUids(nc))
+                && (onlyImmutable || satisfiedBySSID(nc)));
     }
 
     /**
@@ -1264,7 +1337,8 @@
                 && equalsLinkBandwidths(that)
                 && equalsSignalStrength(that)
                 && equalsSpecifier(that)
-                && equalsUids(that));
+                && equalsUids(that)
+                && equalsSSID(that));
     }
 
     @Override
@@ -1279,7 +1353,8 @@
                 + (mLinkDownBandwidthKbps * 19)
                 + Objects.hashCode(mNetworkSpecifier) * 23
                 + (mSignalStrength * 29)
-                + Objects.hashCode(mUids) * 31;
+                + Objects.hashCode(mUids) * 31
+                + Objects.hashCode(mSSID) * 37;
     }
 
     @Override
@@ -1296,6 +1371,7 @@
         dest.writeParcelable((Parcelable) mNetworkSpecifier, flags);
         dest.writeInt(mSignalStrength);
         dest.writeArraySet(mUids);
+        dest.writeString(mSSID);
     }
 
     public static final Creator<NetworkCapabilities> CREATOR =
@@ -1313,6 +1389,7 @@
                 netCap.mSignalStrength = in.readInt();
                 netCap.mUids = (ArraySet<UidRange>) in.readArraySet(
                         null /* ClassLoader, null for default */);
+                netCap.mSSID = in.readString();
                 return netCap;
             }
             @Override
@@ -1363,6 +1440,10 @@
             sb.append(" EstablishingAppUid: ").append(mEstablishingVpnAppUid);
         }
 
+        if (null != mSSID) {
+            sb.append(" SSID: ").append(mSSID);
+        }
+
         sb.append("]");
         return sb.toString();
     }
diff --git a/core/java/android/net/NetworkRequest.java b/core/java/android/net/NetworkRequest.java
index 1ee0ed7..f3669fb 100644
--- a/core/java/android/net/NetworkRequest.java
+++ b/core/java/android/net/NetworkRequest.java
@@ -17,6 +17,8 @@
 package android.net;
 
 import android.annotation.NonNull;
+import android.net.NetworkCapabilities.NetCapability;
+import android.net.NetworkCapabilities.Transport;
 import android.os.Parcel;
 import android.os.Parcelable;
 import android.os.Process;
@@ -165,9 +167,6 @@
          * the requested network's required capabilities.  Note that when searching
          * for a network to satisfy a request, all capabilities requested must be
          * satisfied.
-         * <p>
-         * If the given capability was previously added to the list of unwanted capabilities
-         * then the capability will also be removed from the list of unwanted capabilities.
          *
          * @param capability The capability to add.
          * @return The builder to facilitate chaining
@@ -179,8 +178,7 @@
         }
 
         /**
-         * Removes (if found) the given capability from this builder instance from both required
-         * and unwanted capabilities lists.
+         * Removes (if found) the given capability from this builder instance.
          *
          * @param capability The capability to remove.
          * @return The builder to facilitate chaining.
@@ -199,8 +197,7 @@
          * @hide
          */
         public Builder setCapabilities(NetworkCapabilities nc) {
-            mNetworkCapabilities.clearAll();
-            mNetworkCapabilities.combineCapabilities(nc);
+            mNetworkCapabilities.set(nc);
             return this;
         }
 
@@ -228,6 +225,7 @@
          *
          * @param capability The capability to add to unwanted capability list.
          * @return The builder to facilitate chaining.
+         *
          * @hide
          */
         public Builder addUnwantedCapability(@NetworkCapabilities.NetCapability int capability) {
@@ -426,6 +424,29 @@
         return type == Type.BACKGROUND_REQUEST;
     }
 
+    /**
+     * @see Builder#addCapability(int)
+     */
+    public boolean hasCapability(@NetCapability int capability) {
+        return networkCapabilities.hasCapability(capability);
+    }
+
+    /**
+     * @see Builder#addUnwantedCapability(int)
+     *
+     * @hide
+     */
+    public boolean hasUnwantedCapability(@NetCapability int capability) {
+        return networkCapabilities.hasUnwantedCapability(capability);
+    }
+
+    /**
+     * @see Builder#addTransportType(int)
+     */
+    public boolean hasTransport(@Transport int transportType) {
+        return networkCapabilities.hasTransport(transportType);
+    }
+
     public String toString() {
         return "NetworkRequest [ " + type + " id=" + requestId +
                 (legacyType != ConnectivityManager.TYPE_NONE ? ", legacyType=" + legacyType : "") +
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 9838de1..2593690 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -102,6 +102,8 @@
 import android.os.ResultReceiver;
 import android.os.ServiceManager;
 import android.os.ServiceSpecificException;
+import android.os.ShellCallback;
+import android.os.ShellCommand;
 import android.os.SystemClock;
 import android.os.UserHandle;
 import android.os.UserManager;
@@ -234,8 +236,9 @@
 
     private KeyStore mKeyStore;
 
+    @VisibleForTesting
     @GuardedBy("mVpns")
-    private final SparseArray<Vpn> mVpns = new SparseArray<Vpn>();
+    protected final SparseArray<Vpn> mVpns = new SparseArray<Vpn>();
 
     // TODO: investigate if mLockdownEnabled can be removed and replaced everywhere by
     // a direct call to LockdownVpnTracker.isEnabled().
@@ -494,24 +497,24 @@
     private static final int MAX_VALIDATION_LOGS = 10;
     private static class ValidationLog {
         final Network mNetwork;
-        final String mNetworkExtraInfo;
+        final String mName;
         final ReadOnlyLocalLog mLog;
 
-        ValidationLog(Network network, String networkExtraInfo, ReadOnlyLocalLog log) {
+        ValidationLog(Network network, String name, ReadOnlyLocalLog log) {
             mNetwork = network;
-            mNetworkExtraInfo = networkExtraInfo;
+            mName = name;
             mLog = log;
         }
     }
     private final ArrayDeque<ValidationLog> mValidationLogs =
             new ArrayDeque<ValidationLog>(MAX_VALIDATION_LOGS);
 
-    private void addValidationLogs(ReadOnlyLocalLog log, Network network, String networkExtraInfo) {
+    private void addValidationLogs(ReadOnlyLocalLog log, Network network, String name) {
         synchronized (mValidationLogs) {
             while (mValidationLogs.size() >= MAX_VALIDATION_LOGS) {
                 mValidationLogs.removeLast();
             }
-            mValidationLogs.addFirst(new ValidationLog(network, networkExtraInfo, log));
+            mValidationLogs.addFirst(new ValidationLog(network, name, log));
         }
     }
 
@@ -898,6 +901,15 @@
                 deps);
     }
 
+    private static NetworkCapabilities createDefaultNetworkCapabilitiesForUid(int uid) {
+        final NetworkCapabilities netCap = new NetworkCapabilities();
+        netCap.addCapability(NET_CAPABILITY_INTERNET);
+        netCap.addCapability(NET_CAPABILITY_NOT_RESTRICTED);
+        netCap.removeCapability(NET_CAPABILITY_NOT_VPN);
+        netCap.setSingleUid(uid);
+        return netCap;
+    }
+
     private NetworkRequest createDefaultInternetRequestForTransport(
             int transportType, NetworkRequest.Type type) {
         NetworkCapabilities netCap = new NetworkCapabilities();
@@ -1150,12 +1162,20 @@
         int vpnNetId = NETID_UNSET;
         synchronized (mVpns) {
             final Vpn vpn = mVpns.get(user);
+            // TODO : now that capabilities contain the UID, the appliesToUid test should
+            // be removed as the satisfying test below should be enough.
             if (vpn != null && vpn.appliesToUid(uid)) vpnNetId = vpn.getNetId();
         }
         NetworkAgentInfo nai;
         if (vpnNetId != NETID_UNSET) {
             nai = getNetworkAgentInfoForNetId(vpnNetId);
-            if (nai != null) return nai.network;
+            if (nai != null) {
+                final NetworkCapabilities requiredCaps =
+                    createDefaultNetworkCapabilitiesForUid(uid);
+                if (requiredCaps.satisfiedByNetworkCapabilities(nai.networkCapabilities)) {
+                    return nai.network;
+                }
+            }
         }
         nai = getDefaultNetwork();
         if (nai != null
@@ -1352,7 +1372,8 @@
         if (nai != null) {
             synchronized (nai) {
                 if (nai.networkCapabilities != null) {
-                    return networkCapabilitiesWithoutUidsUnlessAllowed(nai.networkCapabilities,
+                    return networkCapabilitiesRestrictedForCallerPermissions(
+                            nai.networkCapabilities,
                             Binder.getCallingPid(), Binder.getCallingUid());
                 }
             }
@@ -1366,10 +1387,14 @@
         return getNetworkCapabilitiesInternal(getNetworkAgentInfoForNetwork(network));
     }
 
-    private NetworkCapabilities networkCapabilitiesWithoutUidsUnlessAllowed(
+    private NetworkCapabilities networkCapabilitiesRestrictedForCallerPermissions(
             NetworkCapabilities nc, int callerPid, int callerUid) {
-        if (checkSettingsPermission(callerPid, callerUid)) return new NetworkCapabilities(nc);
-        return new NetworkCapabilities(nc).setUids(null);
+        final NetworkCapabilities newNc = new NetworkCapabilities(nc);
+        if (!checkSettingsPermission(callerPid, callerUid)) {
+            newNc.setUids(null);
+            newNc.setSSID(null);
+        }
+        return newNc;
     }
 
     private void restrictRequestUidsForCaller(NetworkCapabilities nc) {
@@ -1415,7 +1440,8 @@
     public boolean isActiveNetworkMetered() {
         enforceAccessPermission();
 
-        final NetworkCapabilities caps = getNetworkCapabilities(getActiveNetwork());
+        final int uid = Binder.getCallingUid();
+        final NetworkCapabilities caps = getUnfilteredActiveNetworkState(uid).networkCapabilities;
         if (caps != null) {
             return !caps.hasCapability(NetworkCapabilities.NET_CAPABILITY_NOT_METERED);
         } else {
@@ -1737,6 +1763,7 @@
                             ni != null ? ni.getState().toString() : "?");
                 } catch (RemoteException e) {
                 }
+                intent.addFlags(Intent.FLAG_RECEIVER_VISIBLE_TO_INSTANT_APPS);
             }
             try {
                 mContext.sendStickyBroadcastAsUser(intent, UserHandle.ALL, options);
@@ -2032,7 +2059,7 @@
             synchronized (mValidationLogs) {
                 pw.println("mValidationLogs (most recent first):");
                 for (ValidationLog p : mValidationLogs) {
-                    pw.println(p.mNetwork + " - " + p.mNetworkExtraInfo);
+                    pw.println(p.mNetwork + " - " + p.mName);
                     pw.increaseIndent();
                     p.mLog.dump(fd, pw, args);
                     pw.decreaseIndent();
@@ -2377,94 +2404,107 @@
         }
     }
 
+    // This is a no-op if it's called with a message designating a network that has
+    // already been destroyed, because its reference will not be found in the relevant
+    // maps.
     private void handleAsyncChannelDisconnected(Message msg) {
         NetworkAgentInfo nai = mNetworkAgentInfos.get(msg.replyTo);
         if (nai != null) {
-            if (DBG) {
-                log(nai.name() + " got DISCONNECTED, was satisfying " + nai.numNetworkRequests());
-            }
-            // A network agent has disconnected.
-            // TODO - if we move the logic to the network agent (have them disconnect
-            // because they lost all their requests or because their score isn't good)
-            // then they would disconnect organically, report their new state and then
-            // disconnect the channel.
-            if (nai.networkInfo.isConnected()) {
-                nai.networkInfo.setDetailedState(NetworkInfo.DetailedState.DISCONNECTED,
-                        null, null);
-            }
-            final boolean wasDefault = isDefaultNetwork(nai);
-            if (wasDefault) {
-                mDefaultInetConditionPublished = 0;
-                // Log default network disconnection before required book-keeping.
-                // Let rematchAllNetworksAndRequests() below record a new default network event
-                // if there is a fallback. Taken together, the two form a X -> 0, 0 -> Y sequence
-                // whose timestamps tell how long it takes to recover a default network.
-                long now = SystemClock.elapsedRealtime();
-                metricsLogger().defaultNetworkMetrics().logDefaultNetworkEvent(now, null, nai);
-            }
-            notifyIfacesChangedForNetworkStats();
-            // TODO - we shouldn't send CALLBACK_LOST to requests that can be satisfied
-            // by other networks that are already connected. Perhaps that can be done by
-            // sending all CALLBACK_LOST messages (for requests, not listens) at the end
-            // of rematchAllNetworksAndRequests
-            notifyNetworkCallbacks(nai, ConnectivityManager.CALLBACK_LOST);
-            mKeepaliveTracker.handleStopAllKeepalives(nai,
-                    ConnectivityManager.PacketKeepalive.ERROR_INVALID_NETWORK);
-            for (String iface : nai.linkProperties.getAllInterfaceNames()) {
-                // Disable wakeup packet monitoring for each interface.
-                wakeupModifyInterface(iface, nai.networkCapabilities, false);
-            }
-            nai.networkMonitor.sendMessage(NetworkMonitor.CMD_NETWORK_DISCONNECTED);
-            mNetworkAgentInfos.remove(msg.replyTo);
-            nai.maybeStopClat();
-            synchronized (mNetworkForNetId) {
-                // Remove the NetworkAgent, but don't mark the netId as
-                // available until we've told netd to delete it below.
-                mNetworkForNetId.remove(nai.network.netId);
-            }
-            // Remove all previously satisfied requests.
-            for (int i = 0; i < nai.numNetworkRequests(); i++) {
-                NetworkRequest request = nai.requestAt(i);
-                NetworkAgentInfo currentNetwork = getNetworkForRequest(request.requestId);
-                if (currentNetwork != null && currentNetwork.network.netId == nai.network.netId) {
-                    clearNetworkForRequest(request.requestId);
-                    sendUpdatedScoreToFactories(request, 0);
-                }
-            }
-            nai.clearLingerState();
-            if (nai.isSatisfyingRequest(mDefaultRequest.requestId)) {
-                removeDataActivityTracking(nai);
-                notifyLockdownVpn(nai);
-                ensureNetworkTransitionWakelock(nai.name());
-            }
-            mLegacyTypeTracker.remove(nai, wasDefault);
-            rematchAllNetworksAndRequests(null, 0);
-            mLingerMonitor.noteDisconnect(nai);
-            if (nai.created) {
-                // Tell netd to clean up the configuration for this network
-                // (routing rules, DNS, etc).
-                // This may be slow as it requires a lot of netd shelling out to ip and
-                // ip[6]tables to flush routes and remove the incoming packet mark rule, so do it
-                // after we've rematched networks with requests which should make a potential
-                // fallback network the default or requested a new network from the
-                // NetworkFactories, so network traffic isn't interrupted for an unnecessarily
-                // long time.
-                try {
-                    mNetd.removeNetwork(nai.network.netId);
-                } catch (Exception e) {
-                    loge("Exception removing network: " + e);
-                }
-                mDnsManager.removeNetwork(nai.network);
-            }
-            synchronized (mNetworkForNetId) {
-                mNetIdInUse.delete(nai.network.netId);
-            }
+            disconnectAndDestroyNetwork(nai);
         } else {
             NetworkFactoryInfo nfi = mNetworkFactoryInfos.remove(msg.replyTo);
             if (DBG && nfi != null) log("unregisterNetworkFactory for " + nfi.name);
         }
     }
 
+    // Destroys a network, remove references to it from the internal state managed by
+    // ConnectivityService, free its interfaces and clean up.
+    // Must be called on the Handler thread.
+    private void disconnectAndDestroyNetwork(NetworkAgentInfo nai) {
+        if (DBG) {
+            log(nai.name() + " got DISCONNECTED, was satisfying " + nai.numNetworkRequests());
+        }
+        // A network agent has disconnected.
+        // TODO - if we move the logic to the network agent (have them disconnect
+        // because they lost all their requests or because their score isn't good)
+        // then they would disconnect organically, report their new state and then
+        // disconnect the channel.
+        if (nai.networkInfo.isConnected()) {
+            nai.networkInfo.setDetailedState(NetworkInfo.DetailedState.DISCONNECTED,
+                    null, null);
+        }
+        final boolean wasDefault = isDefaultNetwork(nai);
+        if (wasDefault) {
+            mDefaultInetConditionPublished = 0;
+            // Log default network disconnection before required book-keeping.
+            // Let rematchAllNetworksAndRequests() below record a new default network event
+            // if there is a fallback. Taken together, the two form a X -> 0, 0 -> Y sequence
+            // whose timestamps tell how long it takes to recover a default network.
+            long now = SystemClock.elapsedRealtime();
+            metricsLogger().defaultNetworkMetrics().logDefaultNetworkEvent(now, null, nai);
+        }
+        notifyIfacesChangedForNetworkStats();
+        // TODO - we shouldn't send CALLBACK_LOST to requests that can be satisfied
+        // by other networks that are already connected. Perhaps that can be done by
+        // sending all CALLBACK_LOST messages (for requests, not listens) at the end
+        // of rematchAllNetworksAndRequests
+        notifyNetworkCallbacks(nai, ConnectivityManager.CALLBACK_LOST);
+        mKeepaliveTracker.handleStopAllKeepalives(nai,
+                ConnectivityManager.PacketKeepalive.ERROR_INVALID_NETWORK);
+        for (String iface : nai.linkProperties.getAllInterfaceNames()) {
+            // Disable wakeup packet monitoring for each interface.
+            wakeupModifyInterface(iface, nai.networkCapabilities, false);
+        }
+        nai.networkMonitor.sendMessage(NetworkMonitor.CMD_NETWORK_DISCONNECTED);
+        mNetworkAgentInfos.remove(nai.messenger);
+        nai.maybeStopClat();
+        synchronized (mNetworkForNetId) {
+            // Remove the NetworkAgent, but don't mark the netId as
+            // available until we've told netd to delete it below.
+            mNetworkForNetId.remove(nai.network.netId);
+        }
+        // Remove all previously satisfied requests.
+        for (int i = 0; i < nai.numNetworkRequests(); i++) {
+            NetworkRequest request = nai.requestAt(i);
+            NetworkAgentInfo currentNetwork = getNetworkForRequest(request.requestId);
+            if (currentNetwork != null && currentNetwork.network.netId == nai.network.netId) {
+                clearNetworkForRequest(request.requestId);
+                sendUpdatedScoreToFactories(request, 0);
+            }
+        }
+        nai.clearLingerState();
+        if (nai.isSatisfyingRequest(mDefaultRequest.requestId)) {
+            removeDataActivityTracking(nai);
+            notifyLockdownVpn(nai);
+            ensureNetworkTransitionWakelock(nai.name());
+        }
+        mLegacyTypeTracker.remove(nai, wasDefault);
+        if (!nai.networkCapabilities.hasTransport(TRANSPORT_VPN)) {
+            updateAllVpnsCapabilities();
+        }
+        rematchAllNetworksAndRequests(null, 0);
+        mLingerMonitor.noteDisconnect(nai);
+        if (nai.created) {
+            // Tell netd to clean up the configuration for this network
+            // (routing rules, DNS, etc).
+            // This may be slow as it requires a lot of netd shelling out to ip and
+            // ip[6]tables to flush routes and remove the incoming packet mark rule, so do it
+            // after we've rematched networks with requests which should make a potential
+            // fallback network the default or requested a new network from the
+            // NetworkFactories, so network traffic isn't interrupted for an unnecessarily
+            // long time.
+            try {
+                mNetd.removeNetwork(nai.network.netId);
+            } catch (Exception e) {
+                loge("Exception removing network: " + e);
+            }
+            mDnsManager.removeNetwork(nai.network);
+        }
+        synchronized (mNetworkForNetId) {
+            mNetIdInUse.delete(nai.network.netId);
+        }
+    }
+
     // If this method proves to be too slow then we can maintain a separate
     // pendingIntent => NetworkRequestInfo map.
     // This method assumes that every non-null PendingIntent maps to exactly 1 NetworkRequestInfo.
@@ -3707,6 +3747,26 @@
         }
     }
 
+    /**
+     * Ask all VPN objects to recompute and update their capabilities.
+     *
+     * When underlying networks change, VPNs may have to update capabilities to reflect things
+     * like the metered bit, their transports, and so on. This asks the VPN objects to update
+     * their capabilities, and as this will cause them to send messages to the ConnectivityService
+     * handler thread through their agent, this is asynchronous. When the capabilities objects
+     * are computed they will be up-to-date as they are computed synchronously from here and
+     * this is running on the ConnectivityService thread.
+     * TODO : Fix this and call updateCapabilities inline to remove out-of-order events.
+     */
+    private void updateAllVpnsCapabilities() {
+        synchronized (mVpns) {
+            for (int i = 0; i < mVpns.size(); i++) {
+                final Vpn vpn = mVpns.valueAt(i);
+                vpn.updateCapabilities();
+            }
+        }
+    }
+
     @Override
     public boolean updateLockdownVpn() {
         if (Binder.getCallingUid() != Process.SYSTEM_UID) {
@@ -4181,6 +4241,15 @@
         }
     }
 
+    // This checks that the passed capabilities either do not request a specific SSID, or the
+    // calling app has permission to do so.
+    private void ensureSufficientPermissionsForRequest(NetworkCapabilities nc,
+            int callerPid, int callerUid) {
+        if (null != nc.getSSID() && !checkSettingsPermission(callerPid, callerUid)) {
+            throw new SecurityException("Insufficient permissions to request a specific SSID");
+        }
+    }
+
     private ArrayList<Integer> getSignalStrengthThresholds(NetworkAgentInfo nai) {
         final SortedSet<Integer> thresholds = new TreeSet();
         synchronized (nai) {
@@ -4238,8 +4307,7 @@
         // the default network request. This allows callers to keep track of
         // the system default network.
         if (type == NetworkRequest.Type.TRACK_DEFAULT) {
-            networkCapabilities = new NetworkCapabilities(mDefaultRequest.networkCapabilities);
-            networkCapabilities.removeCapability(NET_CAPABILITY_NOT_VPN);
+            networkCapabilities = createDefaultNetworkCapabilitiesForUid(Binder.getCallingUid());
             enforceAccessPermission();
         } else {
             networkCapabilities = new NetworkCapabilities(networkCapabilities);
@@ -4250,6 +4318,8 @@
             enforceMeteredApnPolicy(networkCapabilities);
         }
         ensureRequestableCapabilities(networkCapabilities);
+        ensureSufficientPermissionsForRequest(networkCapabilities,
+                Binder.getCallingPid(), Binder.getCallingUid());
         // Set the UID range for this request to the single UID of the requester, or to an empty
         // set of UIDs if the caller has the appropriate permission and UIDs have not been set.
         // This will overwrite any allowed UIDs in the requested capabilities. Though there
@@ -4328,6 +4398,8 @@
         enforceNetworkRequestPermissions(networkCapabilities);
         enforceMeteredApnPolicy(networkCapabilities);
         ensureRequestableCapabilities(networkCapabilities);
+        ensureSufficientPermissionsForRequest(networkCapabilities,
+                Binder.getCallingPid(), Binder.getCallingUid());
         ensureValidNetworkSpecifier(networkCapabilities);
         restrictRequestUidsForCaller(networkCapabilities);
 
@@ -4383,6 +4455,8 @@
         }
 
         NetworkCapabilities nc = new NetworkCapabilities(networkCapabilities);
+        ensureSufficientPermissionsForRequest(networkCapabilities,
+                Binder.getCallingPid(), Binder.getCallingUid());
         restrictRequestUidsForCaller(nc);
         // Apps without the CHANGE_NETWORK_STATE permission can't use background networks, so
         // make all their listens include NET_CAPABILITY_FOREGROUND. That way, they will get
@@ -4409,6 +4483,8 @@
             enforceAccessPermission();
         }
         ensureValidNetworkSpecifier(networkCapabilities);
+        ensureSufficientPermissionsForRequest(networkCapabilities,
+                Binder.getCallingPid(), Binder.getCallingUid());
 
         final NetworkCapabilities nc = new NetworkCapabilities(networkCapabilities);
         restrictRequestUidsForCaller(nc);
@@ -4541,8 +4617,10 @@
         synchronized (this) {
             nai.networkMonitor.systemReady = mSystemReady;
         }
-        addValidationLogs(nai.networkMonitor.getValidationLogs(), nai.network,
-                networkInfo.getExtraInfo());
+        final String extraInfo = networkInfo.getExtraInfo();
+        final String name = TextUtils.isEmpty(extraInfo)
+                ? nai.networkCapabilities.getSSID() : extraInfo;
+        addValidationLogs(nai.networkMonitor.getValidationLogs(), nai.network, name);
         if (DBG) log("registerNetworkAgent " + nai);
         mHandler.sendMessage(mHandler.obtainMessage(EVENT_REGISTER_NETWORK_AGENT, nai));
         return nai.network.netId;
@@ -4562,7 +4640,7 @@
     }
 
     private void updateLinkProperties(NetworkAgentInfo networkAgent, LinkProperties oldLp) {
-        LinkProperties newLp = networkAgent.linkProperties;
+        LinkProperties newLp = new LinkProperties(networkAgent.linkProperties);
         int netId = networkAgent.network.netId;
 
         // The NetworkAgentInfo does not know whether clatd is running on its network or not. Before
@@ -4596,6 +4674,9 @@
         }
         // TODO - move this check to cover the whole function
         if (!Objects.equals(newLp, oldLp)) {
+            synchronized (networkAgent) {
+                networkAgent.linkProperties = newLp;
+            }
             notifyIfacesChangedForNetworkStats();
             notifyNetworkCallbacks(networkAgent, ConnectivityManager.CALLBACK_IP_CHANGED);
         }
@@ -4846,12 +4927,7 @@
         if (!newNc.hasTransport(TRANSPORT_VPN)) {
             // Tell VPNs about updated capabilities, since they may need to
             // bubble those changes through.
-            synchronized (mVpns) {
-                for (int i = 0; i < mVpns.size(); i++) {
-                    final Vpn vpn = mVpns.valueAt(i);
-                    vpn.updateCapabilities();
-                }
-            }
+            updateAllVpnsCapabilities();
         }
     }
 
@@ -4980,7 +5056,7 @@
             }
             case ConnectivityManager.CALLBACK_CAP_CHANGED: {
                 // networkAgent can't be null as it has been accessed a few lines above.
-                final NetworkCapabilities nc = networkCapabilitiesWithoutUidsUnlessAllowed(
+                final NetworkCapabilities nc = networkCapabilitiesRestrictedForCallerPermissions(
                         networkAgent.networkCapabilities, nri.mPid, nri.mUid);
                 putParcelable(bundle, nc);
                 break;
@@ -5511,6 +5587,7 @@
                 }
                 updateUids(networkAgent, networkAgent.networkCapabilities, null);
             }
+            disconnectAndDestroyNetwork(networkAgent);
         } else if ((oldInfo != null && oldInfo.getState() == NetworkInfo.State.SUSPENDED) ||
                 state == NetworkInfo.State.SUSPENDED) {
             // going into or coming out of SUSPEND: rescore and notify
@@ -5812,4 +5889,61 @@
     private static int encodeBool(boolean b) {
         return b ? 1 : 0;
     }
-}
+
+    @Override
+    public void onShellCommand(FileDescriptor in, FileDescriptor out,
+            FileDescriptor err, String[] args, ShellCallback callback,
+            ResultReceiver resultReceiver) {
+        (new ShellCmd()).exec(this, in, out, err, args, callback, resultReceiver);
+    }
+
+    private class ShellCmd extends ShellCommand {
+
+        @Override
+        public int onCommand(String cmd) {
+            if (cmd == null) {
+                return handleDefaultCommands(cmd);
+            }
+            final PrintWriter pw = getOutPrintWriter();
+            try {
+                switch (cmd) {
+                    case "airplane-mode":
+                        final String action = getNextArg();
+                        if ("enable".equals(action)) {
+                            setAirplaneMode(true);
+                            return 0;
+                        } else if ("disable".equals(action)) {
+                            setAirplaneMode(false);
+                            return 0;
+                        } else if (action == null) {
+                            final ContentResolver cr = mContext.getContentResolver();
+                            final int enabled = Settings.Global.getInt(cr,
+                                    Settings.Global.AIRPLANE_MODE_ON);
+                            pw.println(enabled == 0 ? "disabled" : "enabled");
+                            return 0;
+                        } else {
+                            onHelp();
+                            return -1;
+                        }
+                    default:
+                        return handleDefaultCommands(cmd);
+                }
+            } catch (Exception e) {
+                pw.println(e);
+            }
+            return -1;
+        }
+
+        @Override
+        public void onHelp() {
+            PrintWriter pw = getOutPrintWriter();
+            pw.println("Connectivity service commands:");
+            pw.println("  help");
+            pw.println("    Print this help text.");
+            pw.println("  airplane-mode [enable|disable]");
+            pw.println("    Turn airplane mode on or off.");
+            pw.println("  airplane-mode");
+            pw.println("    Get airplane mode.");
+        }
+    }
+}
\ No newline at end of file
diff --git a/services/core/java/com/android/server/connectivity/NetworkNotificationManager.java b/services/core/java/com/android/server/connectivity/NetworkNotificationManager.java
index 0d935db..36a2476 100644
--- a/services/core/java/com/android/server/connectivity/NetworkNotificationManager.java
+++ b/services/core/java/com/android/server/connectivity/NetworkNotificationManager.java
@@ -23,8 +23,10 @@
 import android.content.Intent;
 import android.content.res.Resources;
 import android.net.NetworkCapabilities;
+import android.net.wifi.WifiInfo;
 import android.os.UserHandle;
 import android.telephony.TelephonyManager;
+import android.text.TextUtils;
 import android.util.Slog;
 import android.util.SparseArray;
 import android.util.SparseIntArray;
@@ -130,16 +132,17 @@
         final String tag = tagFor(id);
         final int eventId = notifyType.eventId;
         final int transportType;
-        final String extraInfo;
+        final String name;
         if (nai != null) {
             transportType = getFirstTransportType(nai);
-            extraInfo = nai.networkInfo.getExtraInfo();
+            final String extraInfo = nai.networkInfo.getExtraInfo();
+            name = TextUtils.isEmpty(extraInfo) ? nai.networkCapabilities.getSSID() : extraInfo;
             // Only notify for Internet-capable networks.
             if (!nai.networkCapabilities.hasCapability(NET_CAPABILITY_INTERNET)) return;
         } else {
             // Legacy notifications.
             transportType = TRANSPORT_CELLULAR;
-            extraInfo = null;
+            name = null;
         }
 
         // Clear any previous notification with lower priority, otherwise return. http://b/63676954.
@@ -156,9 +159,8 @@
 
         if (DBG) {
             Slog.d(TAG, String.format(
-                    "showNotification tag=%s event=%s transport=%s extraInfo=%s highPrioriy=%s",
-                    tag, nameOf(eventId), getTransportName(transportType), extraInfo,
-                    highPriority));
+                    "showNotification tag=%s event=%s transport=%s name=%s highPriority=%s",
+                    tag, nameOf(eventId), getTransportName(transportType), name, highPriority));
         }
 
         Resources r = Resources.getSystem();
@@ -176,7 +178,8 @@
             switch (transportType) {
                 case TRANSPORT_WIFI:
                     title = r.getString(R.string.wifi_available_sign_in, 0);
-                    details = r.getString(R.string.network_available_sign_in_detailed, extraInfo);
+                    details = r.getString(R.string.network_available_sign_in_detailed,
+                            WifiInfo.removeDoubleQuotes(nai.networkCapabilities.getSSID()));
                     break;
                 case TRANSPORT_CELLULAR:
                     title = r.getString(R.string.network_available_sign_in, 0);
@@ -186,7 +189,7 @@
                     break;
                 default:
                     title = r.getString(R.string.network_available_sign_in, 0);
-                    details = r.getString(R.string.network_available_sign_in_detailed, extraInfo);
+                    details = r.getString(R.string.network_available_sign_in_detailed, name);
                     break;
             }
         } else if (notifyType == NotificationType.NETWORK_SWITCH) {
diff --git a/services/net/java/android/net/apf/ApfCapabilities.java b/services/net/java/android/net/apf/ApfCapabilities.java
index 703b415..dec8ca2 100644
--- a/services/net/java/android/net/apf/ApfCapabilities.java
+++ b/services/net/java/android/net/apf/ApfCapabilities.java
@@ -49,4 +49,14 @@
         return String.format("%s{version: %d, maxSize: %d, format: %d}", getClass().getSimpleName(),
                 apfVersionSupported, maximumApfProgramSize, apfPacketFormat);
     }
+
+    /**
+     * Returns true if the APF interpreter advertises support for the data buffer access opcodes
+     * LDDW and STDW.
+     *
+     * Full LDDW and STDW support is present from APFv4 on.
+     */
+    public boolean hasDataAccess() {
+        return apfVersionSupported >= 4;
+    }
 }
diff --git a/tests/net/java/android/net/NetworkCapabilitiesTest.java b/tests/net/java/android/net/NetworkCapabilitiesTest.java
index 0696592..2f4d69e 100644
--- a/tests/net/java/android/net/NetworkCapabilitiesTest.java
+++ b/tests/net/java/android/net/NetworkCapabilitiesTest.java
@@ -40,12 +40,14 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 import android.os.Parcel;
 import android.support.test.runner.AndroidJUnit4;
 import android.test.suitebuilder.annotation.SmallTest;
 import android.util.ArraySet;
 
+
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -54,6 +56,9 @@
 @RunWith(AndroidJUnit4.class)
 @SmallTest
 public class NetworkCapabilitiesTest {
+    private static final String TEST_SSID = "TEST_SSID";
+    private static final String DIFFERENT_TEST_SSID = "DIFFERENT_TEST_SSID";
+
     @Test
     public void testMaybeMarkCapabilitiesRestricted() {
         // verify EIMS is restricted
@@ -268,6 +273,8 @@
             .addCapability(NET_CAPABILITY_EIMS)
             .addCapability(NET_CAPABILITY_NOT_METERED);
         assertEqualsThroughMarshalling(netCap);
+        netCap.setSSID(TEST_SSID);
+        assertEqualsThroughMarshalling(netCap);
     }
 
     @Test
@@ -362,6 +369,27 @@
     }
 
     @Test
+    public void testSSID() {
+        NetworkCapabilities nc1 = new NetworkCapabilities();
+        NetworkCapabilities nc2 = new NetworkCapabilities();
+        assertTrue(nc2.satisfiedBySSID(nc1));
+
+        nc1.setSSID(TEST_SSID);
+        assertTrue(nc2.satisfiedBySSID(nc1));
+        nc2.setSSID("different " + TEST_SSID);
+        assertFalse(nc2.satisfiedBySSID(nc1));
+
+        assertTrue(nc1.satisfiedByImmutableNetworkCapabilities(nc2));
+        assertFalse(nc1.satisfiedByNetworkCapabilities(nc2));
+    }
+
+    private ArraySet<UidRange> uidRange(int from, int to) {
+        final ArraySet<UidRange> range = new ArraySet<>(1);
+        range.add(new UidRange(from, to));
+        return range;
+    }
+
+    @Test
     public void testCombineCapabilities() {
         NetworkCapabilities nc1 = new NetworkCapabilities();
         NetworkCapabilities nc2 = new NetworkCapabilities();
@@ -382,6 +410,35 @@
         // will never be satisfied.
         assertTrue(nc2.hasCapability(NET_CAPABILITY_NOT_ROAMING));
         assertTrue(nc2.hasUnwantedCapability(NET_CAPABILITY_NOT_ROAMING));
+
+        nc1.setSSID(TEST_SSID);
+        nc2.combineCapabilities(nc1);
+        assertTrue(TEST_SSID.equals(nc2.getSSID()));
+
+        // Because they now have the same SSID, the following call should not throw
+        nc2.combineCapabilities(nc1);
+
+        nc1.setSSID(DIFFERENT_TEST_SSID);
+        try {
+            nc2.combineCapabilities(nc1);
+            fail("Expected IllegalStateException: can't combine different SSIDs");
+        } catch (IllegalStateException expected) {}
+        nc1.setSSID(TEST_SSID);
+
+        nc1.setUids(uidRange(10, 13));
+        assertNotEquals(nc1, nc2);
+        nc2.combineCapabilities(nc1);  // Everything + 10~13 is still everything.
+        assertNotEquals(nc1, nc2);
+        nc1.combineCapabilities(nc2);  // 10~13 + everything is everything.
+        assertEquals(nc1, nc2);
+        nc1.setUids(uidRange(10, 13));
+        nc2.setUids(uidRange(20, 23));
+        assertNotEquals(nc1, nc2);
+        nc1.combineCapabilities(nc2);
+        assertTrue(nc1.appliesToUid(12));
+        assertFalse(nc2.appliesToUid(12));
+        assertTrue(nc1.appliesToUid(22));
+        assertTrue(nc2.appliesToUid(22));
     }
 
     @Test
@@ -420,4 +477,38 @@
         p.setDataPosition(0);
         assertEquals(NetworkCapabilities.CREATOR.createFromParcel(p), netCap);
     }
+
+    @Test
+    public void testSet() {
+        NetworkCapabilities nc1 = new NetworkCapabilities();
+        NetworkCapabilities nc2 = new NetworkCapabilities();
+
+        nc1.addUnwantedCapability(NET_CAPABILITY_CAPTIVE_PORTAL);
+        nc1.addCapability(NET_CAPABILITY_NOT_ROAMING);
+        assertNotEquals(nc1, nc2);
+        nc2.set(nc1);
+        assertEquals(nc1, nc2);
+        assertTrue(nc2.hasCapability(NET_CAPABILITY_NOT_ROAMING));
+        assertTrue(nc2.hasUnwantedCapability(NET_CAPABILITY_CAPTIVE_PORTAL));
+
+        // This will effectively move NOT_ROAMING capability from required to unwanted for nc1.
+        nc1.addUnwantedCapability(NET_CAPABILITY_NOT_ROAMING);
+        nc1.setSSID(TEST_SSID);
+        nc2.set(nc1);
+        assertEquals(nc1, nc2);
+        // Contrary to combineCapabilities, set() will have removed the NOT_ROAMING capability
+        // from nc2.
+        assertFalse(nc2.hasCapability(NET_CAPABILITY_NOT_ROAMING));
+        assertTrue(nc2.hasUnwantedCapability(NET_CAPABILITY_NOT_ROAMING));
+        assertTrue(TEST_SSID.equals(nc2.getSSID()));
+
+        nc1.setSSID(DIFFERENT_TEST_SSID);
+        nc2.set(nc1);
+        assertEquals(nc1, nc2);
+        assertTrue(DIFFERENT_TEST_SSID.equals(nc2.getSSID()));
+
+        nc1.setUids(uidRange(10, 13));
+        nc2.set(nc1);  // Overwrites, as opposed to combineCapabilities
+        assertEquals(nc1, nc2);
+    }
 }
diff --git a/tests/net/java/android/net/NetworkStatsHistoryTest.java b/tests/net/java/android/net/NetworkStatsHistoryTest.java
index 1c0c14e..301d04d 100644
--- a/tests/net/java/android/net/NetworkStatsHistoryTest.java
+++ b/tests/net/java/android/net/NetworkStatsHistoryTest.java
@@ -32,9 +32,14 @@
 import static android.text.format.DateUtils.SECOND_IN_MILLIS;
 import static android.text.format.DateUtils.WEEK_IN_MILLIS;
 import static android.text.format.DateUtils.YEAR_IN_MILLIS;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
-import android.test.AndroidTestCase;
-import android.test.suitebuilder.annotation.SmallTest;
+import android.content.Context;
+import android.support.test.InstrumentationRegistry;
+import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
 import android.test.suitebuilder.annotation.Suppress;
 import android.util.Log;
 
@@ -46,25 +51,31 @@
 import java.io.DataOutputStream;
 import java.util.Random;
 
+import org.junit.After;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(AndroidJUnit4.class)
 @SmallTest
-public class NetworkStatsHistoryTest extends AndroidTestCase {
+public class NetworkStatsHistoryTest {
     private static final String TAG = "NetworkStatsHistoryTest";
 
     private static final long TEST_START = 1194220800000L;
 
     private NetworkStatsHistory stats;
 
-    @Override
-    protected void tearDown() throws Exception {
-        super.tearDown();
+    @After
+    public void tearDown() throws Exception {
         if (stats != null) {
             assertConsistent(stats);
         }
     }
 
+    @Test
     public void testReadOriginalVersion() throws Exception {
-        final DataInputStream in = new DataInputStream(
-                getContext().getResources().openRawResource(R.raw.history_v1));
+        final Context context = InstrumentationRegistry.getContext();
+        final DataInputStream in =
+                new DataInputStream(context.getResources().openRawResource(R.raw.history_v1));
 
         NetworkStatsHistory.Entry entry = null;
         try {
@@ -88,6 +99,7 @@
         }
     }
 
+    @Test
     public void testRecordSingleBucket() throws Exception {
         final long BUCKET_SIZE = HOUR_IN_MILLIS;
         stats = new NetworkStatsHistory(BUCKET_SIZE);
@@ -100,6 +112,7 @@
         assertValues(stats, 0, SECOND_IN_MILLIS, 1024L, 10L, 2048L, 20L, 2L);
     }
 
+    @Test
     public void testRecordEqualBuckets() throws Exception {
         final long bucketDuration = HOUR_IN_MILLIS;
         stats = new NetworkStatsHistory(bucketDuration);
@@ -114,6 +127,7 @@
         assertValues(stats, 1, HOUR_IN_MILLIS / 2, 512L, 5L, 64L, 1L, 1L);
     }
 
+    @Test
     public void testRecordTouchingBuckets() throws Exception {
         final long BUCKET_SIZE = 15 * MINUTE_IN_MILLIS;
         stats = new NetworkStatsHistory(BUCKET_SIZE);
@@ -134,6 +148,7 @@
         assertValues(stats, 2, 4 * MINUTE_IN_MILLIS, 200L, 400L, 1000L, 2000L, 20L);
     }
 
+    @Test
     public void testRecordGapBuckets() throws Exception {
         final long BUCKET_SIZE = HOUR_IN_MILLIS;
         stats = new NetworkStatsHistory(BUCKET_SIZE);
@@ -165,6 +180,7 @@
         assertValues(stats, 3, SECOND_IN_MILLIS, 64L, 1L, 512L, 8L, 2L);
     }
 
+    @Test
     public void testRecordOverlapBuckets() throws Exception {
         final long BUCKET_SIZE = HOUR_IN_MILLIS;
         stats = new NetworkStatsHistory(BUCKET_SIZE);
@@ -182,6 +198,7 @@
         assertValues(stats, 1, (HOUR_IN_MILLIS / 2), 512L, 5L, 512L, 5L, 5L);
     }
 
+    @Test
     public void testRecordEntireGapIdentical() throws Exception {
         // first, create two separate histories far apart
         final NetworkStatsHistory stats1 = new NetworkStatsHistory(HOUR_IN_MILLIS);
@@ -206,6 +223,7 @@
         assertValues(stats, 3, 500L, 250L);
     }
 
+    @Test
     public void testRecordEntireOverlapVaryingBuckets() throws Exception {
         // create history just over hour bucket boundary
         final NetworkStatsHistory stats1 = new NetworkStatsHistory(HOUR_IN_MILLIS);
@@ -247,6 +265,7 @@
         assertValues(stats, 3, 150L, 150L);
     }
 
+    @Test
     public void testRemove() throws Exception {
         stats = new NetworkStatsHistory(HOUR_IN_MILLIS);
 
@@ -280,6 +299,7 @@
         assertEquals(0, stats.size());
     }
 
+    @Test
     public void testTotalData() throws Exception {
         final long BUCKET_SIZE = HOUR_IN_MILLIS;
         stats = new NetworkStatsHistory(BUCKET_SIZE);
@@ -304,7 +324,7 @@
 
     }
 
-    @Suppress
+    @Test
     public void testFuzzing() throws Exception {
         try {
             // fuzzing with random events, looking for crashes
@@ -341,6 +361,7 @@
         return value < 0 ? -value : value;
     }
 
+    @Test
     public void testIgnoreFields() throws Exception {
         final NetworkStatsHistory history = new NetworkStatsHistory(
                 MINUTE_IN_MILLIS, 0, FIELD_RX_BYTES | FIELD_TX_BYTES);
@@ -353,6 +374,7 @@
         assertFullValues(history, UNKNOWN, 1026L, UNKNOWN, 2050L, UNKNOWN, UNKNOWN);
     }
 
+    @Test
     public void testIgnoreFieldsRecordIn() throws Exception {
         final NetworkStatsHistory full = new NetworkStatsHistory(MINUTE_IN_MILLIS, 0, FIELD_ALL);
         final NetworkStatsHistory partial = new NetworkStatsHistory(
@@ -365,6 +387,7 @@
         assertFullValues(partial, UNKNOWN, UNKNOWN, 10L, UNKNOWN, UNKNOWN, 4L);
     }
 
+    @Test
     public void testIgnoreFieldsRecordOut() throws Exception {
         final NetworkStatsHistory full = new NetworkStatsHistory(MINUTE_IN_MILLIS, 0, FIELD_ALL);
         final NetworkStatsHistory partial = new NetworkStatsHistory(
@@ -377,6 +400,7 @@
         assertFullValues(full, MINUTE_IN_MILLIS, 0L, 10L, 0L, 0L, 4L);
     }
 
+    @Test
     public void testSerialize() throws Exception {
         final NetworkStatsHistory before = new NetworkStatsHistory(MINUTE_IN_MILLIS, 40, FIELD_ALL);
         before.recordData(0, 4 * MINUTE_IN_MILLIS,
@@ -396,6 +420,7 @@
         assertFullValues(after, 5 * MINUTE_IN_MILLIS, 1034L, 30L, 2078L, 60L, 54L);
     }
 
+    @Test
     public void testVarLong() throws Exception {
         assertEquals(0L, performVarLong(0L));
         assertEquals(-1L, performVarLong(-1L));
@@ -409,6 +434,7 @@
         assertEquals(Long.MAX_VALUE - 40, performVarLong(Long.MAX_VALUE - 40));
     }
 
+    @Test
     public void testIndexBeforeAfter() throws Exception {
         final long BUCKET_SIZE = HOUR_IN_MILLIS;
         stats = new NetworkStatsHistory(BUCKET_SIZE);
@@ -451,6 +477,7 @@
         assertIndexBeforeAfter(stats, 4, 4, Long.MAX_VALUE);
     }
 
+    @Test
     public void testIntersects() throws Exception {
         final long BUCKET_SIZE = HOUR_IN_MILLIS;
         stats = new NetworkStatsHistory(BUCKET_SIZE);
@@ -485,6 +512,7 @@
         assertTrue(stats.intersects(Long.MIN_VALUE, TEST_START + 1));
     }
 
+    @Test
     public void testSetValues() throws Exception {
         stats = new NetworkStatsHistory(HOUR_IN_MILLIS);
         stats.recordData(TEST_START, TEST_START + 1,
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index e4df974..d6018f1 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -84,6 +84,7 @@
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
+import android.content.pm.UserInfo;
 import android.content.res.Resources;
 import android.net.CaptivePortal;
 import android.net.ConnectivityManager;
@@ -111,6 +112,8 @@
 import android.net.RouteInfo;
 import android.net.StringNetworkSpecifier;
 import android.net.UidRange;
+import android.net.VpnService;
+import android.net.captiveportal.CaptivePortalProbeResult;
 import android.net.metrics.IpConnectivityLog;
 import android.net.util.MultinetworkPolicyTracker;
 import android.os.ConditionVariable;
@@ -132,6 +135,7 @@
 import android.util.ArraySet;
 import android.util.Log;
 
+import com.android.internal.net.VpnConfig;
 import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.WakeupMessage;
 import com.android.internal.util.test.BroadcastInterceptingContext;
@@ -195,6 +199,7 @@
     private MockNetworkAgent mWiFiNetworkAgent;
     private MockNetworkAgent mCellNetworkAgent;
     private MockNetworkAgent mEthernetNetworkAgent;
+    private MockVpn mMockVpn;
     private Context mContext;
 
     @Mock IpConnectivityMetrics.Logger mMetricsService;
@@ -476,6 +481,14 @@
             mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
         }
 
+        public void setNetworkCapabilities(NetworkCapabilities nc,
+                boolean sendToConnectivityService) {
+            mNetworkCapabilities.set(nc);
+            if (sendToConnectivityService) {
+                mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
+            }
+        }
+
         public void connectWithoutInternet() {
             mNetworkInfo.setDetailedState(DetailedState.CONNECTED, null, null);
             mNetworkAgent.sendNetworkInfo(mNetworkInfo);
@@ -505,6 +518,7 @@
                 mWrappedNetworkMonitor.gen204ProbeResult = 204;
                 NetworkRequest request = new NetworkRequest.Builder()
                         .addTransportType(mNetworkCapabilities.getTransportTypes()[0])
+                        .clearCapabilities()
                         .build();
                 callback = new NetworkCallback() {
                     public void onCapabilitiesChanged(Network network,
@@ -590,6 +604,10 @@
             return mRedirectUrl;
         }
 
+        public NetworkAgent getNetworkAgent() {
+            return mNetworkAgent;
+        }
+
         public NetworkCapabilities getNetworkCapabilities() {
             return mNetworkCapabilities;
         }
@@ -722,6 +740,87 @@
         }
     }
 
+    private static Looper startHandlerThreadAndReturnLooper() {
+        final HandlerThread handlerThread = new HandlerThread("MockVpnThread");
+        handlerThread.start();
+        return handlerThread.getLooper();
+    }
+
+    private class MockVpn extends Vpn {
+        // TODO : the interactions between this mock and the mock network agent are too
+        // hard to get right at this moment, because it's unclear in which case which
+        // target needs to get a method call or both, and in what order. It's because
+        // MockNetworkAgent wants to manage its own NetworkCapabilities, but the Vpn
+        // parent class of MockVpn agent wants that responsibility.
+        // That being said inside the test it should be possible to make the interactions
+        // harder to get wrong with precise speccing, judicious comments, helper methods
+        // and a few sprinkled assertions.
+
+        private boolean mConnected = false;
+        // Careful ! This is different from mNetworkAgent, because MockNetworkAgent does
+        // not inherit from NetworkAgent.
+        private MockNetworkAgent mMockNetworkAgent;
+
+        public MockVpn(int userId) {
+            super(startHandlerThreadAndReturnLooper(), mServiceContext, mNetworkManagementService,
+                    userId);
+        }
+
+        public void setNetworkAgent(MockNetworkAgent agent) {
+            waitForIdle(agent, TIMEOUT_MS);
+            mMockNetworkAgent = agent;
+            mNetworkAgent = agent.getNetworkAgent();
+            mNetworkCapabilities.set(agent.getNetworkCapabilities());
+        }
+
+        public void setUids(Set<UidRange> uids) {
+            mNetworkCapabilities.setUids(uids);
+            updateCapabilities();
+        }
+
+        @Override
+        public int getNetId() {
+            return mMockNetworkAgent.getNetwork().netId;
+        }
+
+        @Override
+        public boolean appliesToUid(int uid) {
+            return mConnected;  // Trickery to simplify testing.
+        }
+
+        @Override
+        protected boolean isCallerEstablishedOwnerLocked() {
+            return mConnected;  // Similar trickery
+        }
+
+        public void connect() {
+            mNetworkCapabilities.set(mMockNetworkAgent.getNetworkCapabilities());
+            mConnected = true;
+            mConfig = new VpnConfig();
+        }
+
+        @Override
+        public void updateCapabilities() {
+            if (!mConnected) return;
+            super.updateCapabilities();
+            // Because super.updateCapabilities will update the capabilities of the agent but not
+            // the mock agent, the mock agent needs to know about them.
+            copyCapabilitiesToNetworkAgent();
+        }
+
+        private void copyCapabilitiesToNetworkAgent() {
+            if (null != mMockNetworkAgent) {
+                mMockNetworkAgent.setNetworkCapabilities(mNetworkCapabilities,
+                        false /* sendToConnectivityService */);
+            }
+        }
+
+        public void disconnect() {
+            mConnected = false;
+            mConfig = null;
+        }
+    }
+
     private class FakeWakeupMessage extends WakeupMessage {
         private static final int UNREASONABLY_LONG_WAIT = 1000;
 
@@ -888,6 +987,17 @@
             return mLastCreatedNetworkMonitor;
         }
 
+        public void mockVpn(int uid) {
+            synchronized (mVpns) {
+                int userId = UserHandle.getUserId(uid);
+                mMockVpn = new MockVpn(userId);
+                // This has no effect unless the VPN is actually connected, because things like
+                // getActiveNetworkForUidInternal call getNetworkAgentInfoForNetId on the VPN
+                // netId, and check if that network is actually connected.
+                mVpns.put(userId, mMockVpn);
+            }
+        }
+
         public void waitForIdle(int timeoutMs) {
             waitForIdleHandler(mHandlerThread, timeoutMs);
         }
@@ -931,8 +1041,9 @@
                 mock(INetworkPolicyManager.class),
                 mock(IpConnectivityLog.class));
 
-        mService.systemReady();
         mCm = new WrappedConnectivityManager(InstrumentationRegistry.getContext(), mService);
+        mService.systemReady();
+        mService.mockVpn(Process.myUid());
         mCm.bindProcessToNetwork(null);
 
         // Ensure that the default setting for Captive Portals is used for most tests
@@ -1345,6 +1456,7 @@
         private final static int TIMEOUT_MS = 100;
 
         private final LinkedBlockingQueue<CallbackInfo> mCallbacks = new LinkedBlockingQueue<>();
+        private Network mLastAvailableNetwork;
 
         protected void setLastCallback(CallbackState state, Network network, Object o) {
             mCallbacks.offer(new CallbackInfo(state, network, o));
@@ -1352,6 +1464,7 @@
 
         @Override
         public void onAvailable(Network network) {
+            mLastAvailableNetwork = network;
             setLastCallback(CallbackState.AVAILABLE, network, null);
         }
 
@@ -1387,9 +1500,14 @@
 
         @Override
         public void onLost(Network network) {
+            mLastAvailableNetwork = null;
             setLastCallback(CallbackState.LOST, network, null);
         }
 
+        public Network getLastAvailableNetwork() {
+            return mLastAvailableNetwork;
+        }
+
         CallbackInfo nextCallback(int timeoutMs) {
             CallbackInfo cb = null;
             try {
@@ -1525,7 +1643,8 @@
 
         void expectCapabilitiesLike(Predicate<NetworkCapabilities> fn, MockNetworkAgent agent) {
             CallbackInfo cbi = expectCallback(CallbackState.NETWORK_CAPABILITIES, agent);
-            assertTrue(fn.test((NetworkCapabilities) cbi.arg));
+            assertTrue("Received capabilities don't match expectations : " + cbi.arg,
+                    fn.test((NetworkCapabilities) cbi.arg));
         }
 
         void assertNoCallback() {
@@ -1656,6 +1775,7 @@
         callback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
         defaultCallback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         mWiFiNetworkAgent.connect(true);
         // We get AVAILABLE on wifi when wifi connects and satisfies our unmetered request.
@@ -1666,6 +1786,7 @@
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         mEthernetNetworkAgent.connect(true);
         callback.expectAvailableCallbacksUnvalidated(mEthernetNetworkAgent);
@@ -1674,11 +1795,13 @@
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mEthernetNetworkAgent);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mEthernetNetworkAgent);
         assertEquals(mEthernetNetworkAgent.getNetwork(), mCm.getActiveNetwork());
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         mEthernetNetworkAgent.disconnect();
         callback.expectCallback(CallbackState.LOST, mEthernetNetworkAgent);
         defaultCallback.expectCallback(CallbackState.LOST, mEthernetNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         for (int i = 0; i < 4; i++) {
             MockNetworkAgent oldNetwork, newNetwork;
@@ -1707,6 +1830,7 @@
         defaultCallback.expectCapabilitiesWithout(NET_CAPABILITY_NOT_METERED, mWiFiNetworkAgent);
         defaultCallback.assertNoCallback();
         callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Wifi no longer satisfies our listen, which is for an unmetered network.
         // But because its score is 55, it's still up (and the default network).
@@ -1716,8 +1840,11 @@
         mWiFiNetworkAgent.disconnect();
         defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
         mCellNetworkAgent.disconnect();
         defaultCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        waitForIdle();
+        assertEquals(null, mCm.getActiveNetwork());
 
         mCm.unregisterNetworkCallback(callback);
         waitForIdle();
@@ -1734,6 +1861,7 @@
         callback.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Bring up wifi with a score of 20.
         // Cell stays up because it would satisfy the default request if it validated.
@@ -1742,12 +1870,14 @@
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         mWiFiNetworkAgent.disconnect();
         callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
         defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Bring up wifi with a score of 70.
         // Cell is lingered because it would not satisfy any request, even if it validated.
@@ -1758,6 +1888,7 @@
         callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Tear down wifi.
         mWiFiNetworkAgent.disconnect();
@@ -1765,6 +1896,7 @@
         defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Bring up wifi, then validate it. Previous versions would immediately tear down cell, but
         // it's arguably correct to linger it, since it was the default network before it validated.
@@ -1776,6 +1908,7 @@
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         defaultCallback.expectAvailableThenValidatedCallbacks(mWiFiNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         mWiFiNetworkAgent.disconnect();
         callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
@@ -1784,12 +1917,15 @@
         mCellNetworkAgent.disconnect();
         callback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
         defaultCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        waitForIdle();
+        assertEquals(null, mCm.getActiveNetwork());
 
         // If a network is lingering, and we add and remove a request from it, resume lingering.
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
         callback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
         defaultCallback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
@@ -1797,6 +1933,7 @@
         // TODO: Investigate sending validated before losing.
         callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         NetworkRequest cellRequest = new NetworkRequest.Builder()
                 .addTransportType(TRANSPORT_CELLULAR).build();
@@ -1813,6 +1950,7 @@
         callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
         defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Cell is now the default network. Pin it with a cell-specific request.
         noopCallback = new NetworkCallback();  // Can't reuse NetworkCallbacks. http://b/20701525
@@ -1823,6 +1961,7 @@
         mWiFiNetworkAgent.connect(true);
         callback.expectAvailableThenValidatedCallbacks(mWiFiNetworkAgent);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
         // The default request is lingering on cell, but nothing happens to cell, and we send no
         // callbacks for it, because it's kept up by cellRequest.
         callback.assertNoCallback();
@@ -1846,6 +1985,7 @@
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mEthernetNetworkAgent);
         trackDefaultCallback.expectAvailableDoubleValidatedCallbacks(mEthernetNetworkAgent);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mEthernetNetworkAgent);
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Let linger run its course.
         callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent, lingerTimeoutMs);
@@ -2494,23 +2634,27 @@
         mCellNetworkAgent.connect(true);
         cellNetworkCallback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
         defaultNetworkCallback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
+        assertEquals(defaultNetworkCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Bring up wifi and expect CALLBACK_AVAILABLE.
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
         cellNetworkCallback.assertNoCallback();
         defaultNetworkCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
+        assertEquals(defaultNetworkCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Bring down cell. Expect no default network callback, since it wasn't the default.
         mCellNetworkAgent.disconnect();
         cellNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
         defaultNetworkCallback.assertNoCallback();
+        assertEquals(defaultNetworkCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Bring up cell. Expect no default network callback, since it won't be the default.
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
         cellNetworkCallback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
         defaultNetworkCallback.assertNoCallback();
+        assertEquals(defaultNetworkCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Bring down wifi. Expect the default network callback to notified of LOST wifi
         // followed by AVAILABLE cell.
@@ -2521,6 +2665,24 @@
         mCellNetworkAgent.disconnect();
         cellNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
         defaultNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        waitForIdle();
+        assertEquals(null, mCm.getActiveNetwork());
+
+        final int uid = Process.myUid();
+        final MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        final ArraySet<UidRange> ranges = new ArraySet<>();
+        ranges.add(new UidRange(uid, uid));
+        mMockVpn.setNetworkAgent(vpnNetworkAgent);
+        mMockVpn.setUids(ranges);
+        vpnNetworkAgent.connect(true);
+        mMockVpn.connect();
+        defaultNetworkCallback.expectAvailableThenValidatedCallbacks(vpnNetworkAgent);
+        assertEquals(defaultNetworkCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
+
+        vpnNetworkAgent.disconnect();
+        defaultNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        waitForIdle();
+        assertEquals(null, mCm.getActiveNetwork());
     }
 
     @Test
@@ -4011,6 +4173,7 @@
         final TestNetworkCallback genericNotVpnNetworkCallback = new TestNetworkCallback();
         final TestNetworkCallback wifiNetworkCallback = new TestNetworkCallback();
         final TestNetworkCallback vpnNetworkCallback = new TestNetworkCallback();
+        final TestNetworkCallback defaultCallback = new TestNetworkCallback();
         final NetworkRequest genericNotVpnRequest = new NetworkRequest.Builder().build();
         final NetworkRequest genericRequest = new NetworkRequest.Builder()
                 .removeCapability(NET_CAPABILITY_NOT_VPN).build();
@@ -4023,6 +4186,8 @@
         mCm.registerNetworkCallback(genericNotVpnRequest, genericNotVpnNetworkCallback);
         mCm.registerNetworkCallback(wifiRequest, wifiNetworkCallback);
         mCm.registerNetworkCallback(vpnNetworkRequest, vpnNetworkCallback);
+        mCm.registerDefaultNetworkCallback(defaultCallback);
+        defaultCallback.assertNoCallback();
 
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
@@ -4030,26 +4195,30 @@
         genericNetworkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         genericNotVpnNetworkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         wifiNetworkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
+        defaultCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         vpnNetworkCallback.assertNoCallback();
-
-        // TODO : check callbacks agree with the return value of mCm.getActiveNetwork().
-        // Right now this is not possible because establish() is not adequately instrumented
-        // in this test.
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         final MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
         final ArraySet<UidRange> ranges = new ArraySet<>();
         ranges.add(new UidRange(uid, uid));
-        vpnNetworkAgent.setUids(ranges);
+        mMockVpn.setNetworkAgent(vpnNetworkAgent);
+        mMockVpn.setUids(ranges);
         vpnNetworkAgent.connect(false);
+        mMockVpn.connect();
 
         genericNetworkCallback.expectAvailableCallbacksUnvalidated(vpnNetworkAgent);
         genericNotVpnNetworkCallback.assertNoCallback();
         wifiNetworkCallback.assertNoCallback();
         vpnNetworkCallback.expectAvailableCallbacksUnvalidated(vpnNetworkAgent);
+        defaultCallback.expectAvailableCallbacksUnvalidated(vpnNetworkAgent);
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         genericNetworkCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES, vpnNetworkAgent);
         genericNotVpnNetworkCallback.assertNoCallback();
         vpnNetworkCallback.expectCapabilitiesLike(nc -> null == nc.getUids(), vpnNetworkAgent);
+        defaultCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES, vpnNetworkAgent);
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         ranges.clear();
         vpnNetworkAgent.setUids(ranges);
@@ -4059,13 +4228,24 @@
         wifiNetworkCallback.assertNoCallback();
         vpnNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
 
+        // TODO : The default network callback should actually get a LOST call here (also see the
+        // comment below for AVAILABLE). This is because ConnectivityService does not look at UID
+        // ranges at all when determining whether a network should be rematched. In practice, VPNs
+        // can't currently update their UIDs without disconnecting, so this does not matter too
+        // much, but that is the reason the test here has to check for an update to the
+        // capabilities instead of the expected LOST then AVAILABLE.
+        defaultCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES, vpnNetworkAgent);
+
         ranges.add(new UidRange(uid, uid));
-        vpnNetworkAgent.setUids(ranges);
+        mMockVpn.setUids(ranges);
 
         genericNetworkCallback.expectAvailableCallbacksValidated(vpnNetworkAgent);
         genericNotVpnNetworkCallback.assertNoCallback();
         wifiNetworkCallback.assertNoCallback();
         vpnNetworkCallback.expectAvailableCallbacksValidated(vpnNetworkAgent);
+        // TODO : Here like above, AVAILABLE would be correct, but because this can't actually
+        // happen outside of the test, ConnectivityService does not rematch callbacks.
+        defaultCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES, vpnNetworkAgent);
 
         mWiFiNetworkAgent.disconnect();
 
@@ -4073,6 +4253,7 @@
         genericNotVpnNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
         wifiNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
         vpnNetworkCallback.assertNoCallback();
+        defaultCallback.assertNoCallback();
 
         vpnNetworkAgent.disconnect();
 
@@ -4080,9 +4261,161 @@
         genericNotVpnNetworkCallback.assertNoCallback();
         wifiNetworkCallback.assertNoCallback();
         vpnNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        defaultCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        assertEquals(null, mCm.getActiveNetwork());
 
         mCm.unregisterNetworkCallback(genericNetworkCallback);
         mCm.unregisterNetworkCallback(wifiNetworkCallback);
         mCm.unregisterNetworkCallback(vpnNetworkCallback);
+        mCm.unregisterNetworkCallback(defaultCallback);
+    }
+
+    @Test
+    public void testVpnWithAndWithoutInternet() {
+        final int uid = Process.myUid();
+
+        final TestNetworkCallback defaultCallback = new TestNetworkCallback();
+        mCm.registerDefaultNetworkCallback(defaultCallback);
+        defaultCallback.assertNoCallback();
+
+        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent.connect(true);
+
+        defaultCallback.expectAvailableThenValidatedCallbacks(mWiFiNetworkAgent);
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
+
+        MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        final ArraySet<UidRange> ranges = new ArraySet<>();
+        ranges.add(new UidRange(uid, uid));
+        mMockVpn.setNetworkAgent(vpnNetworkAgent);
+        mMockVpn.setUids(ranges);
+        vpnNetworkAgent.connect(true /* validated */, false /* hasInternet */);
+        mMockVpn.connect();
+
+        defaultCallback.assertNoCallback();
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
+
+        vpnNetworkAgent.disconnect();
+        defaultCallback.assertNoCallback();
+
+        vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        mMockVpn.setNetworkAgent(vpnNetworkAgent);
+        mMockVpn.setUids(ranges);
+        vpnNetworkAgent.connect(true /* validated */, true /* hasInternet */);
+        mMockVpn.connect();
+        defaultCallback.expectAvailableThenValidatedCallbacks(vpnNetworkAgent);
+        assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
+
+        vpnNetworkAgent.disconnect();
+        defaultCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        defaultCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
+
+        vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        ranges.clear();
+        mMockVpn.setNetworkAgent(vpnNetworkAgent);
+        mMockVpn.setUids(ranges);
+        vpnNetworkAgent.connect(false /* validated */, true /* hasInternet */);
+        mMockVpn.connect();
+        defaultCallback.assertNoCallback();
+
+        mCm.unregisterNetworkCallback(defaultCallback);
+    }
+
+    @Test
+    public void testVpnSetUnderlyingNetworks() {
+        final int uid = Process.myUid();
+
+        final TestNetworkCallback vpnNetworkCallback = new TestNetworkCallback();
+        final NetworkRequest vpnNetworkRequest = new NetworkRequest.Builder()
+                .removeCapability(NET_CAPABILITY_NOT_VPN)
+                .addTransportType(TRANSPORT_VPN)
+                .build();
+        NetworkCapabilities nc;
+        mCm.registerNetworkCallback(vpnNetworkRequest, vpnNetworkCallback);
+        vpnNetworkCallback.assertNoCallback();
+
+        final MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        final ArraySet<UidRange> ranges = new ArraySet<>();
+        ranges.add(new UidRange(uid, uid));
+        mMockVpn.setNetworkAgent(vpnNetworkAgent);
+        mMockVpn.connect();
+        mMockVpn.setUids(ranges);
+        vpnNetworkAgent.connect(true /* validated */, false /* hasInternet */);
+
+        vpnNetworkCallback.expectAvailableThenValidatedCallbacks(vpnNetworkAgent);
+        nc = mCm.getNetworkCapabilities(vpnNetworkAgent.getNetwork());
+        assertTrue(nc.hasTransport(TRANSPORT_VPN));
+        assertFalse(nc.hasTransport(TRANSPORT_CELLULAR));
+        assertFalse(nc.hasTransport(TRANSPORT_WIFI));
+        // For safety reasons a VPN without underlying networks is considered metered.
+        assertFalse(nc.hasCapability(NET_CAPABILITY_NOT_METERED));
+
+        // Connect cell and use it as an underlying network.
+        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent.connect(true);
+
+        mService.setUnderlyingNetworksForVpn(
+                new Network[] { mCellNetworkAgent.getNetwork() });
+
+        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+                && caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
+                vpnNetworkAgent);
+
+        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent.addCapability(NET_CAPABILITY_NOT_METERED);
+        mWiFiNetworkAgent.connect(true);
+
+        mService.setUnderlyingNetworksForVpn(
+                new Network[] { mCellNetworkAgent.getNetwork(), mWiFiNetworkAgent.getNetwork() });
+
+        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+                && caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
+                vpnNetworkAgent);
+
+        // Don't disconnect, but note the VPN is not using wifi any more.
+        mService.setUnderlyingNetworksForVpn(
+                new Network[] { mCellNetworkAgent.getNetwork() });
+
+        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+                && caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
+                vpnNetworkAgent);
+
+        // Use Wifi but not cell. Note the VPN is now unmetered.
+        mService.setUnderlyingNetworksForVpn(
+                new Network[] { mWiFiNetworkAgent.getNetwork() });
+
+        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+                && !caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
+                && caps.hasCapability(NET_CAPABILITY_NOT_METERED),
+                vpnNetworkAgent);
+
+        // Use both again.
+        mService.setUnderlyingNetworksForVpn(
+                new Network[] { mCellNetworkAgent.getNetwork(), mWiFiNetworkAgent.getNetwork() });
+
+        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+                && caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
+                vpnNetworkAgent);
+
+        // Disconnect cell. Receive update without even removing the dead network from the
+        // underlying networks – it's dead anyway. Not metered any more.
+        mCellNetworkAgent.disconnect();
+        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+                && !caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
+                && caps.hasCapability(NET_CAPABILITY_NOT_METERED),
+                vpnNetworkAgent);
+
+        // Disconnect wifi too. No underlying networks means this is now metered.
+        mWiFiNetworkAgent.disconnect();
+        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+                && !caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
+                vpnNetworkAgent);
+
+        mMockVpn.disconnect();
     }
 }
diff --git a/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java b/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java
index da0a48a..6f14332 100644
--- a/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java
+++ b/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java
@@ -26,6 +26,9 @@
 import static android.os.Process.myUid;
 import static android.text.format.DateUtils.HOUR_IN_MILLIS;
 import static android.text.format.DateUtils.MINUTE_IN_MILLIS;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.fail;
 
 import static com.android.server.net.NetworkStatsCollection.multiplySafe;
 
@@ -37,11 +40,12 @@
 import android.net.NetworkTemplate;
 import android.os.Process;
 import android.os.UserHandle;
+import android.support.test.InstrumentationRegistry;
+import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
 import android.telephony.SubscriptionPlan;
 import android.telephony.TelephonyManager;
-import android.test.AndroidTestCase;
 import android.test.MoreAsserts;
-import android.test.suitebuilder.annotation.SmallTest;
 import android.text.format.DateUtils;
 import android.util.RecurrenceRule;
 
@@ -64,11 +68,17 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
 /**
  * Tests for {@link NetworkStatsCollection}.
  */
+@RunWith(AndroidJUnit4.class)
 @SmallTest
-public class NetworkStatsCollectionTest extends AndroidTestCase {
+public class NetworkStatsCollectionTest {
 
     private static final String TEST_FILE = "test.bin";
     private static final String TEST_IMSI = "310260000000000";
@@ -79,18 +89,15 @@
 
     private static Clock sOriginalClock;
 
-    @Override
+    @Before
     public void setUp() throws Exception {
-        super.setUp();
         sOriginalClock = RecurrenceRule.sClock;
-
         // ignore any device overlay while testing
         NetworkTemplate.forceAllNetworkTypes();
     }
 
-    @Override
-    protected void tearDown() throws Exception {
-        super.tearDown();
+    @After
+    public void tearDown() throws Exception {
         RecurrenceRule.sClock = sOriginalClock;
     }
 
@@ -98,8 +105,10 @@
         RecurrenceRule.sClock = Clock.fixed(instant, ZoneId.systemDefault());
     }
 
+    @Test
     public void testReadLegacyNetwork() throws Exception {
-        final File testFile = new File(getContext().getFilesDir(), TEST_FILE);
+        final File testFile =
+                new File(InstrumentationRegistry.getContext().getFilesDir(), TEST_FILE);
         stageFile(R.raw.netstats_v1, testFile);
 
         final NetworkStatsCollection collection = new NetworkStatsCollection(30 * MINUTE_IN_MILLIS);
@@ -124,8 +133,10 @@
                 636016770L, 709306L, 88038768L, 518836L, NetworkStatsAccess.Level.DEVICE);
     }
 
+    @Test
     public void testReadLegacyUid() throws Exception {
-        final File testFile = new File(getContext().getFilesDir(), TEST_FILE);
+        final File testFile =
+                new File(InstrumentationRegistry.getContext().getFilesDir(), TEST_FILE);
         stageFile(R.raw.netstats_uid_v4, testFile);
 
         final NetworkStatsCollection collection = new NetworkStatsCollection(30 * MINUTE_IN_MILLIS);
@@ -150,8 +161,10 @@
                 637076152L, 711413L, 88343717L, 521022L, NetworkStatsAccess.Level.DEVICE);
     }
 
+    @Test
     public void testReadLegacyUidTags() throws Exception {
-        final File testFile = new File(getContext().getFilesDir(), TEST_FILE);
+        final File testFile =
+                new File(InstrumentationRegistry.getContext().getFilesDir(), TEST_FILE);
         stageFile(R.raw.netstats_uid_v4, testFile);
 
         final NetworkStatsCollection collection = new NetworkStatsCollection(30 * MINUTE_IN_MILLIS);
@@ -176,6 +189,7 @@
                 77017831L, 100995L, 35436758L, 92344L);
     }
 
+    @Test
     public void testStartEndAtomicBuckets() throws Exception {
         final NetworkStatsCollection collection = new NetworkStatsCollection(HOUR_IN_MILLIS);
 
@@ -190,6 +204,7 @@
         assertEquals(2 * HOUR_IN_MILLIS, collection.getEndMillis());
     }
 
+    @Test
     public void testAccessLevels() throws Exception {
         final NetworkStatsCollection collection = new NetworkStatsCollection(HOUR_IN_MILLIS);
         final NetworkStats.Entry entry = new NetworkStats.Entry();
@@ -250,8 +265,10 @@
                 0, NetworkStatsAccess.Level.DEVICE);
     }
 
+    @Test
     public void testAugmentPlan() throws Exception {
-        final File testFile = new File(getContext().getFilesDir(), TEST_FILE);
+        final File testFile =
+                new File(InstrumentationRegistry.getContext().getFilesDir(), TEST_FILE);
         stageFile(R.raw.netstats_v1, testFile);
 
         final NetworkStatsCollection emptyCollection = new NetworkStatsCollection(30 * MINUTE_IN_MILLIS);
@@ -439,6 +456,7 @@
         }
     }
 
+    @Test
     public void testAugmentPlanGigantic() throws Exception {
         // We're in the future, but not that far off
         setClock(Instant.parse("2012-06-01T00:00:00.00Z"));
@@ -461,6 +479,7 @@
         assertEquals(4_939_212_386L, getHistory(large, plan, TIME_A, TIME_C).getTotalBytes());
     }
 
+    @Test
     public void testRounding() throws Exception {
         final NetworkStatsCollection coll = new NetworkStatsCollection(HOUR_IN_MILLIS);
 
@@ -482,6 +501,7 @@
         assertEquals(TIME_A - HOUR_IN_MILLIS, coll.roundDown(TIME_A - 1));
     }
 
+    @Test
     public void testMultiplySafe() {
         assertEquals(25, multiplySafe(50, 1, 2));
         assertEquals(100, multiplySafe(50, 2, 1));
@@ -510,7 +530,7 @@
         InputStream in = null;
         OutputStream out = null;
         try {
-            in = getContext().getResources().openRawResource(rawId);
+            in = InstrumentationRegistry.getContext().getResources().openRawResource(rawId);
             out = new FileOutputStream(file);
             Streams.copy(in, out);
         } finally {
diff --git a/tests/net/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/net/java/com/android/server/net/NetworkStatsServiceTest.java
index 2fb9faf..6836626 100644
--- a/tests/net/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/net/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -93,8 +93,8 @@
 import android.os.PowerManager;
 import android.support.test.InstrumentationRegistry;
 import android.support.test.runner.AndroidJUnit4;
+import android.support.test.filters.SmallTest;
 import android.telephony.TelephonyManager;
-import android.test.suitebuilder.annotation.SmallTest;
 import android.util.Log;
 import android.util.TrustedTime;
 
@@ -216,7 +216,6 @@
               ArgumentCaptor.forClass(INetworkManagementEventObserver.class);
         verify(mNetManager).registerObserver(networkObserver.capture());
         mNetworkObserver = networkObserver.getValue();
-
     }
 
     @After