[automerger skipped] Merge "Get resource based on subId for multi-SIM" into qt-r1-dev
am: 559e7a3de1 -s ours
am skip reason: change_id Ib5b085d97103889600773d269e03b939c29ca47d with SHA1 5fed8d9ee3 is in history

Change-Id: I405bb51787bce13441aa4c3b985ef2b92df18826
diff --git a/core/java/android/net/NetworkAgent.java b/core/java/android/net/NetworkAgent.java
index 9cd4f53..43ea589 100644
--- a/core/java/android/net/NetworkAgent.java
+++ b/core/java/android/net/NetworkAgent.java
@@ -178,7 +178,21 @@
      */
     public static final int EVENT_SOCKET_KEEPALIVE = BASE + 13;
 
-    // TODO: move the above 2 constants down so they are in order once merge conflicts are resolved
+    /**
+     * Sent by ConnectivityService to inform this network transport of signal strength thresholds
+     * that when crossed should trigger a system wakeup and a NetworkCapabilities update.
+     *
+     *   obj = int[] describing signal strength thresholds.
+     */
+    public static final int CMD_SET_SIGNAL_STRENGTH_THRESHOLDS = BASE + 14;
+
+    /**
+     * Sent by ConnectivityService to the NeworkAgent to inform the agent to avoid
+     * automatically reconnecting to this network (e.g. via autojoin).  Happens
+     * when user selects "No" option on the "Stay connected?" dialog box.
+     */
+    public static final int CMD_PREVENT_AUTOMATIC_RECONNECT = BASE + 15;
+
     /**
      * Sent by the KeepaliveTracker to NetworkAgent to add a packet filter.
      *
@@ -198,21 +212,6 @@
      */
     public static final int CMD_REMOVE_KEEPALIVE_PACKET_FILTER = BASE + 17;
 
-    /**
-     * Sent by ConnectivityService to inform this network transport of signal strength thresholds
-     * that when crossed should trigger a system wakeup and a NetworkCapabilities update.
-     *
-     *   obj = int[] describing signal strength thresholds.
-     */
-    public static final int CMD_SET_SIGNAL_STRENGTH_THRESHOLDS = BASE + 14;
-
-    /**
-     * Sent by ConnectivityService to the NeworkAgent to inform the agent to avoid
-     * automatically reconnecting to this network (e.g. via autojoin).  Happens
-     * when user selects "No" option on the "Stay connected?" dialog box.
-     */
-    public static final int CMD_PREVENT_AUTOMATIC_RECONNECT = BASE + 15;
-
     // TODO : remove these two constructors. They are a stopgap measure to help sheperding a number
     // of dependent changes that would conflict throughout the automerger graph. Having these
     // temporarily helps with the process of going through with all these dependent changes across
@@ -438,15 +437,23 @@
     }
 
     /**
-     * Called by the bearer to indicate this network was manually selected by the user.
+     * Called by the bearer to indicate whether the network was manually selected by the user.
      * This should be called before the NetworkInfo is marked CONNECTED so that this
-     * Network can be given special treatment at that time. If {@code acceptUnvalidated} is
-     * {@code true}, then the system will switch to this network. If it is {@code false} and the
-     * network cannot be validated, the system will ask the user whether to switch to this network.
-     * If the user confirms and selects "don't ask again", then the system will call
-     * {@link #saveAcceptUnvalidated} to persist the user's choice. Thus, if the transport ever
-     * calls this method with {@code acceptUnvalidated} set to {@code false}, it must also implement
-     * {@link #saveAcceptUnvalidated} to respect the user's choice.
+     * Network can be given special treatment at that time.
+     *
+     * If {@code explicitlySelected} is {@code true}, and {@code acceptUnvalidated} is {@code true},
+     * then the system will switch to this network. If {@code explicitlySelected} is {@code true}
+     * and {@code acceptUnvalidated} is {@code false}, and the  network cannot be validated, the
+     * system will ask the user whether to switch to this network.  If the user confirms and selects
+     * "don't ask again", then the system will call {@link #saveAcceptUnvalidated} to persist the
+     * user's choice. Thus, if the transport ever calls this method with {@code explicitlySelected}
+     * set to {@code true} and {@code acceptUnvalidated} set to {@code false}, it must also
+     * implement {@link #saveAcceptUnvalidated} to respect the user's choice.
+     *
+     * If  {@code explicitlySelected} is {@code false} and {@code acceptUnvalidated} is
+     * {@code true}, the system will interpret this as the user having accepted partial connectivity
+     * on this network. Thus, the system will switch to the network and consider it validated even
+     * if it only provides partial connectivity, but the network is not otherwise treated specially.
      */
     public void explicitlySelected(boolean explicitlySelected, boolean acceptUnvalidated) {
         queueOrSendMessage(EVENT_SET_EXPLICITLY_SELECTED,
diff --git a/core/java/android/net/StaticIpConfiguration.java b/core/java/android/net/StaticIpConfiguration.java
index 050fcfa..d6deba5 100644
--- a/core/java/android/net/StaticIpConfiguration.java
+++ b/core/java/android/net/StaticIpConfiguration.java
@@ -113,8 +113,8 @@
     }
 
     /**
-     * Get a {@link String} listing in priority order of the comma separated domains to search when
-     * resolving host names on the link.
+     * Get a {@link String} containing the comma separated domains to search when resolving host
+     * names on this link, in priority order.
      */
     public @Nullable String getDomains() {
         return domains;
diff --git a/core/java/android/net/TestNetworkManager.java b/core/java/android/net/TestNetworkManager.java
index e274005..4ac4a69 100644
--- a/core/java/android/net/TestNetworkManager.java
+++ b/core/java/android/net/TestNetworkManager.java
@@ -56,6 +56,26 @@
     /**
      * Sets up a capability-limited, testing-only network for a given interface
      *
+     * @param lp The LinkProperties for the TestNetworkService to use for this test network. Note
+     *     that the interface name and link addresses will be overwritten, and the passed-in values
+     *     discarded.
+     * @param isMetered Whether or not the network should be considered metered.
+     * @param binder A binder object guarding the lifecycle of this test network.
+     * @hide
+     */
+    public void setupTestNetwork(
+            @NonNull LinkProperties lp, boolean isMetered, @NonNull IBinder binder) {
+        Preconditions.checkNotNull(lp, "Invalid LinkProperties");
+        try {
+            mService.setupTestNetwork(lp.getInterfaceName(), lp, isMetered, binder);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * Sets up a capability-limited, testing-only network for a given interface
+     *
      * @param iface the name of the interface to be used for the Network LinkProperties.
      * @param binder A binder object guarding the lifecycle of this test network.
      * @hide
@@ -63,7 +83,7 @@
     @TestApi
     public void setupTestNetwork(@NonNull String iface, @NonNull IBinder binder) {
         try {
-            mService.setupTestNetwork(iface, binder);
+            mService.setupTestNetwork(iface, null, true, binder);
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }
diff --git a/core/java/android/net/util/DnsUtils.java b/core/java/android/net/util/DnsUtils.java
index e6abd50..7908353 100644
--- a/core/java/android/net/util/DnsUtils.java
+++ b/core/java/android/net/util/DnsUtils.java
@@ -141,14 +141,17 @@
      */
     public static @NonNull List<InetAddress> rfc6724Sort(@Nullable Network network,
             @NonNull List<InetAddress> answers) {
-        List<SortableAddress> sortableAnswerList = new ArrayList<>();
-        answers.forEach(addr -> sortableAnswerList.add(
-                new SortableAddress(addr, findSrcAddress(network, addr))));
+        final ArrayList<SortableAddress> sortableAnswerList = new ArrayList<>();
+        for (InetAddress addr : answers) {
+            sortableAnswerList.add(new SortableAddress(addr, findSrcAddress(network, addr)));
+        }
 
         Collections.sort(sortableAnswerList, sRfc6724Comparator);
 
         final List<InetAddress> sortedAnswers = new ArrayList<>();
-        sortableAnswerList.forEach(ans -> sortedAnswers.add(ans.address));
+        for (SortableAddress ans : sortableAnswerList) {
+            sortedAnswers.add(ans.address);
+        }
 
         return sortedAnswers;
     }
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index e5fddef..e67ccc4 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -150,7 +150,6 @@
 import android.util.Pair;
 import android.util.Slog;
 import android.util.SparseArray;
-import android.util.SparseBooleanArray;
 import android.util.SparseIntArray;
 import android.util.Xml;
 
@@ -168,7 +167,6 @@
 import com.android.internal.util.DumpUtils;
 import com.android.internal.util.IndentingPrintWriter;
 import com.android.internal.util.MessageUtils;
-import com.android.internal.util.WakeupMessage;
 import com.android.internal.util.XmlUtils;
 import com.android.server.am.BatteryStatsService;
 import com.android.server.connectivity.AutodestructReference;
@@ -193,7 +191,6 @@
 import com.android.server.net.BaseNetworkObserver;
 import com.android.server.net.LockdownVpnTracker;
 import com.android.server.net.NetworkPolicyManagerInternal;
-import com.android.server.net.NetworkStatsFactory;
 import com.android.server.utils.PriorityDump;
 
 import com.google.android.collect.Lists;
@@ -306,7 +303,8 @@
     /** Flag indicating if background data is restricted. */
     private boolean mRestrictBackground;
 
-    final private Context mContext;
+    private final Context mContext;
+    private final Dependencies mDeps;
     // 0 is full bad, 100 is full good
     private int mDefaultInetConditionPublished = 0;
 
@@ -587,11 +585,6 @@
     private NetworkNotificationManager mNotifier;
     private LingerMonitor mLingerMonitor;
 
-    // sequence number for Networks; keep in sync with system/netd/NetworkController.cpp
-    private static final int MIN_NET_ID = 100; // some reserved marks
-    private static final int MAX_NET_ID = 65535 - 0x0400; // Top 1024 bits reserved by IpSecService
-    private int mNextNetId = MIN_NET_ID;
-
     // sequence number of NetworkRequests
     private int mNextNetworkRequestId = 1;
 
@@ -835,19 +828,113 @@
         }
     };
 
+    /**
+     * Dependencies of ConnectivityService, for injection in tests.
+     */
+    @VisibleForTesting
+    public static class Dependencies {
+        /**
+         * Get system properties to use in ConnectivityService.
+         */
+        public MockableSystemProperties getSystemProperties() {
+            return new MockableSystemProperties();
+        }
+
+        /**
+         * Create a HandlerThread to use in ConnectivityService.
+         */
+        public HandlerThread makeHandlerThread() {
+            return new HandlerThread("ConnectivityServiceThread");
+        }
+
+        /**
+         * Get a reference to the NetworkStackClient.
+         */
+        public NetworkStackClient getNetworkStack() {
+            return NetworkStackClient.getInstance();
+        }
+
+        /**
+         * @see Tethering
+         */
+        public Tethering makeTethering(@NonNull Context context,
+                @NonNull INetworkManagementService nms,
+                @NonNull INetworkStatsService statsService,
+                @NonNull INetworkPolicyManager policyManager,
+                @NonNull TetheringDependencies tetheringDeps) {
+            return new Tethering(context, nms, statsService, policyManager,
+                    IoThread.get().getLooper(), getSystemProperties(), tetheringDeps);
+        }
+
+        /**
+         * @see ProxyTracker
+         */
+        public ProxyTracker makeProxyTracker(@NonNull Context context,
+                @NonNull Handler connServiceHandler) {
+            return new ProxyTracker(context, connServiceHandler, EVENT_PROXY_HAS_CHANGED);
+        }
+
+        /**
+         * @see NetIdManager
+         */
+        public NetIdManager makeNetIdManager() {
+            return new NetIdManager();
+        }
+
+        /**
+         * @see NetworkUtils#queryUserAccess(int, int)
+         */
+        public boolean queryUserAccess(int uid, int netId) {
+            return NetworkUtils.queryUserAccess(uid, netId);
+        }
+
+        /**
+         * @see MultinetworkPolicyTracker
+         */
+        public MultinetworkPolicyTracker makeMultinetworkPolicyTracker(
+                @NonNull Context c, @NonNull Handler h, @NonNull Runnable r) {
+            return new MultinetworkPolicyTracker(c, h, r);
+        }
+
+        /**
+         * @see ServiceManager#checkService(String)
+         */
+        public boolean hasService(@NonNull String name) {
+            return ServiceManager.checkService(name) != null;
+        }
+
+        /**
+         * @see IpConnectivityMetrics.Logger
+         */
+        public IpConnectivityMetrics.Logger getMetricsLogger() {
+            return checkNotNull(LocalServices.getService(IpConnectivityMetrics.Logger.class),
+                    "no IpConnectivityMetrics service");
+        }
+
+        /**
+         * @see IpConnectivityMetrics
+         */
+        public IIpConnectivityMetrics getIpConnectivityMetrics() {
+            return IIpConnectivityMetrics.Stub.asInterface(
+                    ServiceManager.getService(IpConnectivityLog.SERVICE_NAME));
+        }
+    }
+
     public ConnectivityService(Context context, INetworkManagementService netManager,
             INetworkStatsService statsService, INetworkPolicyManager policyManager) {
-        this(context, netManager, statsService, policyManager,
-            getDnsResolver(), new IpConnectivityLog(), NetdService.getInstance());
+        this(context, netManager, statsService, policyManager, getDnsResolver(),
+                new IpConnectivityLog(), NetdService.getInstance(), new Dependencies());
     }
 
     @VisibleForTesting
     protected ConnectivityService(Context context, INetworkManagementService netManager,
             INetworkStatsService statsService, INetworkPolicyManager policyManager,
-            IDnsResolver dnsresolver, IpConnectivityLog logger, INetd netd) {
+            IDnsResolver dnsresolver, IpConnectivityLog logger, INetd netd, Dependencies deps) {
         if (DBG) log("ConnectivityService starting up");
 
-        mSystemProperties = getSystemProperties();
+        mDeps = checkNotNull(deps, "missing Dependencies");
+        mSystemProperties = mDeps.getSystemProperties();
+        mNetIdManager = mDeps.makeNetIdManager();
 
         mMetricsLog = logger;
         mDefaultRequest = createDefaultInternetRequestForTransport(-1, NetworkRequest.Type.REQUEST);
@@ -864,7 +951,7 @@
         mDefaultWifiRequest = createDefaultInternetRequestForTransport(
                 NetworkCapabilities.TRANSPORT_WIFI, NetworkRequest.Type.BACKGROUND_REQUEST);
 
-        mHandlerThread = new HandlerThread("ConnectivityServiceThread");
+        mHandlerThread = mDeps.makeHandlerThread();
         mHandlerThread.start();
         mHandler = new InternalHandler(mHandlerThread.getLooper());
         mTrackerHandler = new NetworkStateTrackerHandler(mHandlerThread.getLooper());
@@ -882,7 +969,7 @@
                 LocalServices.getService(NetworkPolicyManagerInternal.class),
                 "missing NetworkPolicyManagerInternal");
         mDnsResolver = checkNotNull(dnsresolver, "missing IDnsResolver");
-        mProxyTracker = makeProxyTracker();
+        mProxyTracker = mDeps.makeProxyTracker(mContext, mHandler);
 
         mNetd = netd;
         mKeyStore = KeyStore.getInstance();
@@ -950,7 +1037,7 @@
 
         // Do the same for Ethernet, since it's often not specified in the configs, although many
         // devices can use it via USB host adapters.
-        if (mNetConfigs[TYPE_ETHERNET] == null && hasService(Context.ETHERNET_SERVICE)) {
+        if (mNetConfigs[TYPE_ETHERNET] == null && mDeps.hasService(Context.ETHERNET_SERVICE)) {
             mLegacyTypeTracker.addSupportedType(TYPE_ETHERNET);
             mNetworksDefined++;
         }
@@ -968,7 +1055,10 @@
             }
         }
 
-        mTethering = makeTethering();
+        mUserManager = (UserManager) context.getSystemService(Context.USER_SERVICE);
+
+        mTethering = deps.makeTethering(mContext, mNMS, mStatsService, mPolicyManager,
+                makeTetheringDependencies());
 
         mPermissionMonitor = new PermissionMonitor(mContext, mNetd);
 
@@ -1015,8 +1105,6 @@
         final DataConnectionStats dataConnectionStats = new DataConnectionStats(mContext);
         dataConnectionStats.startMonitoring();
 
-        mUserManager = (UserManager) context.getSystemService(Context.USER_SERVICE);
-
         mKeepaliveTracker = new KeepaliveTracker(mContext, mHandler);
         mNotifier = new NetworkNotificationManager(mContext, mTelephonyManager,
                 mContext.getSystemService(NotificationManager.class));
@@ -1029,7 +1117,7 @@
                 LingerMonitor.DEFAULT_NOTIFICATION_RATE_LIMIT_MILLIS);
         mLingerMonitor = new LingerMonitor(mContext, mNotifier, dailyLimit, rateLimit);
 
-        mMultinetworkPolicyTracker = createMultinetworkPolicyTracker(
+        mMultinetworkPolicyTracker = mDeps.makeMultinetworkPolicyTracker(
                 mContext, mHandler, () -> rematchForAvoidBadWifiUpdate());
         mMultinetworkPolicyTracker.start();
 
@@ -1039,10 +1127,8 @@
         registerPrivateDnsSettingsCallbacks();
     }
 
-    @VisibleForTesting
-    protected Tethering makeTethering() {
-        // TODO: Move other elements into @Overridden getters.
-        final TetheringDependencies deps = new TetheringDependencies() {
+    private TetheringDependencies makeTetheringDependencies() {
+        return new TetheringDependencies() {
             @Override
             public boolean isTetheringSupported() {
                 return ConnectivityService.this.isTetheringSupported();
@@ -1052,14 +1138,6 @@
                 return mDefaultRequest;
             }
         };
-        return new Tethering(mContext, mNMS, mStatsService, mPolicyManager,
-                IoThread.get().getLooper(), new MockableSystemProperties(),
-                deps);
-    }
-
-    @VisibleForTesting
-    protected ProxyTracker makeProxyTracker() {
-        return new ProxyTracker(mContext, mHandler, EVENT_PROXY_HAS_CHANGED);
     }
 
     private static NetworkCapabilities createDefaultNetworkCapabilitiesForUid(int uid) {
@@ -1151,22 +1229,6 @@
         return mNextNetworkRequestId++;
     }
 
-    @VisibleForTesting
-    protected int reserveNetId() {
-        synchronized (mNetworkForNetId) {
-            for (int i = MIN_NET_ID; i <= MAX_NET_ID; i++) {
-                int netId = mNextNetId;
-                if (++mNextNetId > MAX_NET_ID) mNextNetId = MIN_NET_ID;
-                // Make sure NetID unused.  http://b/16815182
-                if (!mNetIdInUse.get(netId)) {
-                    mNetIdInUse.put(netId, true);
-                    return netId;
-                }
-            }
-        }
-        throw new IllegalStateException("No free netIds");
-    }
-
     private NetworkState getFilteredNetworkState(int networkType, int uid) {
         if (mLegacyTypeTracker.isTypeSupported(networkType)) {
             final NetworkAgentInfo nai = mLegacyTypeTracker.getNetworkForType(networkType);
@@ -1783,11 +1845,9 @@
             // the caller thread of registerNetworkAgent. Thus, it's not allowed to register netd
             // event callback for certain nai. e.g. cellular. Register here to pass to
             // NetworkMonitor instead.
-            // TODO: Move the Dns Event to NetworkMonitor. Use Binder.clearCallingIdentity() in
-            // registerNetworkAgent to have NetworkMonitor created with system process as design
-            // expectation. Also, NetdEventListenerService only allow one callback from each
-            // caller type. Need to re-factor NetdEventListenerService to allow multiple
-            // NetworkMonitor registrants.
+            // TODO: Move the Dns Event to NetworkMonitor. NetdEventListenerService only allow one
+            // callback from each caller type. Need to re-factor NetdEventListenerService to allow
+            // multiple NetworkMonitor registrants.
             if (nai != null && nai.satisfies(mDefaultRequest)) {
                 nai.networkMonitor().notifyDnsResponse(returnCode);
             }
@@ -1800,11 +1860,8 @@
         }
     };
 
-    @VisibleForTesting
-    protected void registerNetdEventCallback() {
-        final IIpConnectivityMetrics ipConnectivityMetrics =
-                IIpConnectivityMetrics.Stub.asInterface(
-                        ServiceManager.getService(IpConnectivityLog.SERVICE_NAME));
+    private void registerNetdEventCallback() {
+        final IIpConnectivityMetrics ipConnectivityMetrics = mDeps.getIpConnectivityMetrics();
         if (ipConnectivityMetrics == null) {
             Slog.wtf(TAG, "Missing IIpConnectivityMetrics");
             return;
@@ -2240,12 +2297,6 @@
     protected static final String DEFAULT_TCP_BUFFER_SIZES = "4096,87380,110208,4096,16384,110208";
     private static final String DEFAULT_TCP_RWND_KEY = "net.tcp.default_init_rwnd";
 
-    // Overridden for testing purposes to avoid writing to SystemProperties.
-    @VisibleForTesting
-    protected MockableSystemProperties getSystemProperties() {
-        return new MockableSystemProperties();
-    }
-
     private void updateTcpBufferSizes(String tcpBufferSizes) {
         String[] values = null;
         if (tcpBufferSizes != null) {
@@ -2582,8 +2633,8 @@
                     if (nai.everConnected) {
                         loge("ERROR: cannot call explicitlySelected on already-connected network");
                     }
-                    nai.networkMisc.explicitlySelected = (msg.arg1 == 1);
-                    nai.networkMisc.acceptUnvalidated = (msg.arg1 == 1) && (msg.arg2 == 1);
+                    nai.networkMisc.explicitlySelected = toBool(msg.arg1);
+                    nai.networkMisc.acceptUnvalidated = toBool(msg.arg1) && toBool(msg.arg2);
                     // Mark the network as temporarily accepting partial connectivity so that it
                     // will be validated (and possibly become default) even if it only provides
                     // partial internet access. Note that if user connects to partial connectivity
@@ -2591,7 +2642,7 @@
                     // out of wifi coverage) and if the same wifi is available again, the device
                     // will auto connect to this wifi even though the wifi has "no internet".
                     // TODO: Evaluate using a separate setting in IpMemoryStore.
-                    nai.networkMisc.acceptPartialConnectivity = (msg.arg2 == 1);
+                    nai.networkMisc.acceptPartialConnectivity = toBool(msg.arg2);
                     break;
                 }
                 case NetworkAgent.EVENT_SOCKET_KEEPALIVE: {
@@ -2635,8 +2686,9 @@
                     }
                     if (valid != nai.lastValidated) {
                         if (wasDefault) {
-                            metricsLogger().defaultNetworkMetrics().logDefaultNetworkValidity(
-                                    SystemClock.elapsedRealtime(), valid);
+                            mDeps.getMetricsLogger()
+                                    .defaultNetworkMetrics().logDefaultNetworkValidity(
+                                            SystemClock.elapsedRealtime(), valid);
                         }
                         final int oldScore = nai.getCurrentScore();
                         nai.lastValidated = valid;
@@ -2971,8 +3023,8 @@
                     final boolean wasDefault = isDefaultNetwork(nai);
                     synchronized (mNetworkForNetId) {
                         mNetworkForNetId.remove(nai.network.netId);
-                        mNetIdInUse.delete(nai.network.netId);
                     }
+                    mNetIdManager.releaseNetId(nai.network.netId);
                     // Just in case.
                     mLegacyTypeTracker.remove(nai, wasDefault);
                 }
@@ -3019,7 +3071,7 @@
             // 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);
+            mDeps.getMetricsLogger().defaultNetworkMetrics().logDefaultNetworkEvent(now, null, nai);
         }
         notifyIfacesChangedForNetworkStats();
         // TODO - we shouldn't send CALLBACK_LOST to requests that can be satisfied
@@ -3073,9 +3125,7 @@
             destroyNativeNetwork(nai);
             mDnsManager.removeNetwork(nai.network);
         }
-        synchronized (mNetworkForNetId) {
-            mNetIdInUse.delete(nai.network.netId);
-        }
+        mNetIdManager.releaseNetId(nai.network.netId);
     }
 
     private boolean createNativeNetwork(@NonNull NetworkAgentInfo networkAgent) {
@@ -4159,7 +4209,7 @@
                 return null;
             }
             return getLinkPropertiesProxyInfo(activeNetwork);
-        } else if (queryUserAccess(Binder.getCallingUid(), network.netId)) {
+        } else if (mDeps.queryUserAccess(Binder.getCallingUid(), network.netId)) {
             // Don't call getLinkProperties() as it requires ACCESS_NETWORK_STATE permission, which
             // caller may not have.
             return getLinkPropertiesProxyInfo(network);
@@ -4168,10 +4218,6 @@
         return null;
     }
 
-    @VisibleForTesting
-    protected boolean queryUserAccess(int uid, int netId) {
-        return NetworkUtils.queryUserAccess(uid, netId);
-    }
 
     private ProxyInfo getLinkPropertiesProxyInfo(Network network) {
         final NetworkAgentInfo nai = getNetworkAgentInfoForNetwork(network);
@@ -4764,7 +4810,7 @@
         final long ident = Binder.clearCallingIdentity();
         try {
             // Concatenate the range of types onto the range of NetIDs.
-            int id = MAX_NET_ID + 1 + (networkType - ConnectivityManager.TYPE_NONE);
+            int id = NetIdManager.MAX_NET_ID + 1 + (networkType - ConnectivityManager.TYPE_NONE);
             mNotifier.setProvNotificationVisible(visible, id, action);
         } finally {
             Binder.restoreCallingIdentity(ident);
@@ -5375,10 +5421,9 @@
     @GuardedBy("mNetworkForNetId")
     private final SparseArray<NetworkAgentInfo> mNetworkForNetId = new SparseArray<>();
     // NOTE: Accessed on multiple threads, synchronized with mNetworkForNetId.
-    // An entry is first added to mNetIdInUse, prior to mNetworkForNetId, so
+    // An entry is first reserved with NetIdManager, prior to being added to mNetworkForNetId, so
     // there may not be a strict 1:1 correlation between the two.
-    @GuardedBy("mNetworkForNetId")
-    private final SparseBooleanArray mNetIdInUse = new SparseBooleanArray();
+    private final NetIdManager mNetIdManager;
 
     // NetworkAgentInfo keyed off its connecting messenger
     // TODO - eval if we can reduce the number of lists/hashmaps/sparsearrays
@@ -5480,9 +5525,9 @@
         // satisfies mDefaultRequest.
         final NetworkCapabilities nc = new NetworkCapabilities(networkCapabilities);
         final NetworkAgentInfo nai = new NetworkAgentInfo(messenger, new AsyncChannel(),
-                new Network(reserveNetId()), new NetworkInfo(networkInfo), lp, nc, currentScore,
-                mContext, mTrackerHandler, new NetworkMisc(networkMisc), this, mNetd, mDnsResolver,
-                mNMS, factorySerialNumber);
+                new Network(mNetIdManager.reserveNetId()), new NetworkInfo(networkInfo), lp, nc,
+                currentScore, mContext, mTrackerHandler, new NetworkMisc(networkMisc), this, mNetd,
+                mDnsResolver, mNMS, factorySerialNumber);
         // Make sure the network capabilities reflect what the agent info says.
         nai.setNetworkCapabilities(mixInCapabilities(nai, nc));
         final String extraInfo = networkInfo.getExtraInfo();
@@ -5491,7 +5536,7 @@
         if (DBG) log("registerNetworkAgent " + nai);
         final long token = Binder.clearCallingIdentity();
         try {
-            getNetworkStack().makeNetworkMonitor(
+            mDeps.getNetworkStack().makeNetworkMonitor(
                     nai.network, name, new NetworkMonitorCallbacks(nai));
         } finally {
             Binder.restoreCallingIdentity(token);
@@ -5503,11 +5548,6 @@
         return nai.network.netId;
     }
 
-    @VisibleForTesting
-    protected NetworkStackClient getNetworkStack() {
-        return NetworkStackClient.getInstance();
-    }
-
     private void handleRegisterNetworkAgent(NetworkAgentInfo nai, INetworkMonitor networkMonitor) {
         nai.onNetworkMonitorCreated(networkMonitor);
         if (VDBG) log("Got NetworkAgent Messenger");
@@ -5523,7 +5563,6 @@
         }
         nai.asyncChannel.connect(mContext, mTrackerHandler, nai.messenger);
         NetworkInfo networkInfo = nai.networkInfo;
-        nai.networkInfo = null;
         updateNetworkInfo(nai, networkInfo);
         updateUids(nai, null, nai.networkCapabilities);
     }
@@ -6317,7 +6356,7 @@
             // Notify system services that this network is up.
             makeDefault(newNetwork);
             // Log 0 -> X and Y -> X default network transitions, where X is the new default.
-            metricsLogger().defaultNetworkMetrics().logDefaultNetworkEvent(
+            mDeps.getMetricsLogger().defaultNetworkMetrics().logDefaultNetworkEvent(
                     now, newNetwork, oldDefaultNetwork);
             // Have a new default network, release the transition wakelock in
             scheduleReleaseNetworkTransitionWakelock();
@@ -6522,8 +6561,7 @@
 
         if (DBG) {
             log(networkAgent.name() + " EVENT_NETWORK_INFO_CHANGED, going from " +
-                    (oldInfo == null ? "null" : oldInfo.getState()) +
-                    " to " + state);
+                    oldInfo.getState() + " to " + state);
         }
 
         if (!networkAgent.created
@@ -6596,8 +6634,8 @@
                 // TODO(b/122649188): send the broadcast only to VPN users.
                 mProxyTracker.sendProxyBroadcast();
             }
-        } else if ((oldInfo != null && oldInfo.getState() == NetworkInfo.State.SUSPENDED) ||
-                state == NetworkInfo.State.SUSPENDED) {
+        } else if (networkAgent.created && (oldInfo.getState() == NetworkInfo.State.SUSPENDED ||
+                state == NetworkInfo.State.SUSPENDED)) {
             // going into or coming out of SUSPEND: re-score and notify
             if (networkAgent.getCurrentScore() != oldScore) {
                 rematchAllNetworksAndRequests(networkAgent, oldScore);
@@ -6803,8 +6841,8 @@
     }
 
     /**
-     * Notify NetworkStatsService and NetworkStatsFactory that the set of active ifaces has changed,
-     * or that one of the active iface's trackedproperties has changed.
+     * Notify NetworkStatsService that the set of active ifaces has changed, or that one of the
+     * active iface's tracked properties has changed.
      */
     private void notifyIfacesChangedForNetworkStats() {
         ensureRunningOnConnectivityServiceThread();
@@ -6814,16 +6852,12 @@
             activeIface = activeLinkProperties.getInterfaceName();
         }
 
-        // CAUTION: Ordering matters between updateVpnInfos() and forceUpdateIfaces(), which
-        // triggers a new poll. Trigger the poll first to ensure a snapshot is taken before
-        // switching to the new state. This ensures that traffic does not get mis-attributed to
-        // incorrect apps (including VPN app).
+        final VpnInfo[] vpnInfos = getAllVpnInfo();
         try {
             mStatsService.forceUpdateIfaces(
-                    getDefaultNetworks(), getAllNetworkState(), activeIface);
+                    getDefaultNetworks(), getAllNetworkState(), activeIface, vpnInfos);
         } catch (Exception ignored) {
         }
-        NetworkStatsFactory.updateVpnInfos(getAllVpnInfo());
     }
 
     @Override
@@ -6999,27 +7033,6 @@
         return nwm.getWatchlistConfigHash();
     }
 
-    @VisibleForTesting
-    MultinetworkPolicyTracker createMultinetworkPolicyTracker(Context c, Handler h, Runnable r) {
-        return new MultinetworkPolicyTracker(c, h, r);
-    }
-
-    @VisibleForTesting
-    public WakeupMessage makeWakeupMessage(Context c, Handler h, String s, int cmd, Object obj) {
-        return new WakeupMessage(c, h, s, cmd, 0, 0, obj);
-    }
-
-    @VisibleForTesting
-    public boolean hasService(String name) {
-        return ServiceManager.checkService(name) != null;
-    }
-
-    @VisibleForTesting
-    protected IpConnectivityMetrics.Logger metricsLogger() {
-        return checkNotNull(LocalServices.getService(IpConnectivityMetrics.Logger.class),
-                "no IpConnectivityMetrics service");
-    }
-
     private void logNetworkEvent(NetworkAgentInfo nai, int evtype) {
         int[] transports = nai.networkCapabilities.getTransportTypes();
         mMetricsLog.log(nai.network.netId, transports, new NetworkEvent(evtype));
diff --git a/services/core/java/com/android/server/NetIdManager.java b/services/core/java/com/android/server/NetIdManager.java
new file mode 100644
index 0000000..11533be
--- /dev/null
+++ b/services/core/java/com/android/server/NetIdManager.java
@@ -0,0 +1,76 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server;
+
+import android.annotation.NonNull;
+import android.util.SparseBooleanArray;
+
+import com.android.internal.annotations.GuardedBy;
+
+/**
+ * Class used to reserve and release net IDs.
+ *
+ * <p>Instances of this class are thread-safe.
+ */
+public class NetIdManager {
+    // Sequence number for Networks; keep in sync with system/netd/NetworkController.cpp
+    public static final int MIN_NET_ID = 100; // some reserved marks
+    // Top IDs reserved by IpSecService
+    public static final int MAX_NET_ID = 65535 - IpSecService.TUN_INTF_NETID_RANGE;
+
+    @GuardedBy("mNetIdInUse")
+    private final SparseBooleanArray mNetIdInUse = new SparseBooleanArray();
+
+    @GuardedBy("mNetIdInUse")
+    private int mLastNetId = MIN_NET_ID - 1;
+
+    /**
+     * Get the first netId that follows the provided lastId and is available.
+     */
+    private static int getNextAvailableNetIdLocked(
+            int lastId, @NonNull SparseBooleanArray netIdInUse) {
+        int netId = lastId;
+        for (int i = MIN_NET_ID; i <= MAX_NET_ID; i++) {
+            netId = netId < MAX_NET_ID ? netId + 1 : MIN_NET_ID;
+            if (!netIdInUse.get(netId)) {
+                return netId;
+            }
+        }
+        throw new IllegalStateException("No free netIds");
+    }
+
+    /**
+     * Reserve a new ID for a network.
+     */
+    public int reserveNetId() {
+        synchronized (mNetIdInUse) {
+            mLastNetId = getNextAvailableNetIdLocked(mLastNetId, mNetIdInUse);
+            // Make sure NetID unused.  http://b/16815182
+            mNetIdInUse.put(mLastNetId, true);
+            return mLastNetId;
+        }
+    }
+
+    /**
+     * Clear a previously reserved ID for a network.
+     */
+    public void releaseNetId(int id) {
+        synchronized (mNetIdInUse) {
+            mNetIdInUse.delete(id);
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/TestNetworkService.java b/services/core/java/com/android/server/TestNetworkService.java
index 40bf7bc..d19d2dd 100644
--- a/services/core/java/com/android/server/TestNetworkService.java
+++ b/services/core/java/com/android/server/TestNetworkService.java
@@ -19,6 +19,7 @@
 import static com.android.internal.util.Preconditions.checkNotNull;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.content.Context;
 import android.net.ConnectivityManager;
 import android.net.INetd;
@@ -53,6 +54,7 @@
 import java.net.InterfaceAddress;
 import java.net.NetworkInterface;
 import java.net.SocketException;
+import java.util.ArrayList;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /** @hide */
@@ -226,6 +228,8 @@
             @NonNull Looper looper,
             @NonNull Context context,
             @NonNull String iface,
+            @Nullable LinkProperties lp,
+            boolean isMetered,
             int callingUid,
             @NonNull IBinder binder)
             throws RemoteException, SocketException {
@@ -245,9 +249,19 @@
         nc.addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_SUSPENDED);
         nc.addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED);
         nc.setNetworkSpecifier(new StringNetworkSpecifier(iface));
+        if (!isMetered) {
+            nc.addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_METERED);
+        }
 
         // Build LinkProperties
-        LinkProperties lp = new LinkProperties();
+        if (lp == null) {
+            lp = new LinkProperties();
+        } else {
+            lp = new LinkProperties(lp);
+            // Use LinkAddress(es) from the interface itself to minimize how much the caller
+            // is trusted.
+            lp.setLinkAddresses(new ArrayList<>());
+        }
         lp.setInterfaceName(iface);
 
         // Find the currently assigned addresses, and add them to LinkProperties
@@ -284,7 +298,11 @@
      * <p>This method provides a Network that is useful only for testing.
      */
     @Override
-    public void setupTestNetwork(@NonNull String iface, @NonNull IBinder binder) {
+    public void setupTestNetwork(
+            @NonNull String iface,
+            @Nullable LinkProperties lp,
+            boolean isMetered,
+            @NonNull IBinder binder) {
         enforceTestNetworkPermissions(mContext);
 
         checkNotNull(iface, "missing Iface");
@@ -315,6 +333,8 @@
                                             mHandler.getLooper(),
                                             mContext,
                                             iface,
+                                            lp,
+                                            isMetered,
                                             callingUid,
                                             binder);
 
diff --git a/services/core/java/com/android/server/connectivity/KeepaliveTracker.java b/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
index cb64245..9bae902 100644
--- a/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
+++ b/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
@@ -473,8 +473,6 @@
                 cleanupStoppedKeepalive(nai, ki.mSlot);
             }
         }
-        // Clean up keepalives will be done as a result of calling ki.stop() after the slots are
-        // freed.
     }
 
     public void handleStopKeepalive(NetworkAgentInfo nai, int slot, int reason) {
diff --git a/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java b/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java
index 5b04379..96b7cb3 100644
--- a/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java
+++ b/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java
@@ -16,6 +16,7 @@
 
 package com.android.server.connectivity;
 
+import android.annotation.NonNull;
 import android.content.Context;
 import android.net.IDnsResolver;
 import android.net.INetd;
@@ -116,7 +117,7 @@
 // not, ConnectivityService disconnects the NetworkAgent's AsyncChannel.
 public class NetworkAgentInfo implements Comparable<NetworkAgentInfo> {
 
-    public NetworkInfo networkInfo;
+    @NonNull public NetworkInfo networkInfo;
     // This Network object should always be used if possible, so as to encourage reuse of the
     // enclosed socket factory and connection pool.  Avoid creating other Network objects.
     // This Network object is always valid.
@@ -579,10 +580,12 @@
         }
 
         if (newExpiry > 0) {
-            mLingerMessage = mConnService.makeWakeupMessage(
+            mLingerMessage = new WakeupMessage(
                     mContext, mHandler,
-                    "NETWORK_LINGER_COMPLETE." + network.netId,
-                    EVENT_NETWORK_LINGER_COMPLETE, this);
+                    "NETWORK_LINGER_COMPLETE." + network.netId /* cmdName */,
+                    EVENT_NETWORK_LINGER_COMPLETE /* cmd */,
+                    0 /* arg1 (unused) */, 0 /* arg2 (unused) */,
+                    this /* obj (NetworkAgentInfo) */);
             mLingerMessage.schedule(newExpiry);
         }
 
diff --git a/tests/net/Android.bp b/tests/net/Android.bp
index 306cc51..502aa97 100644
--- a/tests/net/Android.bp
+++ b/tests/net/Android.bp
@@ -3,22 +3,6 @@
 //########################################################################
 java_defaults {
     name: "FrameworksNetTests-jni-defaults",
-    static_libs: [
-        "FrameworksNetCommonTests",
-        "frameworks-base-testutils",
-        "frameworks-net-testutils",
-        "framework-protos",
-        "androidx.test.rules",
-        "mockito-target-minus-junit4",
-        "platform-test-annotations",
-        "services.core",
-        "services.net",
-    ],
-    libs: [
-        "android.test.runner",
-        "android.test.base",
-        "android.test.mock",
-    ],
     jni_libs: [
         "ld-android",
         "libartbase",
@@ -44,20 +28,20 @@
         "libnativehelper",
         "libnetdbpf",
         "libnetdutils",
+        "libnetworkstatsfactorytestjni",
         "libpackagelistparser",
         "libpcre2",
         "libprocessgroup",
         "libselinux",
-        "libui",
-        "libutils",
-        "libvndksupport",
         "libtinyxml2",
+        "libui",
         "libunwindstack",
+        "libutils",
         "libutilscallstack",
+        "libvndksupport",
         "libziparchive",
         "libz",
         "netd_aidl_interface-V2-cpp",
-        "libnetworkstatsfactorytestjni",
     ],
 }
 
@@ -68,4 +52,21 @@
     platform_apis: true,
     test_suites: ["device-tests"],
     certificate: "platform",
+    static_libs: [
+        "androidx.test.rules",
+        "FrameworksNetCommonTests",
+        "frameworks-base-testutils",
+        "frameworks-net-integration-testutils",
+        "framework-protos",
+        "mockito-target-minus-junit4",
+        "net-tests-utils",
+        "platform-test-annotations",
+        "services.core",
+        "services.net",
+    ],
+    libs: [
+        "android.test.runner",
+        "android.test.base",
+        "android.test.mock",
+    ],
 }
diff --git a/tests/net/common/Android.bp b/tests/net/common/Android.bp
index db1ccb4..e44d460 100644
--- a/tests/net/common/Android.bp
+++ b/tests/net/common/Android.bp
@@ -21,12 +21,12 @@
     srcs: ["java/**/*.java", "java/**/*.kt"],
     static_libs: [
         "androidx.test.rules",
-        "frameworks-net-testutils",
         "junit",
         "mockito-target-minus-junit4",
+        "net-tests-utils",
         "platform-test-annotations",
     ],
     libs: [
         "android.test.base.stubs",
     ],
-}
\ No newline at end of file
+}
diff --git a/tests/net/common/java/android/net/IpPrefixTest.java b/tests/net/common/java/android/net/IpPrefixTest.java
index 719960d..985e10d 100644
--- a/tests/net/common/java/android/net/IpPrefixTest.java
+++ b/tests/net/common/java/android/net/IpPrefixTest.java
@@ -16,16 +16,18 @@
 
 package android.net;
 
+import static com.android.testutils.MiscAssertsKt.assertEqualBothWays;
+import static com.android.testutils.MiscAssertsKt.assertFieldCountEquals;
+import static com.android.testutils.MiscAssertsKt.assertNotEqualEitherWay;
+import static com.android.testutils.ParcelUtilsKt.assertParcelingIsLossless;
+
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
-import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
-import android.os.Parcel;
-
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
@@ -171,56 +173,46 @@
 
     }
 
-    private void assertAreEqual(Object o1, Object o2) {
-        assertTrue(o1.equals(o2));
-        assertTrue(o2.equals(o1));
-    }
-
-    private void assertAreNotEqual(Object o1, Object o2) {
-        assertFalse(o1.equals(o2));
-        assertFalse(o2.equals(o1));
-    }
-
     @Test
     public void testEquals() {
         IpPrefix p1, p2;
 
         p1 = new IpPrefix("192.0.2.251/23");
         p2 = new IpPrefix(new byte[]{(byte) 192, (byte) 0, (byte) 2, (byte) 251}, 23);
-        assertAreEqual(p1, p2);
+        assertEqualBothWays(p1, p2);
 
         p1 = new IpPrefix("192.0.2.5/23");
-        assertAreEqual(p1, p2);
+        assertEqualBothWays(p1, p2);
 
         p1 = new IpPrefix("192.0.2.5/24");
-        assertAreNotEqual(p1, p2);
+        assertNotEqualEitherWay(p1, p2);
 
         p1 = new IpPrefix("192.0.4.5/23");
-        assertAreNotEqual(p1, p2);
+        assertNotEqualEitherWay(p1, p2);
 
 
         p1 = new IpPrefix("2001:db8:dead:beef:f00::80/122");
         p2 = new IpPrefix(IPV6_BYTES, 122);
         assertEquals("2001:db8:dead:beef:f00::80/122", p2.toString());
-        assertAreEqual(p1, p2);
+        assertEqualBothWays(p1, p2);
 
         p1 = new IpPrefix("2001:db8:dead:beef:f00::bf/122");
-        assertAreEqual(p1, p2);
+        assertEqualBothWays(p1, p2);
 
         p1 = new IpPrefix("2001:db8:dead:beef:f00::8:0/123");
-        assertAreNotEqual(p1, p2);
+        assertNotEqualEitherWay(p1, p2);
 
         p1 = new IpPrefix("2001:db8:dead:beef::/122");
-        assertAreNotEqual(p1, p2);
+        assertNotEqualEitherWay(p1, p2);
 
         // 192.0.2.4/32 != c000:0204::/32.
         byte[] ipv6bytes = new byte[16];
         System.arraycopy(IPV4_BYTES, 0, ipv6bytes, 0, IPV4_BYTES.length);
         p1 = new IpPrefix(ipv6bytes, 32);
-        assertAreEqual(p1, new IpPrefix("c000:0204::/32"));
+        assertEqualBothWays(p1, new IpPrefix("c000:0204::/32"));
 
         p2 = new IpPrefix(IPV4_BYTES, 32);
-        assertAreNotEqual(p1, p2);
+        assertNotEqualEitherWay(p1, p2);
     }
 
     @Test
@@ -356,25 +348,6 @@
         assertEquals(InetAddress.parseNumericAddress("192.0.2.0"), p.getAddress());
     }
 
-    public IpPrefix passThroughParcel(IpPrefix p) {
-        Parcel parcel = Parcel.obtain();
-        IpPrefix p2 = null;
-        try {
-            p.writeToParcel(parcel, 0);
-            parcel.setDataPosition(0);
-            p2 = IpPrefix.CREATOR.createFromParcel(parcel);
-        } finally {
-            parcel.recycle();
-        }
-        assertNotNull(p2);
-        return p2;
-    }
-
-    public void assertParcelingIsLossless(IpPrefix p) {
-        IpPrefix p2 = passThroughParcel(p);
-        assertEquals(p, p2);
-    }
-
     @Test
     public void testParceling() {
         IpPrefix p;
@@ -386,5 +359,7 @@
         p = new IpPrefix("192.0.2.0/25");
         assertParcelingIsLossless(p);
         assertTrue(p.isIPv4());
+
+        assertFieldCountEquals(2, IpPrefix.class);
     }
 }
diff --git a/tests/net/common/java/android/net/LinkAddressTest.java b/tests/net/common/java/android/net/LinkAddressTest.java
index d462441..b2e573b 100644
--- a/tests/net/common/java/android/net/LinkAddressTest.java
+++ b/tests/net/common/java/android/net/LinkAddressTest.java
@@ -27,15 +27,17 @@
 import static android.system.OsConstants.RT_SCOPE_SITE;
 import static android.system.OsConstants.RT_SCOPE_UNIVERSE;
 
+import static com.android.testutils.MiscAssertsKt.assertEqualBothWays;
+import static com.android.testutils.MiscAssertsKt.assertNotEqualEitherWay;
+import static com.android.testutils.ParcelUtilsKt.assertParcelSane;
+import static com.android.testutils.ParcelUtilsKt.assertParcelingIsLossless;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
-import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
-import android.os.Parcel;
-
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
@@ -217,67 +219,56 @@
                 l1.isSameAddressAs(l2));
     }
 
-    private void assertLinkAddressesEqual(LinkAddress l1, LinkAddress l2) {
-        assertTrue(l1 + " unexpectedly not equal to " + l2, l1.equals(l2));
-        assertTrue(l2 + " unexpectedly not equal to " + l1, l2.equals(l1));
-        assertEquals(l1.hashCode(), l2.hashCode());
-    }
-
-    private void assertLinkAddressesNotEqual(LinkAddress l1, LinkAddress l2) {
-        assertFalse(l1 + " unexpectedly equal to " + l2, l1.equals(l2));
-        assertFalse(l2 + " unexpectedly equal to " + l1, l2.equals(l1));
-    }
-
     @Test
     public void testEqualsAndSameAddressAs() {
         LinkAddress l1, l2, l3;
 
         l1 = new LinkAddress("2001:db8::1/64");
         l2 = new LinkAddress("2001:db8::1/64");
-        assertLinkAddressesEqual(l1, l2);
+        assertEqualBothWays(l1, l2);
         assertIsSameAddressAs(l1, l2);
 
         l2 = new LinkAddress("2001:db8::1/65");
-        assertLinkAddressesNotEqual(l1, l2);
+        assertNotEqualEitherWay(l1, l2);
         assertIsNotSameAddressAs(l1, l2);
 
         l2 = new LinkAddress("2001:db8::2/64");
-        assertLinkAddressesNotEqual(l1, l2);
+        assertNotEqualEitherWay(l1, l2);
         assertIsNotSameAddressAs(l1, l2);
 
 
         l1 = new LinkAddress("192.0.2.1/24");
         l2 = new LinkAddress("192.0.2.1/24");
-        assertLinkAddressesEqual(l1, l2);
+        assertEqualBothWays(l1, l2);
         assertIsSameAddressAs(l1, l2);
 
         l2 = new LinkAddress("192.0.2.1/23");
-        assertLinkAddressesNotEqual(l1, l2);
+        assertNotEqualEitherWay(l1, l2);
         assertIsNotSameAddressAs(l1, l2);
 
         l2 = new LinkAddress("192.0.2.2/24");
-        assertLinkAddressesNotEqual(l1, l2);
+        assertNotEqualEitherWay(l1, l2);
         assertIsNotSameAddressAs(l1, l2);
 
 
         // Check equals() and isSameAddressAs() on identical addresses with different flags.
         l1 = new LinkAddress(V6_ADDRESS, 64);
         l2 = new LinkAddress(V6_ADDRESS, 64, 0, RT_SCOPE_UNIVERSE);
-        assertLinkAddressesEqual(l1, l2);
+        assertEqualBothWays(l1, l2);
         assertIsSameAddressAs(l1, l2);
 
         l2 = new LinkAddress(V6_ADDRESS, 64, IFA_F_DEPRECATED, RT_SCOPE_UNIVERSE);
-        assertLinkAddressesNotEqual(l1, l2);
+        assertNotEqualEitherWay(l1, l2);
         assertIsSameAddressAs(l1, l2);
 
         // Check equals() and isSameAddressAs() on identical addresses with different scope.
         l1 = new LinkAddress(V4_ADDRESS, 24);
         l2 = new LinkAddress(V4_ADDRESS, 24, 0, RT_SCOPE_UNIVERSE);
-        assertLinkAddressesEqual(l1, l2);
+        assertEqualBothWays(l1, l2);
         assertIsSameAddressAs(l1, l2);
 
         l2 = new LinkAddress(V4_ADDRESS, 24, 0, RT_SCOPE_HOST);
-        assertLinkAddressesNotEqual(l1, l2);
+        assertNotEqualEitherWay(l1, l2);
         assertIsSameAddressAs(l1, l2);
 
         // Addresses with the same start or end bytes aren't equal between families.
@@ -291,10 +282,10 @@
         assertTrue(Arrays.equals(ipv4Bytes, l2FirstIPv6Bytes));
         assertTrue(Arrays.equals(ipv4Bytes, l3LastIPv6Bytes));
 
-        assertLinkAddressesNotEqual(l1, l2);
+        assertNotEqualEitherWay(l1, l2);
         assertIsNotSameAddressAs(l1, l2);
 
-        assertLinkAddressesNotEqual(l1, l3);
+        assertNotEqualEitherWay(l1, l3);
         assertIsNotSameAddressAs(l1, l3);
 
         // Because we use InetAddress, an IPv4 address is equal to its IPv4-mapped address.
@@ -302,7 +293,7 @@
         String addressString = V4 + "/24";
         l1 = new LinkAddress(addressString);
         l2 = new LinkAddress("::ffff:" + addressString);
-        assertLinkAddressesEqual(l1, l2);
+        assertEqualBothWays(l1, l2);
         assertIsSameAddressAs(l1, l2);
     }
 
@@ -319,25 +310,6 @@
         assertNotEquals(l1.hashCode(), l2.hashCode());
     }
 
-    private LinkAddress passThroughParcel(LinkAddress l) {
-        Parcel p = Parcel.obtain();
-        LinkAddress l2 = null;
-        try {
-            l.writeToParcel(p, 0);
-            p.setDataPosition(0);
-            l2 = LinkAddress.CREATOR.createFromParcel(p);
-        } finally {
-            p.recycle();
-        }
-        assertNotNull(l2);
-        return l2;
-    }
-
-    private void assertParcelingIsLossless(LinkAddress l) {
-      LinkAddress l2 = passThroughParcel(l);
-      assertEquals(l, l2);
-    }
-
     @Test
     public void testParceling() {
         LinkAddress l;
@@ -346,7 +318,7 @@
         assertParcelingIsLossless(l);
 
         l = new LinkAddress(V4 + "/28", IFA_F_PERMANENT, RT_SCOPE_LINK);
-        assertParcelingIsLossless(l);
+        assertParcelSane(l, 4);
     }
 
     private void assertGlobalPreferred(LinkAddress l, String msg) {
diff --git a/tests/net/common/java/android/net/LinkPropertiesTest.java b/tests/net/common/java/android/net/LinkPropertiesTest.java
index e1c4238..b0464d9 100644
--- a/tests/net/common/java/android/net/LinkPropertiesTest.java
+++ b/tests/net/common/java/android/net/LinkPropertiesTest.java
@@ -16,6 +16,8 @@
 
 package android.net;
 
+import static com.android.testutils.ParcelUtilsKt.assertParcelingIsLossless;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
@@ -31,8 +33,6 @@
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
-import com.android.internal.util.TestUtils;
-
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -942,13 +942,13 @@
 
         source.setNat64Prefix(new IpPrefix("2001:db8:1:2:64:64::/96"));
 
-        TestUtils.assertParcelingIsLossless(source);
+        assertParcelingIsLossless(source);
     }
 
     @Test
     public void testParcelUninitialized() throws Exception {
         LinkProperties empty = new LinkProperties();
-        TestUtils.assertParcelingIsLossless(empty);
+        assertParcelingIsLossless(empty);
     }
 
     @Test
diff --git a/tests/net/common/java/android/net/NetworkCapabilitiesTest.java b/tests/net/common/java/android/net/NetworkCapabilitiesTest.java
index 6bc7c1b..2ca0d1a 100644
--- a/tests/net/common/java/android/net/NetworkCapabilitiesTest.java
+++ b/tests/net/common/java/android/net/NetworkCapabilitiesTest.java
@@ -38,6 +38,9 @@
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 import static android.net.NetworkCapabilities.UNRESTRICTED_CAPABILITIES;
 
+import static com.android.testutils.ParcelUtilsKt.assertParcelSane;
+import static com.android.testutils.ParcelUtilsKt.assertParcelingIsLossless;
+
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -45,7 +48,6 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
-import android.os.Parcel;
 import android.test.suitebuilder.annotation.SmallTest;
 import android.util.ArraySet;
 
@@ -267,9 +269,9 @@
             .setUids(uids)
             .addCapability(NET_CAPABILITY_EIMS)
             .addCapability(NET_CAPABILITY_NOT_METERED);
-        assertEqualsThroughMarshalling(netCap);
+        assertParcelingIsLossless(netCap);
         netCap.setSSID(TEST_SSID);
-        assertEqualsThroughMarshalling(netCap);
+        assertParcelSane(netCap, 11);
     }
 
     @Test
@@ -542,18 +544,6 @@
         nc1.combineCapabilities(nc3);
     }
 
-    private void assertEqualsThroughMarshalling(NetworkCapabilities netCap) {
-        Parcel p = Parcel.obtain();
-        netCap.writeToParcel(p, /* flags */ 0);
-        p.setDataPosition(0);
-        byte[] marshalledData = p.marshall();
-
-        p = Parcel.obtain();
-        p.unmarshall(marshalledData, 0, marshalledData.length);
-        p.setDataPosition(0);
-        assertEquals(NetworkCapabilities.CREATOR.createFromParcel(p), netCap);
-    }
-
     @Test
     public void testSet() {
         NetworkCapabilities nc1 = new NetworkCapabilities();
diff --git a/tests/net/common/java/android/net/NetworkTest.java b/tests/net/common/java/android/net/NetworkTest.java
index 38bc744..11d44b8 100644
--- a/tests/net/common/java/android/net/NetworkTest.java
+++ b/tests/net/common/java/android/net/NetworkTest.java
@@ -18,13 +18,10 @@
 
 import static org.junit.Assert.assertEquals;
 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.net.LocalServerSocket;
-import android.net.LocalSocket;
-import android.net.LocalSocketAddress;
-import android.net.Network;
 import android.platform.test.annotations.AppModeFull;
 
 import androidx.test.filters.SmallTest;
@@ -40,7 +37,6 @@
 import java.net.Inet6Address;
 import java.net.InetAddress;
 import java.net.SocketException;
-import java.util.Objects;
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
@@ -123,29 +119,29 @@
         Network three = new Network(3);
 
         // None of the hashcodes are zero.
-        assertNotEqual(0, one.hashCode());
-        assertNotEqual(0, two.hashCode());
-        assertNotEqual(0, three.hashCode());
+        assertNotEquals(0, one.hashCode());
+        assertNotEquals(0, two.hashCode());
+        assertNotEquals(0, three.hashCode());
 
         // All the hashcodes are distinct.
-        assertNotEqual(one.hashCode(), two.hashCode());
-        assertNotEqual(one.hashCode(), three.hashCode());
-        assertNotEqual(two.hashCode(), three.hashCode());
+        assertNotEquals(one.hashCode(), two.hashCode());
+        assertNotEquals(one.hashCode(), three.hashCode());
+        assertNotEquals(two.hashCode(), three.hashCode());
 
         // None of the handles are zero.
-        assertNotEqual(0, one.getNetworkHandle());
-        assertNotEqual(0, two.getNetworkHandle());
-        assertNotEqual(0, three.getNetworkHandle());
+        assertNotEquals(0, one.getNetworkHandle());
+        assertNotEquals(0, two.getNetworkHandle());
+        assertNotEquals(0, three.getNetworkHandle());
 
         // All the handles are distinct.
-        assertNotEqual(one.getNetworkHandle(), two.getNetworkHandle());
-        assertNotEqual(one.getNetworkHandle(), three.getNetworkHandle());
-        assertNotEqual(two.getNetworkHandle(), three.getNetworkHandle());
+        assertNotEquals(one.getNetworkHandle(), two.getNetworkHandle());
+        assertNotEquals(one.getNetworkHandle(), three.getNetworkHandle());
+        assertNotEquals(two.getNetworkHandle(), three.getNetworkHandle());
 
         // The handles are not equal to the hashcodes.
-        assertNotEqual(one.hashCode(), one.getNetworkHandle());
-        assertNotEqual(two.hashCode(), two.getNetworkHandle());
-        assertNotEqual(three.hashCode(), three.getNetworkHandle());
+        assertNotEquals(one.hashCode(), one.getNetworkHandle());
+        assertNotEquals(two.hashCode(), two.getNetworkHandle());
+        assertNotEquals(three.hashCode(), three.getNetworkHandle());
 
         // Adjust as necessary to test an implementation's specific constants.
         // When running with runtest, "adb logcat -s TestRunner" can be useful.
@@ -154,15 +150,11 @@
         assertEquals(16290598925L, three.getNetworkHandle());
     }
 
-    private static <T> void assertNotEqual(T t1, T t2) {
-        assertFalse(Objects.equals(t1, t2));
-    }
-
     @Test
     public void testGetPrivateDnsBypassingCopy() {
         final Network copy = mNetwork.getPrivateDnsBypassingCopy();
         assertEquals(mNetwork.netId, copy.netId);
-        assertNotEqual(copy.netId, copy.getNetIdForResolv());
-        assertNotEqual(mNetwork.getNetIdForResolv(), copy.getNetIdForResolv());
+        assertNotEquals(copy.netId, copy.getNetIdForResolv());
+        assertNotEquals(mNetwork.getNetIdForResolv(), copy.getNetIdForResolv());
     }
 }
diff --git a/tests/net/common/java/android/net/RouteInfoTest.java b/tests/net/common/java/android/net/RouteInfoTest.java
index 2edbd40..5ce8436 100644
--- a/tests/net/common/java/android/net/RouteInfoTest.java
+++ b/tests/net/common/java/android/net/RouteInfoTest.java
@@ -18,7 +18,11 @@
 
 import static android.net.RouteInfo.RTN_UNREACHABLE;
 
-import android.os.Parcel;
+import static com.android.testutils.MiscAssertsKt.assertEqualBothWays;
+import static com.android.testutils.MiscAssertsKt.assertNotEqualEitherWay;
+import static com.android.testutils.ParcelUtilsKt.assertParcelSane;
+import static com.android.testutils.ParcelUtilsKt.assertParcelingIsLossless;
+
 import android.test.suitebuilder.annotation.SmallTest;
 
 import junit.framework.TestCase;
@@ -109,47 +113,37 @@
         assertFalse(ipv4Default.matches(Address("2001:db8::f00")));
     }
 
-    private void assertAreEqual(Object o1, Object o2) {
-        assertTrue(o1.equals(o2));
-        assertTrue(o2.equals(o1));
-    }
-
-    private void assertAreNotEqual(Object o1, Object o2) {
-        assertFalse(o1.equals(o2));
-        assertFalse(o2.equals(o1));
-    }
-
     public void testEquals() {
         // IPv4
         RouteInfo r1 = new RouteInfo(Prefix("2001:db8:ace::/48"), Address("2001:db8::1"), "wlan0");
         RouteInfo r2 = new RouteInfo(Prefix("2001:db8:ace::/48"), Address("2001:db8::1"), "wlan0");
-        assertAreEqual(r1, r2);
+        assertEqualBothWays(r1, r2);
 
         RouteInfo r3 = new RouteInfo(Prefix("2001:db8:ace::/49"), Address("2001:db8::1"), "wlan0");
         RouteInfo r4 = new RouteInfo(Prefix("2001:db8:ace::/48"), Address("2001:db8::2"), "wlan0");
         RouteInfo r5 = new RouteInfo(Prefix("2001:db8:ace::/48"), Address("2001:db8::1"), "rmnet0");
-        assertAreNotEqual(r1, r3);
-        assertAreNotEqual(r1, r4);
-        assertAreNotEqual(r1, r5);
+        assertNotEqualEitherWay(r1, r3);
+        assertNotEqualEitherWay(r1, r4);
+        assertNotEqualEitherWay(r1, r5);
 
         // IPv6
         r1 = new RouteInfo(Prefix("192.0.2.0/25"), Address("192.0.2.1"), "wlan0");
         r2 = new RouteInfo(Prefix("192.0.2.0/25"), Address("192.0.2.1"), "wlan0");
-        assertAreEqual(r1, r2);
+        assertEqualBothWays(r1, r2);
 
         r3 = new RouteInfo(Prefix("192.0.2.0/24"), Address("192.0.2.1"), "wlan0");
         r4 = new RouteInfo(Prefix("192.0.2.0/25"), Address("192.0.2.2"), "wlan0");
         r5 = new RouteInfo(Prefix("192.0.2.0/25"), Address("192.0.2.1"), "rmnet0");
-        assertAreNotEqual(r1, r3);
-        assertAreNotEqual(r1, r4);
-        assertAreNotEqual(r1, r5);
+        assertNotEqualEitherWay(r1, r3);
+        assertNotEqualEitherWay(r1, r4);
+        assertNotEqualEitherWay(r1, r5);
 
         // Interfaces (but not destinations or gateways) can be null.
         r1 = new RouteInfo(Prefix("192.0.2.0/25"), Address("192.0.2.1"), null);
         r2 = new RouteInfo(Prefix("192.0.2.0/25"), Address("192.0.2.1"), null);
         r3 = new RouteInfo(Prefix("192.0.2.0/24"), Address("192.0.2.1"), "wlan0");
-        assertAreEqual(r1, r2);
-        assertAreNotEqual(r1, r3);
+        assertEqualBothWays(r1, r2);
+        assertNotEqualEitherWay(r1, r3);
     }
 
     public void testHostAndDefaultRoutes() {
@@ -257,25 +251,6 @@
       // No exceptions? Good.
     }
 
-    public RouteInfo passThroughParcel(RouteInfo r) {
-        Parcel p = Parcel.obtain();
-        RouteInfo r2 = null;
-        try {
-            r.writeToParcel(p, 0);
-            p.setDataPosition(0);
-            r2 = RouteInfo.CREATOR.createFromParcel(p);
-        } finally {
-            p.recycle();
-        }
-        assertNotNull(r2);
-        return r2;
-    }
-
-    public void assertParcelingIsLossless(RouteInfo r) {
-      RouteInfo r2 = passThroughParcel(r);
-      assertEquals(r, r2);
-    }
-
     public void testParceling() {
         RouteInfo r;
 
@@ -283,6 +258,6 @@
         assertParcelingIsLossless(r);
 
         r = new RouteInfo(Prefix("192.0.2.0/24"), null, "wlan0");
-        assertParcelingIsLossless(r);
+        assertParcelSane(r, 6);
     }
 }
diff --git a/tests/net/common/java/android/net/StaticIpConfigurationTest.java b/tests/net/common/java/android/net/StaticIpConfigurationTest.java
index 5096be2..b5f23bf 100644
--- a/tests/net/common/java/android/net/StaticIpConfigurationTest.java
+++ b/tests/net/common/java/android/net/StaticIpConfigurationTest.java
@@ -18,6 +18,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
@@ -34,7 +35,6 @@
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Objects;
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
@@ -61,10 +61,6 @@
         assertEquals(0, s.dnsServers.size());
     }
 
-    private static <T> void assertNotEquals(T t1, T t2) {
-        assertFalse(Objects.equals(t1, t2));
-    }
-
     private StaticIpConfiguration makeTestObject() {
         StaticIpConfiguration s = new StaticIpConfiguration();
         s.ipAddress = ADDR;
diff --git a/tests/net/common/java/android/net/apf/ApfCapabilitiesTest.java b/tests/net/common/java/android/net/apf/ApfCapabilitiesTest.java
index 0ce7c91..f4f804a 100644
--- a/tests/net/common/java/android/net/apf/ApfCapabilitiesTest.java
+++ b/tests/net/common/java/android/net/apf/ApfCapabilitiesTest.java
@@ -16,6 +16,8 @@
 
 package android.net.apf;
 
+import static com.android.testutils.ParcelUtilsKt.assertParcelSane;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
@@ -24,9 +26,6 @@
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
-import com.android.internal.util.ParcelableTestUtil;
-import com.android.internal.util.TestUtils;
-
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -40,9 +39,7 @@
         assertEquals(456, caps.maximumApfProgramSize);
         assertEquals(789, caps.apfPacketFormat);
 
-        ParcelableTestUtil.assertFieldCountEquals(3, ApfCapabilities.class);
-
-        TestUtils.assertParcelingIsLossless(caps);
+        assertParcelSane(caps, 3);
     }
 
     @Test
diff --git a/tests/net/common/java/android/net/metrics/ApfProgramEventTest.kt b/tests/net/common/java/android/net/metrics/ApfProgramEventTest.kt
index 8d055c9..0b7b740 100644
--- a/tests/net/common/java/android/net/metrics/ApfProgramEventTest.kt
+++ b/tests/net/common/java/android/net/metrics/ApfProgramEventTest.kt
@@ -16,11 +16,9 @@
 
 package android.net.metrics;
 
-import android.os.Parcelable
 import androidx.test.filters.SmallTest
 import androidx.test.runner.AndroidJUnit4
-import com.android.internal.util.ParcelableTestUtil
-import com.android.internal.util.TestUtils
+import com.android.testutils.assertParcelSane
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertFalse
 import org.junit.Assert.assertTrue
@@ -30,11 +28,6 @@
 @RunWith(AndroidJUnit4::class)
 @SmallTest
 class ApfProgramEventTest {
-    private fun <T: Parcelable> testParcel(obj: T, fieldCount: Int) {
-        ParcelableTestUtil.assertFieldCountEquals(fieldCount, obj::class.java)
-        TestUtils.assertParcelingIsLossless(obj)
-    }
-
     private infix fun Int.hasFlag(flag: Int) = (this and (1 shl flag)) != 0
 
     @Test
@@ -55,7 +48,7 @@
         assertEquals(5, apfProgramEvent.programLength)
         assertEquals(ApfProgramEvent.flagsFor(true, true), apfProgramEvent.flags)
 
-        testParcel(apfProgramEvent, 6)
+        assertParcelSane(apfProgramEvent, 6)
     }
 
     @Test
diff --git a/tests/net/common/java/android/net/metrics/ApfStatsTest.kt b/tests/net/common/java/android/net/metrics/ApfStatsTest.kt
index f8eb40c..46a8c8e 100644
--- a/tests/net/common/java/android/net/metrics/ApfStatsTest.kt
+++ b/tests/net/common/java/android/net/metrics/ApfStatsTest.kt
@@ -16,11 +16,9 @@
 
 package android.net.metrics
 
-import android.os.Parcelable
 import androidx.test.filters.SmallTest
 import androidx.test.runner.AndroidJUnit4
-import com.android.internal.util.ParcelableTestUtil
-import com.android.internal.util.TestUtils
+import com.android.testutils.assertParcelSane
 import org.junit.Assert.assertEquals
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -28,11 +26,6 @@
 @RunWith(AndroidJUnit4::class)
 @SmallTest
 class ApfStatsTest {
-    private fun <T: Parcelable> testParcel(obj: T, fieldCount: Int) {
-        ParcelableTestUtil.assertFieldCountEquals(fieldCount, obj::class.java)
-        TestUtils.assertParcelingIsLossless(obj)
-    }
-
     @Test
     fun testBuilderAndParcel() {
         val apfStats = ApfStats.Builder()
@@ -59,6 +52,6 @@
         assertEquals(8, apfStats.programUpdatesAllowingMulticast)
         assertEquals(9, apfStats.maxProgramSize)
 
-        testParcel(apfStats, 10)
+        assertParcelSane(apfStats, 10)
     }
 }
diff --git a/tests/net/common/java/android/net/metrics/DhcpErrorEventTest.kt b/tests/net/common/java/android/net/metrics/DhcpErrorEventTest.kt
index e9d5e6d..236f72e 100644
--- a/tests/net/common/java/android/net/metrics/DhcpErrorEventTest.kt
+++ b/tests/net/common/java/android/net/metrics/DhcpErrorEventTest.kt
@@ -1,10 +1,10 @@
 package android.net.metrics
 
-import android.net.metrics.DhcpErrorEvent.errorCodeWithOption
 import android.net.metrics.DhcpErrorEvent.DHCP_INVALID_OPTION_LENGTH
+import android.net.metrics.DhcpErrorEvent.errorCodeWithOption
 import androidx.test.filters.SmallTest
 import androidx.test.runner.AndroidJUnit4
-import com.android.internal.util.TestUtils.parcelingRoundTrip
+import com.android.testutils.parcelingRoundTrip
 import java.lang.reflect.Modifier
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertNotNull
@@ -62,4 +62,4 @@
     fun testToString_InvalidErrorCode() {
         assertNotNull(DhcpErrorEvent(TEST_ERROR_CODE).toString())
     }
-}
\ No newline at end of file
+}
diff --git a/tests/net/common/java/android/net/metrics/IpReachabilityEventTest.kt b/tests/net/common/java/android/net/metrics/IpReachabilityEventTest.kt
index d76ebf6..55b5e49 100644
--- a/tests/net/common/java/android/net/metrics/IpReachabilityEventTest.kt
+++ b/tests/net/common/java/android/net/metrics/IpReachabilityEventTest.kt
@@ -16,11 +16,9 @@
 
 package android.net.metrics
 
-import android.os.Parcelable
 import androidx.test.filters.SmallTest
 import androidx.test.runner.AndroidJUnit4
-import com.android.internal.util.ParcelableTestUtil
-import com.android.internal.util.TestUtils
+import com.android.testutils.assertParcelSane
 import org.junit.Assert.assertEquals
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -28,18 +26,13 @@
 @RunWith(AndroidJUnit4::class)
 @SmallTest
 class IpReachabilityEventTest {
-    private fun <T: Parcelable> testParcel(obj: T, fieldCount: Int) {
-        ParcelableTestUtil.assertFieldCountEquals(fieldCount, obj::class.java)
-        TestUtils.assertParcelingIsLossless(obj)
-    }
-
     @Test
     fun testConstructorAndParcel() {
         (IpReachabilityEvent.PROBE..IpReachabilityEvent.PROVISIONING_LOST_ORGANIC).forEach {
             val ipReachabilityEvent = IpReachabilityEvent(it)
             assertEquals(it, ipReachabilityEvent.eventType)
 
-            testParcel(ipReachabilityEvent, 1)
+            assertParcelSane(ipReachabilityEvent, 1)
         }
     }
 }
diff --git a/tests/net/common/java/android/net/metrics/RaEventTest.kt b/tests/net/common/java/android/net/metrics/RaEventTest.kt
index f38d328..d9b7203 100644
--- a/tests/net/common/java/android/net/metrics/RaEventTest.kt
+++ b/tests/net/common/java/android/net/metrics/RaEventTest.kt
@@ -16,11 +16,9 @@
 
 package android.net.metrics
 
-import android.os.Parcelable
 import androidx.test.filters.SmallTest
 import androidx.test.runner.AndroidJUnit4
-import com.android.internal.util.ParcelableTestUtil
-import com.android.internal.util.TestUtils
+import com.android.testutils.assertParcelSane
 import org.junit.Assert.assertEquals
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -30,11 +28,6 @@
 @RunWith(AndroidJUnit4::class)
 @SmallTest
 class RaEventTest {
-    private fun <T: Parcelable> testParcel(obj: T, fieldCount: Int) {
-        ParcelableTestUtil.assertFieldCountEquals(fieldCount, obj::class.java)
-        TestUtils.assertParcelingIsLossless(obj)
-    }
-
     @Test
     fun testConstructorAndParcel() {
         var raEvent = RaEvent.Builder().build()
@@ -74,6 +67,6 @@
         assertEquals(5, raEvent.rdnssLifetime)
         assertEquals(6, raEvent.dnsslLifetime)
 
-        testParcel(raEvent, 6)
+        assertParcelSane(raEvent, 6)
     }
 }
diff --git a/tests/net/common/java/android/net/metrics/ValidationProbeEventTest.kt b/tests/net/common/java/android/net/metrics/ValidationProbeEventTest.kt
index c0cef8f..51c0d41 100644
--- a/tests/net/common/java/android/net/metrics/ValidationProbeEventTest.kt
+++ b/tests/net/common/java/android/net/metrics/ValidationProbeEventTest.kt
@@ -16,11 +16,9 @@
 
 package android.net.metrics
 
-import android.os.Parcelable
 import androidx.test.filters.SmallTest
 import androidx.test.runner.AndroidJUnit4
-import com.android.internal.util.ParcelableTestUtil
-import com.android.internal.util.TestUtils
+import com.android.testutils.assertParcelSane
 import java.lang.reflect.Modifier
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertTrue
@@ -33,11 +31,6 @@
 @RunWith(AndroidJUnit4::class)
 @SmallTest
 class ValidationProbeEventTest {
-    private fun <T: Parcelable> testParcel(obj: T, fieldCount: Int) {
-        ParcelableTestUtil.assertFieldCountEquals(fieldCount, obj::class.java)
-        TestUtils.assertParcelingIsLossless(obj)
-    }
-
     private infix fun Int.hasType(type: Int) = (type and this) == type
 
     @Test
@@ -58,7 +51,7 @@
         assertTrue(validationProbeEvent.probeType hasType FIRST_VALIDATION)
         assertEquals(ValidationProbeEvent.DNS_SUCCESS, validationProbeEvent.returnCode)
 
-        testParcel(validationProbeEvent, 3)
+        assertParcelSane(validationProbeEvent, 3)
     }
 
     @Test
diff --git a/tests/net/common/java/android/net/util/SocketUtilsTest.kt b/tests/net/common/java/android/net/util/SocketUtilsTest.kt
new file mode 100644
index 0000000..9c7cfb0
--- /dev/null
+++ b/tests/net/common/java/android/net/util/SocketUtilsTest.kt
@@ -0,0 +1,76 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.net.util;
+
+import android.system.NetlinkSocketAddress
+import android.system.Os
+import android.system.OsConstants.AF_INET
+import android.system.OsConstants.ETH_P_ALL
+import android.system.OsConstants.IPPROTO_UDP
+import android.system.OsConstants.RTMGRP_NEIGH
+import android.system.OsConstants.SOCK_DGRAM
+import android.system.PacketSocketAddress
+import androidx.test.filters.SmallTest
+import androidx.test.runner.AndroidJUnit4
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertTrue
+import org.junit.Assert.fail
+import org.junit.Test
+import org.junit.runner.RunWith
+
+private const val TEST_INDEX = 123
+private const val TEST_PORT = 555
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+class SocketUtilsTest {
+    @Test
+    fun testMakeNetlinkSocketAddress() {
+        val nlAddress = SocketUtils.makeNetlinkSocketAddress(TEST_PORT, RTMGRP_NEIGH)
+        if (nlAddress is NetlinkSocketAddress) {
+            assertEquals(TEST_PORT, nlAddress.getPortId())
+            assertEquals(RTMGRP_NEIGH, nlAddress.getGroupsMask())
+        } else {
+            fail("Not NetlinkSocketAddress object")
+        }
+    }
+
+    @Test
+    fun testMakePacketSocketAddress() {
+        val pkAddress = SocketUtils.makePacketSocketAddress(ETH_P_ALL, TEST_INDEX)
+        assertTrue("Not PacketSocketAddress object", pkAddress is PacketSocketAddress)
+
+        val ff = 0xff.toByte()
+        val pkAddress2 = SocketUtils.makePacketSocketAddress(TEST_INDEX,
+                byteArrayOf(ff, ff, ff, ff, ff, ff))
+        assertTrue("Not PacketSocketAddress object", pkAddress2 is PacketSocketAddress)
+    }
+
+    @Test
+    fun testCloseSocket() {
+        // Expect no exception happening with null object.
+        SocketUtils.closeSocket(null)
+
+        val fd = Os.socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)
+        assertTrue(fd.valid())
+        SocketUtils.closeSocket(fd)
+        assertFalse(fd.valid())
+        // Expecting socket should be still invalid even closed socket again.
+        SocketUtils.closeSocket(fd)
+        assertFalse(fd.valid())
+    }
+}
diff --git a/tests/net/integration/util/com/android/server/ConnectivityServiceTestUtils.kt b/tests/net/integration/util/com/android/server/ConnectivityServiceTestUtils.kt
new file mode 100644
index 0000000..fa2b99c
--- /dev/null
+++ b/tests/net/integration/util/com/android/server/ConnectivityServiceTestUtils.kt
@@ -0,0 +1,41 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License
+ */
+
+package com.android.server
+
+import android.net.ConnectivityManager.TYPE_BLUETOOTH
+import android.net.ConnectivityManager.TYPE_ETHERNET
+import android.net.ConnectivityManager.TYPE_MOBILE
+import android.net.ConnectivityManager.TYPE_NONE
+import android.net.ConnectivityManager.TYPE_TEST
+import android.net.ConnectivityManager.TYPE_VPN
+import android.net.ConnectivityManager.TYPE_WIFI
+import android.net.NetworkCapabilities.TRANSPORT_BLUETOOTH
+import android.net.NetworkCapabilities.TRANSPORT_CELLULAR
+import android.net.NetworkCapabilities.TRANSPORT_ETHERNET
+import android.net.NetworkCapabilities.TRANSPORT_TEST
+import android.net.NetworkCapabilities.TRANSPORT_VPN
+import android.net.NetworkCapabilities.TRANSPORT_WIFI
+
+fun transportToLegacyType(transport: Int) = when (transport) {
+    TRANSPORT_BLUETOOTH -> TYPE_BLUETOOTH
+    TRANSPORT_CELLULAR -> TYPE_MOBILE
+    TRANSPORT_ETHERNET -> TYPE_ETHERNET
+    TRANSPORT_TEST -> TYPE_TEST
+    TRANSPORT_VPN -> TYPE_VPN
+    TRANSPORT_WIFI -> TYPE_WIFI
+    else -> TYPE_NONE
+}
\ No newline at end of file
diff --git a/tests/net/integration/util/com/android/server/NetworkAgentWrapper.java b/tests/net/integration/util/com/android/server/NetworkAgentWrapper.java
new file mode 100644
index 0000000..1e8d83c
--- /dev/null
+++ b/tests/net/integration/util/com/android/server/NetworkAgentWrapper.java
@@ -0,0 +1,267 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server;
+
+import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_VPN;
+import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
+import static android.net.NetworkCapabilities.TRANSPORT_ETHERNET;
+import static android.net.NetworkCapabilities.TRANSPORT_VPN;
+import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
+import static android.net.NetworkCapabilities.TRANSPORT_WIFI_AWARE;
+
+import static com.android.server.ConnectivityServiceTestUtilsKt.transportToLegacyType;
+
+import static junit.framework.Assert.assertTrue;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+
+import android.content.Context;
+import android.net.ConnectivityManager;
+import android.net.LinkProperties;
+import android.net.Network;
+import android.net.NetworkAgent;
+import android.net.NetworkCapabilities;
+import android.net.NetworkFactory;
+import android.net.NetworkInfo;
+import android.net.NetworkMisc;
+import android.net.NetworkSpecifier;
+import android.net.SocketKeepalive;
+import android.net.UidRange;
+import android.os.ConditionVariable;
+import android.os.HandlerThread;
+import android.os.Message;
+import android.util.Log;
+
+import com.android.server.connectivity.ConnectivityConstants;
+import com.android.testutils.HandlerUtilsKt;
+import com.android.testutils.TestableNetworkCallback;
+
+import java.util.Set;
+
+public class NetworkAgentWrapper implements TestableNetworkCallback.HasNetwork {
+    private final NetworkInfo mNetworkInfo;
+    private final NetworkCapabilities mNetworkCapabilities;
+    private final HandlerThread mHandlerThread;
+    private final Context mContext;
+    private final String mLogTag;
+
+    private final ConditionVariable mDisconnected = new ConditionVariable();
+    private final ConditionVariable mPreventReconnectReceived = new ConditionVariable();
+    private int mScore;
+    private NetworkAgent mNetworkAgent;
+    private int mStartKeepaliveError = SocketKeepalive.ERROR_UNSUPPORTED;
+    private int mStopKeepaliveError = SocketKeepalive.NO_KEEPALIVE;
+    private Integer mExpectedKeepaliveSlot = null;
+
+    public NetworkAgentWrapper(int transport, LinkProperties linkProperties, Context context)
+            throws Exception {
+        final int type = transportToLegacyType(transport);
+        final String typeName = ConnectivityManager.getNetworkTypeName(type);
+        mNetworkInfo = new NetworkInfo(type, 0, typeName, "Mock");
+        mNetworkCapabilities = new NetworkCapabilities();
+        mNetworkCapabilities.addTransportType(transport);
+        switch (transport) {
+            case TRANSPORT_ETHERNET:
+                mScore = 70;
+                break;
+            case TRANSPORT_WIFI:
+                mScore = 60;
+                break;
+            case TRANSPORT_CELLULAR:
+                mScore = 50;
+                break;
+            case TRANSPORT_WIFI_AWARE:
+                mScore = 20;
+                break;
+            case TRANSPORT_VPN:
+                mNetworkCapabilities.removeCapability(NET_CAPABILITY_NOT_VPN);
+                mScore = ConnectivityConstants.VPN_DEFAULT_SCORE;
+                break;
+            default:
+                throw new UnsupportedOperationException("unimplemented network type");
+        }
+        mContext = context;
+        mLogTag = "Mock-" + typeName;
+        mHandlerThread = new HandlerThread(mLogTag);
+        mHandlerThread.start();
+
+        mNetworkAgent = makeNetworkAgent(linkProperties);
+    }
+
+    protected InstrumentedNetworkAgent makeNetworkAgent(LinkProperties linkProperties)
+            throws Exception {
+        return new InstrumentedNetworkAgent(this, linkProperties);
+    }
+
+    public static class InstrumentedNetworkAgent extends NetworkAgent {
+        private final NetworkAgentWrapper mWrapper;
+
+        public InstrumentedNetworkAgent(NetworkAgentWrapper wrapper, LinkProperties lp) {
+            super(wrapper.mHandlerThread.getLooper(), wrapper.mContext, wrapper.mLogTag,
+                    wrapper.mNetworkInfo, wrapper.mNetworkCapabilities, lp, wrapper.mScore,
+                    new NetworkMisc(), NetworkFactory.SerialNumber.NONE);
+            mWrapper = wrapper;
+        }
+
+        @Override
+        public void unwanted() {
+            mWrapper.mDisconnected.open();
+        }
+
+        @Override
+        public void startSocketKeepalive(Message msg) {
+            int slot = msg.arg1;
+            if (mWrapper.mExpectedKeepaliveSlot != null) {
+                assertEquals((int) mWrapper.mExpectedKeepaliveSlot, slot);
+            }
+            onSocketKeepaliveEvent(slot, mWrapper.mStartKeepaliveError);
+        }
+
+        @Override
+        public void stopSocketKeepalive(Message msg) {
+            onSocketKeepaliveEvent(msg.arg1, mWrapper.mStopKeepaliveError);
+        }
+
+        @Override
+        protected void preventAutomaticReconnect() {
+            mWrapper.mPreventReconnectReceived.open();
+        }
+
+        @Override
+        protected void addKeepalivePacketFilter(Message msg) {
+            Log.i(mWrapper.mLogTag, "Add keepalive packet filter.");
+        }
+
+        @Override
+        protected void removeKeepalivePacketFilter(Message msg) {
+            Log.i(mWrapper.mLogTag, "Remove keepalive packet filter.");
+        }
+    }
+
+    public void adjustScore(int change) {
+        mScore += change;
+        mNetworkAgent.sendNetworkScore(mScore);
+    }
+
+    public int getScore() {
+        return mScore;
+    }
+
+    public void explicitlySelected(boolean explicitlySelected, boolean acceptUnvalidated) {
+        mNetworkAgent.explicitlySelected(explicitlySelected, acceptUnvalidated);
+    }
+
+    public void addCapability(int capability) {
+        mNetworkCapabilities.addCapability(capability);
+        mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
+    }
+
+    public void removeCapability(int capability) {
+        mNetworkCapabilities.removeCapability(capability);
+        mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
+    }
+
+    public void setUids(Set<UidRange> uids) {
+        mNetworkCapabilities.setUids(uids);
+        mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
+    }
+
+    public void setSignalStrength(int signalStrength) {
+        mNetworkCapabilities.setSignalStrength(signalStrength);
+        mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
+    }
+
+    public void setNetworkSpecifier(NetworkSpecifier networkSpecifier) {
+        mNetworkCapabilities.setNetworkSpecifier(networkSpecifier);
+        mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
+    }
+
+    public void setNetworkCapabilities(NetworkCapabilities nc, boolean sendToConnectivityService) {
+        mNetworkCapabilities.set(nc);
+        if (sendToConnectivityService) {
+            mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
+        }
+    }
+
+    public void connect() {
+        assertNotEquals("MockNetworkAgents can only be connected once",
+                getNetworkInfo().getDetailedState(), NetworkInfo.DetailedState.CONNECTED);
+        mNetworkInfo.setDetailedState(NetworkInfo.DetailedState.CONNECTED, null, null);
+        mNetworkAgent.sendNetworkInfo(mNetworkInfo);
+    }
+
+    public void suspend() {
+        mNetworkInfo.setDetailedState(NetworkInfo.DetailedState.SUSPENDED, null, null);
+        mNetworkAgent.sendNetworkInfo(mNetworkInfo);
+    }
+
+    public void resume() {
+        mNetworkInfo.setDetailedState(NetworkInfo.DetailedState.CONNECTED, null, null);
+        mNetworkAgent.sendNetworkInfo(mNetworkInfo);
+    }
+
+    public void disconnect() {
+        mNetworkInfo.setDetailedState(NetworkInfo.DetailedState.DISCONNECTED, null, null);
+        mNetworkAgent.sendNetworkInfo(mNetworkInfo);
+    }
+
+    @Override
+    public Network getNetwork() {
+        return new Network(mNetworkAgent.netId);
+    }
+
+    public void expectPreventReconnectReceived(long timeoutMs) {
+        assertTrue(mPreventReconnectReceived.block(timeoutMs));
+    }
+
+    public void expectDisconnected(long timeoutMs) {
+        assertTrue(mDisconnected.block(timeoutMs));
+    }
+
+    public void sendLinkProperties(LinkProperties lp) {
+        mNetworkAgent.sendLinkProperties(lp);
+    }
+
+    public void setStartKeepaliveEvent(int reason) {
+        mStartKeepaliveError = reason;
+    }
+
+    public void setStopKeepaliveEvent(int reason) {
+        mStopKeepaliveError = reason;
+    }
+
+    public void setExpectedKeepaliveSlot(Integer slot) {
+        mExpectedKeepaliveSlot = slot;
+    }
+
+    public NetworkAgent getNetworkAgent() {
+        return mNetworkAgent;
+    }
+
+    public NetworkInfo getNetworkInfo() {
+        return mNetworkInfo;
+    }
+
+    public NetworkCapabilities getNetworkCapabilities() {
+        return mNetworkCapabilities;
+    }
+
+    public void waitForIdle(long timeoutMs) {
+        HandlerUtilsKt.waitForIdle(mHandlerThread, timeoutMs);
+    }
+}
diff --git a/tests/net/integration/util/com/android/server/TestNetIdManager.kt b/tests/net/integration/util/com/android/server/TestNetIdManager.kt
new file mode 100644
index 0000000..eb290dc
--- /dev/null
+++ b/tests/net/integration/util/com/android/server/TestNetIdManager.kt
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License
+ */
+
+package com.android.server
+
+import java.util.concurrent.atomic.AtomicInteger
+
+/**
+ * A [NetIdManager] that generates ID starting from [NetIdManager.MAX_NET_ID] and decreasing, rather
+ * than starting from [NetIdManager.MIN_NET_ID] and increasing.
+ *
+ * Useful for testing ConnectivityService, to minimize the risk of test ConnectivityService netIDs
+ * overlapping with netIDs used by the real ConnectivityService on the device.
+ *
+ * IDs may still overlap if many networks have been used on the device (so the "real" netIDs
+ * are close to MAX_NET_ID), but this is typically not the case when running unit tests. Also, there
+ * is no way to fully solve the overlap issue without altering ID allocation in non-test code, as
+ * the real ConnectivityService could start using netIds that have been used by the test in the
+ * past.
+ */
+class TestNetIdManager : NetIdManager() {
+    private val nextId = AtomicInteger(MAX_NET_ID)
+    override fun reserveNetId() = nextId.decrementAndGet()
+    override fun releaseNetId(id: Int) = Unit
+}
diff --git a/tests/net/java/android/app/usage/NetworkStatsManagerTest.java b/tests/net/java/android/app/usage/NetworkStatsManagerTest.java
index fd555c1..899295a 100644
--- a/tests/net/java/android/app/usage/NetworkStatsManagerTest.java
+++ b/tests/net/java/android/app/usage/NetworkStatsManagerTest.java
@@ -202,8 +202,7 @@
         assertFalse(stats.hasNextBucket());
     }
 
-    private void assertBucketMatches(Entry expected,
-            NetworkStats.Bucket actual) {
+    private void assertBucketMatches(Entry expected, NetworkStats.Bucket actual) {
         assertEquals(expected.uid, actual.getUid());
         assertEquals(expected.rxBytes, actual.getRxBytes());
         assertEquals(expected.rxPackets, actual.getRxPackets());
diff --git a/tests/net/java/android/net/IpSecConfigTest.java b/tests/net/java/android/net/IpSecConfigTest.java
index 215506c..c9888b2 100644
--- a/tests/net/java/android/net/IpSecConfigTest.java
+++ b/tests/net/java/android/net/IpSecConfigTest.java
@@ -16,12 +16,12 @@
 
 package android.net;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertTrue;
+import static com.android.testutils.ParcelUtilsKt.assertParcelSane;
+import static com.android.testutils.ParcelUtilsKt.assertParcelingIsLossless;
 
-import android.os.Parcel;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertNull;
 
 import androidx.test.filters.SmallTest;
 
@@ -89,23 +89,15 @@
         IpSecConfig original = getSampleConfig();
         IpSecConfig copy = new IpSecConfig(original);
 
-        assertTrue(IpSecConfig.equals(original, copy));
-        assertFalse(original == copy);
+        assertEquals(original, copy);
+        assertNotSame(original, copy);
     }
 
     @Test
-    public void testParcelUnparcel() throws Exception {
+    public void testParcelUnparcel() {
         assertParcelingIsLossless(new IpSecConfig());
 
         IpSecConfig c = getSampleConfig();
-        assertParcelingIsLossless(c);
-    }
-
-    private void assertParcelingIsLossless(IpSecConfig ci) throws Exception {
-        Parcel p = Parcel.obtain();
-        ci.writeToParcel(p, 0);
-        p.setDataPosition(0);
-        IpSecConfig co = IpSecConfig.CREATOR.createFromParcel(p);
-        assertTrue(IpSecConfig.equals(co, ci));
+        assertParcelSane(c, 15);
     }
 }
diff --git a/tests/net/java/android/net/IpSecTransformTest.java b/tests/net/java/android/net/IpSecTransformTest.java
index 2308a3c..424f23d 100644
--- a/tests/net/java/android/net/IpSecTransformTest.java
+++ b/tests/net/java/android/net/IpSecTransformTest.java
@@ -16,8 +16,8 @@
 
 package android.net;
 
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
 
 import androidx.test.filters.SmallTest;
 
@@ -43,7 +43,7 @@
         config.setSpiResourceId(1985);
         IpSecTransform postModification = new IpSecTransform(null, config);
 
-        assertFalse(IpSecTransform.equals(preModification, postModification));
+        assertNotEquals(preModification, postModification);
     }
 
     @Test
@@ -57,6 +57,6 @@
         IpSecTransform config1 = new IpSecTransform(null, config);
         IpSecTransform config2 = new IpSecTransform(null, config);
 
-        assertTrue(IpSecTransform.equals(config1, config2));
+        assertEquals(config1, config2);
     }
 }
diff --git a/tests/net/java/android/net/TcpKeepalivePacketDataTest.java b/tests/net/java/android/net/TcpKeepalivePacketDataTest.java
index 583d3fd..5cb0d7e 100644
--- a/tests/net/java/android/net/TcpKeepalivePacketDataTest.java
+++ b/tests/net/java/android/net/TcpKeepalivePacketDataTest.java
@@ -16,14 +16,14 @@
 
 package android.net;
 
+import static com.android.testutils.ParcelUtilsKt.assertParcelingIsLossless;
+
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
 
 import android.net.SocketKeepalive.InvalidPacketException;
 
-import com.android.internal.util.TestUtils;
-
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -79,7 +79,7 @@
         assertEquals(testInfo.tos, resultData.ipTos);
         assertEquals(testInfo.ttl, resultData.ipTtl);
 
-        TestUtils.assertParcelingIsLossless(resultData);
+        assertParcelingIsLossless(resultData);
 
         final byte[] packet = resultData.getPacket();
         // IP version and IHL
diff --git a/tests/net/java/android/net/nsd/NsdManagerTest.java b/tests/net/java/android/net/nsd/NsdManagerTest.java
index 2d2bccb..cf7587a 100644
--- a/tests/net/java/android/net/nsd/NsdManagerTest.java
+++ b/tests/net/java/android/net/nsd/NsdManagerTest.java
@@ -16,8 +16,6 @@
 
 package android.net.nsd;
 
-import static com.android.internal.util.TestUtils.waitForIdleHandler;
-
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.fail;
@@ -40,6 +38,7 @@
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.util.AsyncChannel;
+import com.android.testutils.HandlerUtilsKt;
 
 import org.junit.After;
 import org.junit.Before;
@@ -74,7 +73,7 @@
 
     @After
     public void tearDown() throws Exception {
-        mServiceHandler.waitForIdle(mTimeoutMs);
+        HandlerUtilsKt.waitForIdle(mServiceHandler, mTimeoutMs);
         mServiceHandler.chan.disconnect();
         mServiceHandler.stop();
         if (mManager != null) {
@@ -334,7 +333,7 @@
     }
 
     int verifyRequest(int expectedMessageType) {
-        mServiceHandler.waitForIdle(mTimeoutMs);
+        HandlerUtilsKt.waitForIdle(mServiceHandler, mTimeoutMs);
         verify(mServiceHandler, timeout(mTimeoutMs)).handleMessage(any());
         reset(mServiceHandler);
         Message received = mServiceHandler.getLastMessage();
@@ -366,10 +365,6 @@
             lastMessage.copyFrom(msg);
         }
 
-        void waitForIdle(long timeoutMs) {
-            waitForIdleHandler(this, timeoutMs);
-        }
-
         @Override
         public void handleMessage(Message msg) {
             setLastMessage(msg);
diff --git a/tests/net/java/android/net/util/DnsUtilsTest.java b/tests/net/java/android/net/util/DnsUtilsTest.java
index 42e340b..b626db8 100644
--- a/tests/net/java/android/net/util/DnsUtilsTest.java
+++ b/tests/net/java/android/net/util/DnsUtilsTest.java
@@ -57,24 +57,38 @@
     @Test
     public void testRfc6724Comparator() {
         final List<DnsUtils.SortableAddress> test = Arrays.asList(
-                makeSortableAddress("216.58.200.36"),             // Ipv4
-                makeSortableAddress("2404:6800:4008:801::2004"),  // global
-                makeSortableAddress("::1"),                       // loop back
-                makeSortableAddress("fe80::c46f:1cff:fe04:39b4"), // link local
-                makeSortableAddress("::ffff:192.168.95.3"),       // IPv4-mapped IPv6
-                makeSortableAddress("2001::47c1"),                // teredo tunneling
-                makeSortableAddress("::216.58.200.36"),           // IPv4-compatible
-                makeSortableAddress("3ffe::1234:5678"));          // 6bone
+                // Ipv4
+                makeSortableAddress("216.58.200.36", "192.168.1.1"),
+                // global with different scope src
+                makeSortableAddress("2404:6800:4008:801::2004", "fe80::1111:2222"),
+                // global without src addr
+                makeSortableAddress("2404:6800:cafe:801::1"),
+                // loop back
+                makeSortableAddress("::1", "::1"),
+                // link local
+                makeSortableAddress("fe80::c46f:1cff:fe04:39b4", "fe80::1"),
+                // teredo tunneling
+                makeSortableAddress("2001::47c1", "2001::2"),
+                // 6bone without src addr
+                makeSortableAddress("3ffe::1234:5678"),
+                // IPv4-compatible
+                makeSortableAddress("::216.58.200.36", "::216.58.200.9"),
+                // 6bone
+                makeSortableAddress("3ffe::1234:5678", "3ffe::1234:1"),
+                // IPv4-mapped IPv6
+                makeSortableAddress("::ffff:192.168.95.7", "::ffff:192.168.95.1"));
 
         final List<InetAddress> expected = Arrays.asList(
                 stringToAddress("::1"),                       // loop back
                 stringToAddress("fe80::c46f:1cff:fe04:39b4"), // link local
-                stringToAddress("2404:6800:4008:801::2004"),  // global
                 stringToAddress("216.58.200.36"),             // Ipv4
-                stringToAddress("::ffff:192.168.95.3"),       // IPv4-mapped IPv6
+                stringToAddress("::ffff:192.168.95.7"),       // IPv4-mapped IPv6
                 stringToAddress("2001::47c1"),                // teredo tunneling
-                stringToAddress("::216.58.200.36"),            // IPv4-compatible
-                stringToAddress("3ffe::1234:5678"));          // 6bone
+                stringToAddress("::216.58.200.36"),           // IPv4-compatible
+                stringToAddress("3ffe::1234:5678"),           // 6bone
+                stringToAddress("2404:6800:4008:801::2004"),  // global with different scope src
+                stringToAddress("2404:6800:cafe:801::1"),     // global without src addr
+                stringToAddress("3ffe::1234:5678"));          // 6bone without src addr
 
         Collections.sort(test, new DnsUtils.Rfc6724Comparator());
 
diff --git a/tests/net/java/com/android/internal/util/RingBufferTest.java b/tests/net/java/com/android/internal/util/RingBufferTest.java
index eff334f..d06095a 100644
--- a/tests/net/java/com/android/internal/util/RingBufferTest.java
+++ b/tests/net/java/com/android/internal/util/RingBufferTest.java
@@ -16,6 +16,7 @@
 
 package com.android.internal.util;
 
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.fail;
 
@@ -25,9 +26,6 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import java.util.Arrays;
-import java.util.Objects;
-
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class RingBufferTest {
@@ -36,7 +34,7 @@
     public void testEmptyRingBuffer() {
         RingBuffer<String> buffer = new RingBuffer<>(String.class, 100);
 
-        assertArraysEqual(new String[0], buffer.toArray());
+        assertArrayEquals(new String[0], buffer.toArray());
     }
 
     @Test
@@ -65,7 +63,7 @@
         buffer.append("e");
 
         String[] expected = {"a", "b", "c", "d", "e"};
-        assertArraysEqual(expected, buffer.toArray());
+        assertArrayEquals(expected, buffer.toArray());
     }
 
     @Test
@@ -73,19 +71,19 @@
         RingBuffer<String> buffer = new RingBuffer<>(String.class, 1);
 
         buffer.append("a");
-        assertArraysEqual(new String[]{"a"}, buffer.toArray());
+        assertArrayEquals(new String[]{"a"}, buffer.toArray());
 
         buffer.append("b");
-        assertArraysEqual(new String[]{"b"}, buffer.toArray());
+        assertArrayEquals(new String[]{"b"}, buffer.toArray());
 
         buffer.append("c");
-        assertArraysEqual(new String[]{"c"}, buffer.toArray());
+        assertArrayEquals(new String[]{"c"}, buffer.toArray());
 
         buffer.append("d");
-        assertArraysEqual(new String[]{"d"}, buffer.toArray());
+        assertArrayEquals(new String[]{"d"}, buffer.toArray());
 
         buffer.append("e");
-        assertArraysEqual(new String[]{"e"}, buffer.toArray());
+        assertArrayEquals(new String[]{"e"}, buffer.toArray());
     }
 
     @Test
@@ -100,7 +98,7 @@
         buffer.append("e");
 
         String[] expected1 = {"a", "b", "c", "d", "e"};
-        assertArraysEqual(expected1, buffer.toArray());
+        assertArrayEquals(expected1, buffer.toArray());
 
         String[] expected2 = new String[capacity];
         int firstIndex = 0;
@@ -111,22 +109,22 @@
             buffer.append("x");
             expected2[i] = "x";
         }
-        assertArraysEqual(expected2, buffer.toArray());
+        assertArrayEquals(expected2, buffer.toArray());
 
         buffer.append("x");
         expected2[firstIndex] = "x";
-        assertArraysEqual(expected2, buffer.toArray());
+        assertArrayEquals(expected2, buffer.toArray());
 
         for (int i = 0; i < 10; i++) {
             for (String s : expected2) {
                 buffer.append(s);
             }
         }
-        assertArraysEqual(expected2, buffer.toArray());
+        assertArrayEquals(expected2, buffer.toArray());
 
         buffer.append("a");
         expected2[lastIndex] = "a";
-        assertArraysEqual(expected2, buffer.toArray());
+        assertArrayEquals(expected2, buffer.toArray());
     }
 
     @Test
@@ -143,7 +141,7 @@
             expected[i] = new DummyClass1();
             expected[i].x = capacity * i;
         }
-        assertArraysEqual(expected, buffer.toArray());
+        assertArrayEquals(expected, buffer.toArray());
 
         for (int i = 0; i < capacity; ++i) {
             if (actual[i] != buffer.getNextSlot()) {
@@ -177,18 +175,4 @@
     }
 
     private static final class DummyClass3 {}
-
-    static <T> void assertArraysEqual(T[] expected, T[] got) {
-        if (expected.length != got.length) {
-            fail(Arrays.toString(expected) + " and " + Arrays.toString(got)
-                    + " did not have the same length");
-        }
-
-        for (int i = 0; i < expected.length; i++) {
-            if (!Objects.equals(expected[i], got[i])) {
-                fail(Arrays.toString(expected) + " and " + Arrays.toString(got)
-                        + " were not equal");
-            }
-        }
-    }
 }
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index 37cc304..f3c735c 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -28,8 +28,6 @@
 import static android.net.ConnectivityManager.TYPE_MOBILE;
 import static android.net.ConnectivityManager.TYPE_MOBILE_FOTA;
 import static android.net.ConnectivityManager.TYPE_MOBILE_MMS;
-import static android.net.ConnectivityManager.TYPE_NONE;
-import static android.net.ConnectivityManager.TYPE_VPN;
 import static android.net.ConnectivityManager.TYPE_WIFI;
 import static android.net.INetworkMonitor.NETWORK_VALIDATION_PROBE_DNS;
 import static android.net.INetworkMonitor.NETWORK_VALIDATION_PROBE_FALLBACK;
@@ -69,11 +67,17 @@
 import static android.net.NetworkPolicyManager.RULE_REJECT_METERED;
 import static android.net.RouteInfo.RTN_UNREACHABLE;
 
-import static com.android.internal.util.TestUtils.waitForIdleHandler;
-import static com.android.internal.util.TestUtils.waitForIdleLooper;
-import static com.android.internal.util.TestUtils.waitForIdleSerialExecutor;
+import static com.android.server.ConnectivityServiceTestUtilsKt.transportToLegacyType;
+import static com.android.testutils.ConcurrentUtilsKt.await;
+import static com.android.testutils.ConcurrentUtilsKt.durationOf;
+import static com.android.testutils.ExceptionUtils.ignoreExceptions;
+import static com.android.testutils.HandlerUtilsKt.waitForIdleSerialExecutor;
+import static com.android.testutils.MiscAssertsKt.assertContainsExactly;
+import static com.android.testutils.MiscAssertsKt.assertEmpty;
+import static com.android.testutils.MiscAssertsKt.assertLength;
+import static com.android.testutils.MiscAssertsKt.assertRunsInAtMost;
+import static com.android.testutils.MiscAssertsKt.assertThrows;
 
-import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
@@ -81,6 +85,7 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Matchers.anyInt;
@@ -88,6 +93,7 @@
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -100,6 +106,7 @@
 import static org.mockito.Mockito.when;
 
 import android.annotation.NonNull;
+import android.app.AlarmManager;
 import android.app.NotificationManager;
 import android.app.PendingIntent;
 import android.content.BroadcastReceiver;
@@ -120,6 +127,7 @@
 import android.net.ConnectivityManager.TooManyRequestsException;
 import android.net.ConnectivityThread;
 import android.net.IDnsResolver;
+import android.net.IIpConnectivityMetrics;
 import android.net.INetd;
 import android.net.INetworkMonitor;
 import android.net.INetworkMonitorCallbacks;
@@ -134,12 +142,8 @@
 import android.net.LinkProperties;
 import android.net.MatchAllNetworkSpecifier;
 import android.net.Network;
-import android.net.NetworkAgent;
 import android.net.NetworkCapabilities;
 import android.net.NetworkFactory;
-import android.net.NetworkInfo;
-import android.net.NetworkInfo.DetailedState;
-import android.net.NetworkMisc;
 import android.net.NetworkRequest;
 import android.net.NetworkSpecifier;
 import android.net.NetworkStack;
@@ -155,6 +159,7 @@
 import android.net.shared.NetworkMonitorUtils;
 import android.net.shared.PrivateDnsConfig;
 import android.net.util.MultinetworkPolicyTracker;
+import android.os.BadParcelableException;
 import android.os.Binder;
 import android.os.Bundle;
 import android.os.ConditionVariable;
@@ -162,7 +167,6 @@
 import android.os.HandlerThread;
 import android.os.INetworkManagementService;
 import android.os.Looper;
-import android.os.Message;
 import android.os.Parcel;
 import android.os.ParcelFileDescriptor;
 import android.os.Parcelable;
@@ -200,7 +204,10 @@
 import com.android.server.connectivity.Vpn;
 import com.android.server.net.NetworkPinner;
 import com.android.server.net.NetworkPolicyManagerInternal;
-import com.android.server.net.NetworkStatsFactory;
+import com.android.testutils.ExceptionUtils;
+import com.android.testutils.HandlerUtilsKt;
+import com.android.testutils.RecorderCallback.CallbackRecord;
+import com.android.testutils.TestableNetworkCallback;
 
 import org.junit.After;
 import org.junit.Before;
@@ -221,14 +228,12 @@
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.Socket;
-import java.net.UnknownHostException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
@@ -237,7 +242,8 @@
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.function.Predicate;
+
+import kotlin.reflect.KClass;
 
 /**
  * Tests for {@link ConnectivityService}.
@@ -257,10 +263,12 @@
     // timeout. For this, our assertions should run fast enough to leave less than
     // (mService.mLingerDelayMs - TEST_CALLBACK_TIMEOUT_MS) between the time callbacks are
     // supposedly fired, and the time we call expectCallback.
-    private final static int TEST_CALLBACK_TIMEOUT_MS = 200;
+    private static final int TEST_CALLBACK_TIMEOUT_MS = 200;
     // Chosen to be less than TEST_CALLBACK_TIMEOUT_MS. This ensures that requests have time to
     // complete before callbacks are verified.
-    private final static int TEST_REQUEST_TIMEOUT_MS = 150;
+    private static final int TEST_REQUEST_TIMEOUT_MS = 150;
+
+    private static final int UNREASONABLY_LONG_ALARM_WAIT_MS = 1000;
 
     private static final String CLAT_PREFIX = "v4-";
     private static final String MOBILE_IFNAME = "test_rmnet_data0";
@@ -268,15 +276,19 @@
     private static final String[] EMPTY_STRING_ARRAY = new String[0];
 
     private MockContext mServiceContext;
-    private WrappedConnectivityService mService;
+    private HandlerThread mCsHandlerThread;
+    private ConnectivityService mService;
     private WrappedConnectivityManager mCm;
-    private MockNetworkAgent mWiFiNetworkAgent;
-    private MockNetworkAgent mCellNetworkAgent;
-    private MockNetworkAgent mEthernetNetworkAgent;
+    private TestNetworkAgentWrapper mWiFiNetworkAgent;
+    private TestNetworkAgentWrapper mCellNetworkAgent;
+    private TestNetworkAgentWrapper mEthernetNetworkAgent;
     private MockVpn mMockVpn;
     private Context mContext;
     private INetworkPolicyListener mPolicyListener;
+    private WrappedMultinetworkPolicyTracker mPolicyTracker;
+    private HandlerThread mAlarmManagerThread;
 
+    @Mock IIpConnectivityMetrics mIpConnectivityMetrics;
     @Mock IpConnectivityMetrics.Logger mMetricsService;
     @Mock DefaultNetworkMetrics mDefaultNetworkMetrics;
     @Mock INetworkManagementService mNetworkManagementService;
@@ -288,6 +300,7 @@
     @Mock PackageManager mPackageManager;
     @Mock UserManager mUserManager;
     @Mock NotificationManager mNotificationManager;
+    @Mock AlarmManager mAlarmManager;
 
     private ArgumentCaptor<ResolverParamsParcel> mResolverParamsParcelCaptor =
             ArgumentCaptor.forClass(ResolverParamsParcel.class);
@@ -360,6 +373,7 @@
             if (Context.NOTIFICATION_SERVICE.equals(name)) return mNotificationManager;
             if (Context.NETWORK_STACK_SERVICE.equals(name)) return mNetworkStack;
             if (Context.USER_SERVICE.equals(name)) return mUserManager;
+            if (Context.ALARM_SERVICE.equals(name)) return mAlarmManager;
             return super.getSystemService(name);
         }
 
@@ -389,29 +403,24 @@
         }
     }
 
-    public void waitForIdle(int timeoutMsAsInt) {
-        long timeoutMs = timeoutMsAsInt;
-        waitForIdleHandler(mService.mHandlerThread, timeoutMs);
-        waitForIdle(mCellNetworkAgent, timeoutMs);
-        waitForIdle(mWiFiNetworkAgent, timeoutMs);
-        waitForIdle(mEthernetNetworkAgent, timeoutMs);
-        waitForIdleHandler(mService.mHandlerThread, timeoutMs);
-        waitForIdleLooper(ConnectivityThread.getInstanceLooper(), timeoutMs);
+    private void waitForIdle() {
+        HandlerUtilsKt.waitForIdle(mCsHandlerThread, TIMEOUT_MS);
+        waitForIdle(mCellNetworkAgent, TIMEOUT_MS);
+        waitForIdle(mWiFiNetworkAgent, TIMEOUT_MS);
+        waitForIdle(mEthernetNetworkAgent, TIMEOUT_MS);
+        HandlerUtilsKt.waitForIdle(mCsHandlerThread, TIMEOUT_MS);
+        HandlerUtilsKt.waitForIdle(ConnectivityThread.get(), TIMEOUT_MS);
     }
 
-    public void waitForIdle(MockNetworkAgent agent, long timeoutMs) {
+    private void waitForIdle(TestNetworkAgentWrapper agent, long timeoutMs) {
         if (agent == null) {
             return;
         }
-        waitForIdleHandler(agent.mHandlerThread, timeoutMs);
-    }
-
-    private void waitForIdle() {
-        waitForIdle(TIMEOUT_MS);
+        agent.waitForIdle(timeoutMs);
     }
 
     @Test
-    public void testWaitForIdle() {
+    public void testWaitForIdle() throws Exception {
         final int attempts = 50;  // Causes the test to take about 200ms on bullhead-eng.
 
         // Tests that waitForIdle returns immediately if the service is already idle.
@@ -421,7 +430,7 @@
 
         // Bring up a network that we can use to send messages to ConnectivityService.
         ConditionVariable cv = waitForConnectivityBroadcasts(1);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
         waitFor(cv);
         Network n = mWiFiNetworkAgent.getNetwork();
@@ -438,10 +447,10 @@
     // This test has an inherent race condition in it, and cannot be enabled for continuous testing
     // or presubmit tests. It is kept for manual runs and documentation purposes.
     @Ignore
-    public void verifyThatNotWaitingForIdleCausesRaceConditions() {
+    public void verifyThatNotWaitingForIdleCausesRaceConditions() throws Exception {
         // Bring up a network that we can use to send messages to ConnectivityService.
         ConditionVariable cv = waitForConnectivityBroadcasts(1);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
         waitFor(cv);
         Network n = mWiFiNetworkAgent.getNetwork();
@@ -461,7 +470,7 @@
         fail("expected race condition at least once in " + attempts + " attempts");
     }
 
-    private class MockNetworkAgent {
+    private class TestNetworkAgentWrapper extends NetworkAgentWrapper {
         private static final int VALIDATION_RESULT_BASE = NETWORK_VALIDATION_PROBE_DNS
                 | NETWORK_VALIDATION_PROBE_HTTP
                 | NETWORK_VALIDATION_PROBE_HTTPS;
@@ -472,97 +481,42 @@
                 | NETWORK_VALIDATION_RESULT_PARTIAL;
         private static final int VALIDATION_RESULT_INVALID = 0;
 
-        private final INetworkMonitor mNetworkMonitor;
-        private final NetworkInfo mNetworkInfo;
-        private final NetworkCapabilities mNetworkCapabilities;
-        private final HandlerThread mHandlerThread;
-        private final ConditionVariable mDisconnected = new ConditionVariable();
-        private final ConditionVariable mNetworkStatusReceived = new ConditionVariable();
-        private final ConditionVariable mPreventReconnectReceived = new ConditionVariable();
-        private int mScore;
-        private NetworkAgent mNetworkAgent;
-        private int mStartKeepaliveError = SocketKeepalive.ERROR_UNSUPPORTED;
-        private int mStopKeepaliveError = SocketKeepalive.NO_KEEPALIVE;
-        private Integer mExpectedKeepaliveSlot = null;
-        // Contains the redirectUrl from networkStatus(). Before reading, wait for
-        // mNetworkStatusReceived.
-        private String mRedirectUrl;
-
+        private INetworkMonitor mNetworkMonitor;
         private INetworkMonitorCallbacks mNmCallbacks;
         private int mNmValidationResult = VALIDATION_RESULT_BASE;
         private String mNmValidationRedirectUrl = null;
         private boolean mNmProvNotificationRequested = false;
 
-        void setNetworkValid() {
-            mNmValidationResult = VALIDATION_RESULT_VALID;
-            mNmValidationRedirectUrl = null;
-        }
+        private final ConditionVariable mNetworkStatusReceived = new ConditionVariable();
+        // Contains the redirectUrl from networkStatus(). Before reading, wait for
+        // mNetworkStatusReceived.
+        private String mRedirectUrl;
 
-        void setNetworkInvalid() {
-            mNmValidationResult = VALIDATION_RESULT_INVALID;
-            mNmValidationRedirectUrl = null;
-        }
-
-        void setNetworkPortal(String redirectUrl) {
-            setNetworkInvalid();
-            mNmValidationRedirectUrl = redirectUrl;
-        }
-
-        void setNetworkPartial() {
-            mNmValidationResult = VALIDATION_RESULT_PARTIAL;
-            mNmValidationRedirectUrl = null;
-        }
-
-        void setNetworkPartialValid() {
-            mNmValidationResult = VALIDATION_RESULT_PARTIAL | VALIDATION_RESULT_VALID;
-            mNmValidationRedirectUrl = null;
-        }
-
-        MockNetworkAgent(int transport) {
+        TestNetworkAgentWrapper(int transport) throws Exception {
             this(transport, new LinkProperties());
         }
 
-        MockNetworkAgent(int transport, LinkProperties linkProperties) {
-            final int type = transportToLegacyType(transport);
-            final String typeName = ConnectivityManager.getNetworkTypeName(type);
-            mNetworkInfo = new NetworkInfo(type, 0, typeName, "Mock");
-            mNetworkCapabilities = new NetworkCapabilities();
-            mNetworkCapabilities.addTransportType(transport);
-            switch (transport) {
-                case TRANSPORT_ETHERNET:
-                    mScore = 70;
-                    break;
-                case TRANSPORT_WIFI:
-                    mScore = 60;
-                    break;
-                case TRANSPORT_CELLULAR:
-                    mScore = 50;
-                    break;
-                case TRANSPORT_WIFI_AWARE:
-                    mScore = 20;
-                    break;
-                case TRANSPORT_VPN:
-                    mNetworkCapabilities.removeCapability(NET_CAPABILITY_NOT_VPN);
-                    mScore = ConnectivityConstants.VPN_DEFAULT_SCORE;
-                    break;
-                default:
-                    throw new UnsupportedOperationException("unimplemented network type");
-            }
-            mHandlerThread = new HandlerThread("Mock-" + typeName);
-            mHandlerThread.start();
+        TestNetworkAgentWrapper(int transport, LinkProperties linkProperties)
+                throws Exception {
+            super(transport, linkProperties, mServiceContext);
 
+            // Waits for the NetworkAgent to be registered, which includes the creation of the
+            // NetworkMonitor.
+            waitForIdle(TIMEOUT_MS);
+        }
+
+        @Override
+        protected InstrumentedNetworkAgent makeNetworkAgent(LinkProperties linkProperties)
+                throws Exception {
             mNetworkMonitor = mock(INetworkMonitor.class);
+
             final Answer validateAnswer = inv -> {
-                new Thread(this::onValidationRequested).start();
+                new Thread(ignoreExceptions(this::onValidationRequested)).start();
                 return null;
             };
 
-            try {
-                doAnswer(validateAnswer).when(mNetworkMonitor).notifyNetworkConnected(any(), any());
-                doAnswer(validateAnswer).when(mNetworkMonitor).forceReevaluation(anyInt());
-            } catch (RemoteException e) {
-                fail(e.getMessage());
-            }
+            doAnswer(validateAnswer).when(mNetworkMonitor).notifyNetworkConnected(any(), any());
+            doAnswer(validateAnswer).when(mNetworkMonitor).forceReevaluation(anyInt());
 
             final ArgumentCaptor<Network> nmNetworkCaptor = ArgumentCaptor.forClass(Network.class);
             final ArgumentCaptor<INetworkMonitorCallbacks> nmCbCaptor =
@@ -572,132 +526,44 @@
                     any() /* name */,
                     nmCbCaptor.capture());
 
-            mNetworkAgent = new NetworkAgent(mHandlerThread.getLooper(), mServiceContext,
-                    "Mock-" + typeName, mNetworkInfo, mNetworkCapabilities,
-                    linkProperties, mScore, new NetworkMisc(), NetworkFactory.SerialNumber.NONE) {
-                @Override
-                public void unwanted() { mDisconnected.open(); }
-
-                @Override
-                public void startSocketKeepalive(Message msg) {
-                    int slot = msg.arg1;
-                    if (mExpectedKeepaliveSlot != null) {
-                        assertEquals((int) mExpectedKeepaliveSlot, slot);
-                    }
-                    onSocketKeepaliveEvent(slot, mStartKeepaliveError);
-                }
-
-                @Override
-                public void stopSocketKeepalive(Message msg) {
-                    onSocketKeepaliveEvent(msg.arg1, mStopKeepaliveError);
-                }
-
+            final InstrumentedNetworkAgent na = new InstrumentedNetworkAgent(this, linkProperties) {
                 @Override
                 public void networkStatus(int status, String redirectUrl) {
                     mRedirectUrl = redirectUrl;
                     mNetworkStatusReceived.open();
                 }
-
-                @Override
-                protected void preventAutomaticReconnect() {
-                    mPreventReconnectReceived.open();
-                }
-
-                @Override
-                protected void addKeepalivePacketFilter(Message msg) {
-                    Log.i(TAG, "Add keepalive packet filter.");
-                }
-
-                @Override
-                protected void removeKeepalivePacketFilter(Message msg) {
-                    Log.i(TAG, "Remove keepalive packet filter.");
-                }
             };
 
-            assertEquals(mNetworkAgent.netId, nmNetworkCaptor.getValue().netId);
+            assertEquals(na.netId, nmNetworkCaptor.getValue().netId);
             mNmCallbacks = nmCbCaptor.getValue();
 
-            try {
-                mNmCallbacks.onNetworkMonitorCreated(mNetworkMonitor);
-            } catch (RemoteException e) {
-                fail(e.getMessage());
-            }
+            mNmCallbacks.onNetworkMonitorCreated(mNetworkMonitor);
 
-            // Waits for the NetworkAgent to be registered, which includes the creation of the
-            // NetworkMonitor.
-            waitForIdle();
+            return na;
         }
 
-        private void onValidationRequested() {
-            try {
-                if (mNmProvNotificationRequested
-                        && ((mNmValidationResult & NETWORK_VALIDATION_RESULT_VALID) != 0)) {
-                    mNmCallbacks.hideProvisioningNotification();
-                    mNmProvNotificationRequested = false;
-                }
+        private void onValidationRequested() throws Exception {
+            if (mNmProvNotificationRequested
+                    && ((mNmValidationResult & NETWORK_VALIDATION_RESULT_VALID) != 0)) {
+                mNmCallbacks.hideProvisioningNotification();
+                mNmProvNotificationRequested = false;
+            }
 
-                mNmCallbacks.notifyNetworkTested(
-                        mNmValidationResult, mNmValidationRedirectUrl);
+            mNmCallbacks.notifyNetworkTested(
+                    mNmValidationResult, mNmValidationRedirectUrl);
 
-                if (mNmValidationRedirectUrl != null) {
-                    mNmCallbacks.showProvisioningNotification(
-                            "test_provisioning_notif_action", "com.android.test.package");
-                    mNmProvNotificationRequested = true;
-                }
-            } catch (RemoteException e) {
-                fail(e.getMessage());
+            if (mNmValidationRedirectUrl != null) {
+                mNmCallbacks.showProvisioningNotification(
+                        "test_provisioning_notif_action", "com.android.test.package");
+                mNmProvNotificationRequested = true;
             }
         }
 
-        public void adjustScore(int change) {
-            mScore += change;
-            mNetworkAgent.sendNetworkScore(mScore);
-        }
-
-        public int getScore() {
-            return mScore;
-        }
-
-        public void explicitlySelected(boolean acceptUnvalidated) {
-            mNetworkAgent.explicitlySelected(acceptUnvalidated);
-        }
-
-        public void addCapability(int capability) {
-            mNetworkCapabilities.addCapability(capability);
-            mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
-        }
-
-        public void removeCapability(int capability) {
-            mNetworkCapabilities.removeCapability(capability);
-            mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
-        }
-
-        public void setUids(Set<UidRange> uids) {
-            mNetworkCapabilities.setUids(uids);
-            mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
-        }
-
-        public void setSignalStrength(int signalStrength) {
-            mNetworkCapabilities.setSignalStrength(signalStrength);
-            mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
-        }
-
-        public void setNetworkSpecifier(NetworkSpecifier networkSpecifier) {
-            mNetworkCapabilities.setNetworkSpecifier(networkSpecifier);
-            mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
-        }
-
-        public void setNetworkCapabilities(NetworkCapabilities nc,
-                boolean sendToConnectivityService) {
-            mNetworkCapabilities.set(nc);
-            if (sendToConnectivityService) {
-                mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
-            }
-        }
-
+        /**
+         * Connect without adding any internet capability.
+         */
         public void connectWithoutInternet() {
-            mNetworkInfo.setDetailedState(DetailedState.CONNECTED, null, null);
-            mNetworkAgent.sendNetworkInfo(mNetworkInfo);
+            super.connect();
         }
 
         /**
@@ -714,23 +580,21 @@
          * @param hasInternet Indicate if network should pretend to have NET_CAPABILITY_INTERNET.
          */
         public void connect(boolean validated, boolean hasInternet) {
-            assertEquals("MockNetworkAgents can only be connected once",
-                    mNetworkInfo.getDetailedState(), DetailedState.IDLE);
-            assertFalse(mNetworkCapabilities.hasCapability(NET_CAPABILITY_INTERNET));
+            assertFalse(getNetworkCapabilities().hasCapability(NET_CAPABILITY_INTERNET));
 
-            NetworkCallback callback = null;
+            ConnectivityManager.NetworkCallback callback = null;
             final ConditionVariable validatedCv = new ConditionVariable();
             if (validated) {
                 setNetworkValid();
                 NetworkRequest request = new NetworkRequest.Builder()
-                        .addTransportType(mNetworkCapabilities.getTransportTypes()[0])
+                        .addTransportType(getNetworkCapabilities().getTransportTypes()[0])
                         .clearCapabilities()
                         .build();
-                callback = new NetworkCallback() {
+                callback = new ConnectivityManager.NetworkCallback() {
                     public void onCapabilitiesChanged(Network network,
                             NetworkCapabilities networkCapabilities) {
                         if (network.equals(getNetwork()) &&
-                            networkCapabilities.hasCapability(NET_CAPABILITY_VALIDATED)) {
+                                networkCapabilities.hasCapability(NET_CAPABILITY_VALIDATED)) {
                             validatedCv.open();
                         }
                     }
@@ -762,47 +626,34 @@
             connect(false);
         }
 
-        public void suspend() {
-            mNetworkInfo.setDetailedState(DetailedState.SUSPENDED, null, null);
-            mNetworkAgent.sendNetworkInfo(mNetworkInfo);
+        public void connectWithPartialValidConnectivity() {
+            setNetworkPartialValid();
+            connect(false);
         }
 
-        public void resume() {
-            mNetworkInfo.setDetailedState(DetailedState.CONNECTED, null, null);
-            mNetworkAgent.sendNetworkInfo(mNetworkInfo);
+        void setNetworkValid() {
+            mNmValidationResult = VALIDATION_RESULT_VALID;
+            mNmValidationRedirectUrl = null;
         }
 
-        public void disconnect() {
-            mNetworkInfo.setDetailedState(DetailedState.DISCONNECTED, null, null);
-            mNetworkAgent.sendNetworkInfo(mNetworkInfo);
+        void setNetworkInvalid() {
+            mNmValidationResult = VALIDATION_RESULT_INVALID;
+            mNmValidationRedirectUrl = null;
         }
 
-        public Network getNetwork() {
-            return new Network(mNetworkAgent.netId);
+        void setNetworkPortal(String redirectUrl) {
+            setNetworkInvalid();
+            mNmValidationRedirectUrl = redirectUrl;
         }
 
-        public ConditionVariable getPreventReconnectReceived() {
-            return mPreventReconnectReceived;
+        void setNetworkPartial() {
+            mNmValidationResult = VALIDATION_RESULT_PARTIAL;
+            mNmValidationRedirectUrl = null;
         }
 
-        public ConditionVariable getDisconnectedCV() {
-            return mDisconnected;
-        }
-
-        public void sendLinkProperties(LinkProperties lp) {
-            mNetworkAgent.sendLinkProperties(lp);
-        }
-
-        public void setStartKeepaliveError(int error) {
-            mStartKeepaliveError = error;
-        }
-
-        public void setStopKeepaliveError(int error) {
-            mStopKeepaliveError = error;
-        }
-
-        public void setExpectedKeepaliveSlot(Integer slot) {
-            mExpectedKeepaliveSlot = slot;
+        void setNetworkPartialValid() {
+            mNmValidationResult = VALIDATION_RESULT_PARTIAL | VALIDATION_RESULT_VALID;
+            mNmValidationRedirectUrl = null;
         }
 
         public String waitForRedirectUrl() {
@@ -810,12 +661,12 @@
             return mRedirectUrl;
         }
 
-        public NetworkAgent getNetworkAgent() {
-            return mNetworkAgent;
+        public void expectDisconnected() {
+            expectDisconnected(TIMEOUT_MS);
         }
 
-        public NetworkCapabilities getNetworkCapabilities() {
-            return mNetworkCapabilities;
+        public void expectPreventReconnectReceived() {
+            expectPreventReconnectReceived(TIMEOUT_MS);
         }
     }
 
@@ -994,15 +845,15 @@
         private boolean mConnected = false;
         // Careful ! This is different from mNetworkAgent, because MockNetworkAgent does
         // not inherit from NetworkAgent.
-        private MockNetworkAgent mMockNetworkAgent;
+        private TestNetworkAgentWrapper mMockNetworkAgent;
 
         public MockVpn(int userId) {
             super(startHandlerThreadAndReturnLooper(), mServiceContext, mNetworkManagementService,
                     userId);
         }
 
-        public void setNetworkAgent(MockNetworkAgent agent) {
-            waitForIdle(agent, TIMEOUT_MS);
+        public void setNetworkAgent(TestNetworkAgentWrapper agent) {
+            agent.waitForIdle(TIMEOUT_MS);
             mMockNetworkAgent = agent;
             mNetworkAgent = agent.getNetworkAgent();
             mNetworkCapabilities.set(agent.getNetworkCapabilities());
@@ -1069,192 +920,45 @@
         }
     }
 
-    private class FakeWakeupMessage extends WakeupMessage {
-        private static final int UNREASONABLY_LONG_WAIT = 1000;
-
-        public FakeWakeupMessage(Context context, Handler handler, String cmdName, int cmd) {
-            super(context, handler, cmdName, cmd);
-        }
-
-        public FakeWakeupMessage(Context context, Handler handler, String cmdName, int cmd,
-                int arg1, int arg2, Object obj) {
-            super(context, handler, cmdName, cmd, arg1, arg2, obj);
-        }
-
-        @Override
-        public void schedule(long when) {
-            long delayMs = when - SystemClock.elapsedRealtime();
-            if (delayMs < 0) delayMs = 0;
-            if (delayMs > UNREASONABLY_LONG_WAIT) {
-                fail("Attempting to send msg more than " + UNREASONABLY_LONG_WAIT +
-                        "ms into the future: " + delayMs);
-            }
-            Message msg = mHandler.obtainMessage(mCmd, mArg1, mArg2, mObj);
-            mHandler.sendMessageDelayed(msg, delayMs);
-        }
-
-        @Override
-        public void cancel() {
-            mHandler.removeMessages(mCmd, mObj);
-        }
-
-        @Override
-        public void onAlarm() {
-            throw new AssertionError("Should never happen. Update this fake.");
+    private void mockVpn(int uid) {
+        synchronized (mService.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.
+            mService.mVpns.put(userId, mMockVpn);
         }
     }
 
-    private class WrappedMultinetworkPolicyTracker extends MultinetworkPolicyTracker {
-        public volatile boolean configRestrictsAvoidBadWifi;
-        public volatile int configMeteredMultipathPreference;
+    private void setUidRulesChanged(int uidRules) throws RemoteException {
+        mPolicyListener.onUidRulesChanged(Process.myUid(), uidRules);
+    }
 
-        public WrappedMultinetworkPolicyTracker(Context c, Handler h, Runnable r) {
+    private void setRestrictBackgroundChanged(boolean restrictBackground) throws RemoteException {
+        mPolicyListener.onRestrictBackgroundChanged(restrictBackground);
+    }
+
+    private Nat464Xlat getNat464Xlat(NetworkAgentWrapper mna) {
+        return mService.getNetworkAgentInfoForNetwork(mna.getNetwork()).clatd;
+    }
+
+    private static class WrappedMultinetworkPolicyTracker extends MultinetworkPolicyTracker {
+        volatile boolean mConfigRestrictsAvoidBadWifi;
+        volatile int mConfigMeteredMultipathPreference;
+
+        WrappedMultinetworkPolicyTracker(Context c, Handler h, Runnable r) {
             super(c, h, r);
         }
 
         @Override
         public boolean configRestrictsAvoidBadWifi() {
-            return configRestrictsAvoidBadWifi;
+            return mConfigRestrictsAvoidBadWifi;
         }
 
         @Override
         public int configMeteredMultipathPreference() {
-            return configMeteredMultipathPreference;
-        }
-    }
-
-    private class WrappedConnectivityService extends ConnectivityService {
-        public WrappedMultinetworkPolicyTracker wrappedMultinetworkPolicyTracker;
-        private MockableSystemProperties mSystemProperties;
-
-        public WrappedConnectivityService(Context context, INetworkManagementService netManager,
-                INetworkStatsService statsService, INetworkPolicyManager policyManager,
-                IpConnectivityLog log, INetd netd, IDnsResolver dnsResolver) {
-            super(context, netManager, statsService, policyManager, dnsResolver, log, netd);
-            mNetd = netd;
-            mLingerDelayMs = TEST_LINGER_DELAY_MS;
-        }
-
-        @Override
-        protected MockableSystemProperties getSystemProperties() {
-            // Minimal approach to overriding system properties: let most calls fall through to real
-            // device values, and only override ones values that are important to this test.
-            mSystemProperties = spy(new MockableSystemProperties());
-            when(mSystemProperties.getInt("net.tcp.default_init_rwnd", 0)).thenReturn(0);
-            when(mSystemProperties.getBoolean("ro.radio.noril", false)).thenReturn(false);
-            return mSystemProperties;
-        }
-
-        @Override
-        protected Tethering makeTethering() {
-            return mock(Tethering.class);
-        }
-
-        @Override
-        protected ProxyTracker makeProxyTracker() {
-            return mock(ProxyTracker.class);
-        }
-
-        @Override
-        protected int reserveNetId() {
-            while (true) {
-                final int netId = super.reserveNetId();
-
-                // Don't overlap test NetIDs with real NetIDs as binding sockets to real networks
-                // can have odd side-effects, like network validations succeeding.
-                Context context = InstrumentationRegistry.getContext();
-                final Network[] networks = ConnectivityManager.from(context).getAllNetworks();
-                boolean overlaps = false;
-                for (Network network : networks) {
-                    if (netId == network.netId) {
-                        overlaps = true;
-                        break;
-                    }
-                }
-                if (overlaps) continue;
-
-                return netId;
-            }
-        }
-
-        @Override
-        protected boolean queryUserAccess(int uid, int netId) {
-            return true;
-        }
-
-        public Nat464Xlat getNat464Xlat(MockNetworkAgent mna) {
-            return getNetworkAgentInfoForNetwork(mna.getNetwork()).clatd;
-        }
-
-        @Override
-        public MultinetworkPolicyTracker createMultinetworkPolicyTracker(
-                Context c, Handler h, Runnable r) {
-            final WrappedMultinetworkPolicyTracker tracker = new WrappedMultinetworkPolicyTracker(c, h, r);
-            return tracker;
-        }
-
-        public WrappedMultinetworkPolicyTracker getMultinetworkPolicyTracker() {
-            return (WrappedMultinetworkPolicyTracker) mMultinetworkPolicyTracker;
-        }
-
-        @Override
-        protected NetworkStackClient getNetworkStack() {
-            return mNetworkStack;
-        }
-
-        @Override
-        public WakeupMessage makeWakeupMessage(
-                Context context, Handler handler, String cmdName, int cmd, Object obj) {
-            return new FakeWakeupMessage(context, handler, cmdName, cmd, 0, 0, obj);
-        }
-
-        @Override
-        public boolean hasService(String name) {
-            // Currenty, the only relevant service that ConnectivityService checks for is
-            // ETHERNET_SERVICE.
-            return Context.ETHERNET_SERVICE.equals(name);
-        }
-
-        @Override
-        protected IpConnectivityMetrics.Logger metricsLogger() {
-            return mMetricsService;
-        }
-
-        @Override
-        protected void registerNetdEventCallback() {
-        }
-
-        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);
-        }
-
-        public void waitForIdle() {
-            waitForIdle(TIMEOUT_MS);
-        }
-
-        public void setUidRulesChanged(int uidRules) {
-            try {
-                mPolicyListener.onUidRulesChanged(Process.myUid(), uidRules);
-            } catch (RemoteException ignored) {
-            }
-        }
-
-        public void setRestrictBackgroundChanged(boolean restrictBackground) {
-            try {
-                mPolicyListener.onRestrictBackgroundChanged(restrictBackground);
-            } catch (RemoteException ignored) {
-            }
+            return mConfigMeteredMultipathPreference;
         }
     }
 
@@ -1300,13 +1004,22 @@
         LocalServices.addService(
                 NetworkPolicyManagerInternal.class, mock(NetworkPolicyManagerInternal.class));
 
-        mService = new WrappedConnectivityService(mServiceContext,
+        mAlarmManagerThread = new HandlerThread("TestAlarmManager");
+        mAlarmManagerThread.start();
+        initAlarmManager(mAlarmManager, mAlarmManagerThread.getThreadHandler());
+
+        mCsHandlerThread = new HandlerThread("TestConnectivityService");
+        final ConnectivityService.Dependencies deps = makeDependencies();
+        mService = new ConnectivityService(mServiceContext,
                 mNetworkManagementService,
                 mStatsService,
                 mNpm,
+                mMockDnsResolver,
                 mock(IpConnectivityLog.class),
                 mMockNetd,
-                mMockDnsResolver);
+                deps);
+        mService.mLingerDelayMs = TEST_LINGER_DELAY_MS;
+        verify(deps).makeMultinetworkPolicyTracker(any(), any(), any());
 
         final ArgumentCaptor<INetworkPolicyListener> policyListenerCaptor =
                 ArgumentCaptor.forClass(INetworkPolicyListener.class);
@@ -1317,7 +1030,7 @@
         // getSystemService() correctly.
         mCm = new WrappedConnectivityManager(InstrumentationRegistry.getContext(), mService);
         mService.systemReady();
-        mService.mockVpn(Process.myUid());
+        mockVpn(Process.myUid());
         mCm.bindProcessToNetwork(null);
 
         // Ensure that the default setting for Captive Portals is used for most tests
@@ -1326,6 +1039,57 @@
         setPrivateDnsSettings(PRIVATE_DNS_MODE_OFF, "ignored.example.com");
     }
 
+    private ConnectivityService.Dependencies makeDependencies() {
+        final MockableSystemProperties systemProperties = spy(new MockableSystemProperties());
+        when(systemProperties.getInt("net.tcp.default_init_rwnd", 0)).thenReturn(0);
+        when(systemProperties.getBoolean("ro.radio.noril", false)).thenReturn(false);
+
+        final ConnectivityService.Dependencies deps = mock(ConnectivityService.Dependencies.class);
+        doReturn(mCsHandlerThread).when(deps).makeHandlerThread();
+        doReturn(new TestNetIdManager()).when(deps).makeNetIdManager();
+        doReturn(mNetworkStack).when(deps).getNetworkStack();
+        doReturn(systemProperties).when(deps).getSystemProperties();
+        doReturn(mock(Tethering.class)).when(deps).makeTethering(any(), any(), any(), any(), any());
+        doReturn(mock(ProxyTracker.class)).when(deps).makeProxyTracker(any(), any());
+        doReturn(mMetricsService).when(deps).getMetricsLogger();
+        doReturn(true).when(deps).queryUserAccess(anyInt(), anyInt());
+        doReturn(mIpConnectivityMetrics).when(deps).getIpConnectivityMetrics();
+        doReturn(true).when(deps).hasService(Context.ETHERNET_SERVICE);
+        doAnswer(inv -> {
+            mPolicyTracker = new WrappedMultinetworkPolicyTracker(
+                    inv.getArgument(0), inv.getArgument(1), inv.getArgument(2));
+            return mPolicyTracker;
+        }).when(deps).makeMultinetworkPolicyTracker(any(), any(), any());
+
+        return deps;
+    }
+
+    private static void initAlarmManager(final AlarmManager am, final Handler alarmHandler) {
+        doAnswer(inv -> {
+            final long when = inv.getArgument(1);
+            final WakeupMessage wakeupMsg = inv.getArgument(3);
+            final Handler handler = inv.getArgument(4);
+
+            long delayMs = when - SystemClock.elapsedRealtime();
+            if (delayMs < 0) delayMs = 0;
+            if (delayMs > UNREASONABLY_LONG_ALARM_WAIT_MS) {
+                fail("Attempting to send msg more than " + UNREASONABLY_LONG_ALARM_WAIT_MS
+                        + "ms into the future: " + delayMs);
+            }
+            alarmHandler.postDelayed(() -> handler.post(wakeupMsg::onAlarm), wakeupMsg /* token */,
+                    delayMs);
+
+            return null;
+        }).when(am).setExact(eq(AlarmManager.ELAPSED_REALTIME_WAKEUP), anyLong(), anyString(),
+                any(WakeupMessage.class), any());
+
+        doAnswer(inv -> {
+            final WakeupMessage wakeupMsg = inv.getArgument(0);
+            alarmHandler.removeCallbacksAndMessages(wakeupMsg /* token */);
+            return null;
+        }).when(am).cancel(any(WakeupMessage.class));
+    }
+
     @After
     public void tearDown() throws Exception {
         setAlwaysOnNetworks(false);
@@ -1342,6 +1106,9 @@
             mEthernetNetworkAgent = null;
         }
         FakeSettingsProvider.clearSettingsProvider();
+
+        mCsHandlerThread.quitSafely();
+        mAlarmManagerThread.quitSafely();
     }
 
     private void mockDefaultPackages() throws Exception {
@@ -1361,21 +1128,6 @@
                 }));
     }
 
-   private static int transportToLegacyType(int transport) {
-        switch (transport) {
-            case TRANSPORT_ETHERNET:
-                return TYPE_ETHERNET;
-            case TRANSPORT_WIFI:
-                return TYPE_WIFI;
-            case TRANSPORT_CELLULAR:
-                return TYPE_MOBILE;
-            case TRANSPORT_VPN:
-                return TYPE_VPN;
-            default:
-                return TYPE_NONE;
-        }
-    }
-
     private void verifyActiveNetwork(int transport) {
         // Test getActiveNetworkInfo()
         assertNotNull(mCm.getActiveNetworkInfo());
@@ -1453,8 +1205,8 @@
     @Test
     public void testLingering() throws Exception {
         verifyNoNetwork();
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         assertNull(mCm.getActiveNetworkInfo());
         assertNull(mCm.getActiveNetwork());
         // Test bringing up validated cellular.
@@ -1478,7 +1230,7 @@
         assertTrue(mCm.getAllNetworks()[0].equals(mCellNetworkAgent.getNetwork()) ||
                 mCm.getAllNetworks()[1].equals(mCellNetworkAgent.getNetwork()));
         // Test cellular linger timeout.
-        waitFor(mCellNetworkAgent.getDisconnectedCV());
+        mCellNetworkAgent.expectDisconnected();
         waitForIdle();
         assertLength(1, mCm.getAllNetworks());
         verifyActiveNetwork(TRANSPORT_WIFI);
@@ -1494,13 +1246,13 @@
     @Test
     public void testValidatedCellularOutscoresUnvalidatedWiFi() throws Exception {
         // Test bringing up unvalidated WiFi
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         ConditionVariable cv = waitForConnectivityBroadcasts(1);
         mWiFiNetworkAgent.connect(false);
         waitFor(cv);
         verifyActiveNetwork(TRANSPORT_WIFI);
         // Test bringing up unvalidated cellular
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(false);
         waitForIdle();
         verifyActiveNetwork(TRANSPORT_WIFI);
@@ -1509,7 +1261,7 @@
         waitForIdle();
         verifyActiveNetwork(TRANSPORT_WIFI);
         // Test bringing up validated cellular
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         cv = waitForConnectivityBroadcasts(2);
         mCellNetworkAgent.connect(true);
         waitFor(cv);
@@ -1529,13 +1281,13 @@
     @Test
     public void testUnvalidatedWifiOutscoresUnvalidatedCellular() throws Exception {
         // Test bringing up unvalidated cellular.
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         ConditionVariable cv = waitForConnectivityBroadcasts(1);
         mCellNetworkAgent.connect(false);
         waitFor(cv);
         verifyActiveNetwork(TRANSPORT_CELLULAR);
         // Test bringing up unvalidated WiFi.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         cv = waitForConnectivityBroadcasts(2);
         mWiFiNetworkAgent.connect(false);
         waitFor(cv);
@@ -1555,7 +1307,7 @@
     @Test
     public void testUnlingeringDoesNotValidate() throws Exception {
         // Test bringing up unvalidated WiFi.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         ConditionVariable cv = waitForConnectivityBroadcasts(1);
         mWiFiNetworkAgent.connect(false);
         waitFor(cv);
@@ -1563,7 +1315,7 @@
         assertFalse(mCm.getNetworkCapabilities(mWiFiNetworkAgent.getNetwork()).hasCapability(
                 NET_CAPABILITY_VALIDATED));
         // Test bringing up validated cellular.
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         cv = waitForConnectivityBroadcasts(2);
         mCellNetworkAgent.connect(true);
         waitFor(cv);
@@ -1583,13 +1335,13 @@
     @Test
     public void testCellularOutscoresWeakWifi() throws Exception {
         // Test bringing up validated cellular.
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         ConditionVariable cv = waitForConnectivityBroadcasts(1);
         mCellNetworkAgent.connect(true);
         waitFor(cv);
         verifyActiveNetwork(TRANSPORT_CELLULAR);
         // Test bringing up validated WiFi.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         cv = waitForConnectivityBroadcasts(2);
         mWiFiNetworkAgent.connect(true);
         waitFor(cv);
@@ -1610,45 +1362,41 @@
     public void testReapingNetwork() throws Exception {
         // Test bringing up WiFi without NET_CAPABILITY_INTERNET.
         // Expect it to be torn down immediately because it satisfies no requests.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
-        ConditionVariable cv = mWiFiNetworkAgent.getDisconnectedCV();
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connectWithoutInternet();
-        waitFor(cv);
+        mWiFiNetworkAgent.expectDisconnected();
         // Test bringing up cellular without NET_CAPABILITY_INTERNET.
         // Expect it to be torn down immediately because it satisfies no requests.
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
-        cv = mCellNetworkAgent.getDisconnectedCV();
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mCellNetworkAgent.connectWithoutInternet();
-        waitFor(cv);
+        mCellNetworkAgent.expectDisconnected();
         // Test bringing up validated WiFi.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
-        cv = waitForConnectivityBroadcasts(1);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
+        final ConditionVariable cv = waitForConnectivityBroadcasts(1);
         mWiFiNetworkAgent.connect(true);
         waitFor(cv);
         verifyActiveNetwork(TRANSPORT_WIFI);
         // Test bringing up unvalidated cellular.
         // Expect it to be torn down because it could never be the highest scoring network
         // satisfying the default request even if it validated.
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
-        cv = mCellNetworkAgent.getDisconnectedCV();
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(false);
-        waitFor(cv);
+        mCellNetworkAgent.expectDisconnected();
         verifyActiveNetwork(TRANSPORT_WIFI);
-        cv = mWiFiNetworkAgent.getDisconnectedCV();
         mWiFiNetworkAgent.disconnect();
-        waitFor(cv);
+        mWiFiNetworkAgent.expectDisconnected();
     }
 
     @Test
     public void testCellularFallback() throws Exception {
         // Test bringing up validated cellular.
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         ConditionVariable cv = waitForConnectivityBroadcasts(1);
         mCellNetworkAgent.connect(true);
         waitFor(cv);
         verifyActiveNetwork(TRANSPORT_CELLULAR);
         // Test bringing up validated WiFi.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         cv = waitForConnectivityBroadcasts(2);
         mWiFiNetworkAgent.connect(true);
         waitFor(cv);
@@ -1680,13 +1428,13 @@
     @Test
     public void testWiFiFallback() throws Exception {
         // Test bringing up unvalidated WiFi.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         ConditionVariable cv = waitForConnectivityBroadcasts(1);
         mWiFiNetworkAgent.connect(false);
         waitFor(cv);
         verifyActiveNetwork(TRANSPORT_WIFI);
         // Test bringing up validated cellular.
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         cv = waitForConnectivityBroadcasts(2);
         mCellNetworkAgent.connect(true);
         waitFor(cv);
@@ -1709,274 +1457,33 @@
                 mCm.getDefaultRequest().networkCapabilities));
     }
 
-    enum CallbackState {
-        NONE,
-        AVAILABLE,
-        NETWORK_CAPABILITIES,
-        LINK_PROPERTIES,
-        SUSPENDED,
-        RESUMED,
-        LOSING,
-        LOST,
-        UNAVAILABLE,
-        BLOCKED_STATUS
-    }
-
-    private static class CallbackInfo {
-        public final CallbackState state;
-        public final Network network;
-        public final Object arg;
-        public CallbackInfo(CallbackState s, Network n, Object o) {
-            state = s; network = n; arg = o;
-        }
-        public String toString() {
-            return String.format("%s (%s) (%s)", state, network, arg);
-        }
-        @Override
-        public boolean equals(Object o) {
-            if (!(o instanceof CallbackInfo)) return false;
-            // Ignore timeMs, since it's unpredictable.
-            CallbackInfo other = (CallbackInfo) o;
-            return (state == other.state) && Objects.equals(network, other.network);
-        }
-        @Override
-        public int hashCode() {
-            return Objects.hash(state, network);
-        }
-    }
-
     /**
      * Utility NetworkCallback for testing. The caller must explicitly test for all the callbacks
      * this class receives, by calling expectCallback() exactly once each time a callback is
      * received. assertNoCallback may be called at any time.
      */
-    private class TestNetworkCallback extends NetworkCallback {
-        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));
+    private class TestNetworkCallback extends TestableNetworkCallback {
+        @Override
+        public void assertNoCallback() {
+            // TODO: better support this use case in TestableNetworkCallback
+            waitForIdle();
+            assertNoCallback(0 /* timeout */);
         }
 
         @Override
-        public void onAvailable(Network network) {
-            mLastAvailableNetwork = network;
-            setLastCallback(CallbackState.AVAILABLE, network, null);
-        }
-
-        @Override
-        public void onCapabilitiesChanged(Network network, NetworkCapabilities netCap) {
-            setLastCallback(CallbackState.NETWORK_CAPABILITIES, network, netCap);
-        }
-
-        @Override
-        public void onLinkPropertiesChanged(Network network, LinkProperties linkProp) {
-            setLastCallback(CallbackState.LINK_PROPERTIES, network, linkProp);
-        }
-
-        @Override
-        public void onUnavailable() {
-            setLastCallback(CallbackState.UNAVAILABLE, null, null);
-        }
-
-        @Override
-        public void onNetworkSuspended(Network network) {
-            setLastCallback(CallbackState.SUSPENDED, network, null);
-        }
-
-        @Override
-        public void onNetworkResumed(Network network) {
-            setLastCallback(CallbackState.RESUMED, network, null);
-        }
-
-        @Override
-        public void onLosing(Network network, int maxMsToLive) {
-            setLastCallback(CallbackState.LOSING, network, maxMsToLive /* autoboxed int */);
-        }
-
-        @Override
-        public void onLost(Network network) {
-            mLastAvailableNetwork = null;
-            setLastCallback(CallbackState.LOST, network, null);
-        }
-
-        @Override
-        public void onBlockedStatusChanged(Network network, boolean blocked) {
-            setLastCallback(CallbackState.BLOCKED_STATUS, network, blocked);
-        }
-
-        public Network getLastAvailableNetwork() {
-            return mLastAvailableNetwork;
-        }
-
-        CallbackInfo nextCallback(int timeoutMs) {
-            CallbackInfo cb = null;
-            try {
-                cb = mCallbacks.poll(timeoutMs, TimeUnit.MILLISECONDS);
-            } catch (InterruptedException e) {
-            }
-            if (cb == null) {
-                // LinkedBlockingQueue.poll() returns null if it timeouts.
-                fail("Did not receive callback after " + timeoutMs + "ms");
-            }
-            return cb;
-        }
-
-        CallbackInfo expectCallback(CallbackState state, MockNetworkAgent agent, int timeoutMs) {
-            final Network expectedNetwork = (agent != null) ? agent.getNetwork() : null;
-            CallbackInfo expected = new CallbackInfo(state, expectedNetwork, 0);
-            CallbackInfo actual = nextCallback(timeoutMs);
-            assertEquals("Unexpected callback:", expected, actual);
-
-            if (state == CallbackState.LOSING) {
+        public <T extends CallbackRecord> T expectCallback(final KClass<T> type, final HasNetwork n,
+                final long timeoutMs) {
+            final T callback = super.expectCallback(type, n, timeoutMs);
+            if (callback instanceof CallbackRecord.Losing) {
+                // TODO : move this to the specific test(s) needing this rather than here.
+                final CallbackRecord.Losing losing = (CallbackRecord.Losing) callback;
+                final int maxMsToLive = losing.getMaxMsToLive();
                 String msg = String.format(
                         "Invalid linger time value %d, must be between %d and %d",
-                        actual.arg, 0, mService.mLingerDelayMs);
-                int maxMsToLive = (Integer) actual.arg;
+                        maxMsToLive, 0, mService.mLingerDelayMs);
                 assertTrue(msg, 0 <= maxMsToLive && maxMsToLive <= mService.mLingerDelayMs);
             }
-
-            return actual;
-        }
-
-        CallbackInfo expectCallback(CallbackState state, MockNetworkAgent agent) {
-            return expectCallback(state, agent, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        CallbackInfo expectCallbackLike(Predicate<CallbackInfo> fn) {
-            return expectCallbackLike(fn, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        CallbackInfo expectCallbackLike(Predicate<CallbackInfo> fn, int timeoutMs) {
-            int timeLeft = timeoutMs;
-            while (timeLeft > 0) {
-                long start = SystemClock.elapsedRealtime();
-                CallbackInfo info = nextCallback(timeLeft);
-                if (fn.test(info)) {
-                    return info;
-                }
-                timeLeft -= (SystemClock.elapsedRealtime() - start);
-            }
-            fail("Did not receive expected callback after " + timeoutMs + "ms");
-            return null;
-        }
-
-        // Expects onAvailable and the callbacks that follow it. These are:
-        // - onSuspended, iff the network was suspended when the callbacks fire.
-        // - onCapabilitiesChanged.
-        // - onLinkPropertiesChanged.
-        // - onBlockedStatusChanged.
-        //
-        // @param agent the network to expect the callbacks on.
-        // @param expectSuspended whether to expect a SUSPENDED callback.
-        // @param expectValidated the expected value of the VALIDATED capability in the
-        //        onCapabilitiesChanged callback.
-        // @param timeoutMs how long to wait for the callbacks.
-        void expectAvailableCallbacks(MockNetworkAgent agent, boolean expectSuspended,
-                boolean expectValidated, boolean expectBlocked, int timeoutMs) {
-            expectCallback(CallbackState.AVAILABLE, agent, timeoutMs);
-            if (expectSuspended) {
-                expectCallback(CallbackState.SUSPENDED, agent, timeoutMs);
-            }
-            if (expectValidated) {
-                expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, agent, timeoutMs);
-            } else {
-                expectCapabilitiesWithout(NET_CAPABILITY_VALIDATED, agent, timeoutMs);
-            }
-            expectCallback(CallbackState.LINK_PROPERTIES, agent, timeoutMs);
-            expectBlockedStatusCallback(expectBlocked, agent);
-        }
-
-        // Expects the available callbacks (validated), plus onSuspended.
-        void expectAvailableAndSuspendedCallbacks(MockNetworkAgent agent, boolean expectValidated) {
-            expectAvailableCallbacks(agent, true, expectValidated, false, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        void expectAvailableCallbacksValidated(MockNetworkAgent agent) {
-            expectAvailableCallbacks(agent, false, true, false, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        void expectAvailableCallbacksValidatedAndBlocked(MockNetworkAgent agent) {
-            expectAvailableCallbacks(agent, false, true, true, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        void expectAvailableCallbacksUnvalidated(MockNetworkAgent agent) {
-            expectAvailableCallbacks(agent, false, false, false, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        void expectAvailableCallbacksUnvalidatedAndBlocked(MockNetworkAgent agent) {
-            expectAvailableCallbacks(agent, false, false, true, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        // Expects the available callbacks (where the onCapabilitiesChanged must contain the
-        // VALIDATED capability), plus another onCapabilitiesChanged which is identical to the
-        // one we just sent.
-        // TODO: this is likely a bug. Fix it and remove this method.
-        void expectAvailableDoubleValidatedCallbacks(MockNetworkAgent agent) {
-            expectCallback(CallbackState.AVAILABLE, agent, TEST_CALLBACK_TIMEOUT_MS);
-            NetworkCapabilities nc1 = expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, agent);
-            expectCallback(CallbackState.LINK_PROPERTIES, agent, TEST_CALLBACK_TIMEOUT_MS);
-            // Implicitly check the network is allowed to use.
-            // TODO: should we need to consider if network is in blocked status in this case?
-            expectBlockedStatusCallback(false, agent);
-            NetworkCapabilities nc2 = expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, agent);
-            assertEquals(nc1, nc2);
-        }
-
-        // Expects the available callbacks where the onCapabilitiesChanged must not have validated,
-        // then expects another onCapabilitiesChanged that has the validated bit set. This is used
-        // when a network connects and satisfies a callback, and then immediately validates.
-        void expectAvailableThenValidatedCallbacks(MockNetworkAgent agent) {
-            expectAvailableCallbacksUnvalidated(agent);
-            expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, agent);
-        }
-
-        NetworkCapabilities expectCapabilitiesWith(int capability, MockNetworkAgent agent) {
-            return expectCapabilitiesWith(capability, agent, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        NetworkCapabilities expectCapabilitiesWith(int capability, MockNetworkAgent agent,
-                int timeoutMs) {
-            CallbackInfo cbi = expectCallback(CallbackState.NETWORK_CAPABILITIES, agent, timeoutMs);
-            NetworkCapabilities nc = (NetworkCapabilities) cbi.arg;
-            assertTrue(nc.hasCapability(capability));
-            return nc;
-        }
-
-        NetworkCapabilities expectCapabilitiesWithout(int capability, MockNetworkAgent agent) {
-            return expectCapabilitiesWithout(capability, agent, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        NetworkCapabilities expectCapabilitiesWithout(int capability, MockNetworkAgent agent,
-                int timeoutMs) {
-            CallbackInfo cbi = expectCallback(CallbackState.NETWORK_CAPABILITIES, agent, timeoutMs);
-            NetworkCapabilities nc = (NetworkCapabilities) cbi.arg;
-            assertFalse(nc.hasCapability(capability));
-            return nc;
-        }
-
-        void expectCapabilitiesLike(Predicate<NetworkCapabilities> fn, MockNetworkAgent agent) {
-            CallbackInfo cbi = expectCallback(CallbackState.NETWORK_CAPABILITIES, agent);
-            assertTrue("Received capabilities don't match expectations : " + cbi.arg,
-                    fn.test((NetworkCapabilities) cbi.arg));
-        }
-
-        void expectLinkPropertiesLike(Predicate<LinkProperties> fn, MockNetworkAgent agent) {
-            CallbackInfo cbi = expectCallback(CallbackState.LINK_PROPERTIES, agent);
-            assertTrue("Received LinkProperties don't match expectations : " + cbi.arg,
-                    fn.test((LinkProperties) cbi.arg));
-        }
-
-        void expectBlockedStatusCallback(boolean expectBlocked, MockNetworkAgent agent) {
-            CallbackInfo cbi = expectCallback(CallbackState.BLOCKED_STATUS, agent);
-            boolean actualBlocked = (boolean) cbi.arg;
-            assertEquals(expectBlocked, actualBlocked);
-        }
-
-        void assertNoCallback() {
-            waitForIdle();
-            CallbackInfo c = mCallbacks.peek();
-            assertNull("Unexpected callback: " + c, c);
+            return callback;
         }
     }
 
@@ -2005,7 +1512,7 @@
 
         // Test unvalidated networks
         ConditionVariable cv = waitForConnectivityBroadcasts(1);
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(false);
         genericNetworkCallback.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
         cellNetworkCallback.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
@@ -2020,7 +1527,7 @@
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
 
         cv = waitForConnectivityBroadcasts(2);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
         genericNetworkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         wifiNetworkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
@@ -2030,21 +1537,21 @@
 
         cv = waitForConnectivityBroadcasts(2);
         mWiFiNetworkAgent.disconnect();
-        genericNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        wifiNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        wifiNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         cellNetworkCallback.assertNoCallback();
         waitFor(cv);
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
 
         cv = waitForConnectivityBroadcasts(1);
         mCellNetworkAgent.disconnect();
-        genericNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         waitFor(cv);
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
 
         // Test validated networks
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
         genericNetworkCallback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
         cellNetworkCallback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
@@ -2057,29 +1564,29 @@
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
         genericNetworkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        genericNetworkCallback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         genericNetworkCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         wifiNetworkCallback.expectAvailableThenValidatedCallbacks(mWiFiNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
 
         mWiFiNetworkAgent.disconnect();
-        genericNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        wifiNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        wifiNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
 
         mCellNetworkAgent.disconnect();
-        genericNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
     }
 
     @Test
-    public void testMultipleLingering() {
+    public void testMultipleLingering() throws Exception {
         // This test would be flaky with the default 120ms timer: that is short enough that
         // lingered networks are torn down before assertions can be run. We don't want to mock the
         // lingering timer to keep the WakeupMessage logic realistic: this has already proven useful
@@ -2095,9 +1602,9 @@
         TestNetworkCallback defaultCallback = new TestNetworkCallback();
         mCm.registerDefaultNetworkCallback(defaultCallback);
 
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
-        mEthernetNetworkAgent = new MockNetworkAgent(TRANSPORT_ETHERNET);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
+        mEthernetNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_ETHERNET);
 
         mCellNetworkAgent.addCapability(NET_CAPABILITY_NOT_METERED);
         mWiFiNetworkAgent.addCapability(NET_CAPABILITY_NOT_METERED);
@@ -2114,7 +1621,7 @@
         // We then get LOSING when wifi validates and cell is outscored.
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         // TODO: Investigate sending validated before losing.
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
@@ -2123,20 +1630,20 @@
         mEthernetNetworkAgent.connect(true);
         callback.expectAvailableCallbacksUnvalidated(mEthernetNetworkAgent);
         // TODO: Investigate sending validated before losing.
-        callback.expectCallback(CallbackState.LOSING, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mWiFiNetworkAgent);
         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);
+        callback.expectCallback(CallbackRecord.LOST, mEthernetNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mEthernetNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         for (int i = 0; i < 4; i++) {
-            MockNetworkAgent oldNetwork, newNetwork;
+            TestNetworkAgentWrapper oldNetwork, newNetwork;
             if (i % 2 == 0) {
                 mWiFiNetworkAgent.adjustScore(-15);
                 oldNetwork = mWiFiNetworkAgent;
@@ -2147,7 +1654,7 @@
                 newNetwork = mWiFiNetworkAgent;
 
             }
-            callback.expectCallback(CallbackState.LOSING, oldNetwork);
+            callback.expectCallback(CallbackRecord.LOSING, oldNetwork);
             // TODO: should we send an AVAILABLE callback to newNetwork, to indicate that it is no
             // longer lingering?
             defaultCallback.expectAvailableCallbacksValidated(newNetwork);
@@ -2161,7 +1668,7 @@
         // We expect a notification about the capabilities change, and nothing else.
         defaultCallback.expectCapabilitiesWithout(NET_CAPABILITY_NOT_METERED, mWiFiNetworkAgent);
         defaultCallback.assertNoCallback();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Wifi no longer satisfies our listen, which is for an unmetered network.
@@ -2170,11 +1677,11 @@
 
         // Disconnect our test networks.
         mWiFiNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
         mCellNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         waitForIdle();
         assertEquals(null, mCm.getActiveNetwork());
 
@@ -2188,7 +1695,7 @@
 
         mCm.registerNetworkCallback(request, callback);
 
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(false);   // Score: 10
         callback.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
@@ -2197,7 +1704,7 @@
 
         // Bring up wifi with a score of 20.
         // Cell stays up because it would satisfy the default request if it validated.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);   // Score: 20
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
@@ -2205,65 +1712,65 @@
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.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.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.adjustScore(50);
         mWiFiNetworkAgent.connect(false);   // Score: 70
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Tear down wifi.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.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.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         // TODO: Investigate sending validated before losing.
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         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);
-        defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
         mCellNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.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 = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
         callback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
         defaultCallback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         // TODO: Investigate sending validated before losing.
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
@@ -2274,13 +1781,13 @@
         // TODO: should this cause an AVAILABLE callback, to indicate that the network is no longer
         // lingering?
         mCm.unregisterNetworkCallback(noopCallback);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
 
         // Similar to the above: lingering can start even after the lingered request is removed.
         // Disconnect wifi and switch to cell.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
@@ -2289,7 +1796,7 @@
         mCm.requestNetwork(cellRequest, noopCallback);
 
         // Now connect wifi, and expect it to become the default network.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
         callback.expectAvailableThenValidatedCallbacks(mWiFiNetworkAgent);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
@@ -2299,34 +1806,34 @@
         callback.assertNoCallback();
         // Now unregister cellRequest and expect cell to start lingering.
         mCm.unregisterNetworkCallback(noopCallback);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
 
         // Let linger run its course.
         callback.assertNoCallback();
         final int lingerTimeoutMs = mService.mLingerDelayMs + mService.mLingerDelayMs / 4;
-        callback.expectCallback(CallbackState.LOST, mCellNetworkAgent, lingerTimeoutMs);
+        callback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent, lingerTimeoutMs);
 
         // Register a TRACK_DEFAULT request and check that it does not affect lingering.
         TestNetworkCallback trackDefaultCallback = new TestNetworkCallback();
         mCm.registerDefaultNetworkCallback(trackDefaultCallback);
         trackDefaultCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
-        mEthernetNetworkAgent = new MockNetworkAgent(TRANSPORT_ETHERNET);
+        mEthernetNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_ETHERNET);
         mEthernetNetworkAgent.connect(true);
         callback.expectAvailableCallbacksUnvalidated(mEthernetNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mWiFiNetworkAgent);
         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);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent, lingerTimeoutMs);
 
         // Clean up.
         mEthernetNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mEthernetNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.LOST, mEthernetNetworkAgent);
-        trackDefaultCallback.expectCallback(CallbackState.LOST, mEthernetNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mEthernetNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mEthernetNetworkAgent);
+        trackDefaultCallback.expectCallback(CallbackRecord.LOST, mEthernetNetworkAgent);
 
         mCm.unregisterNetworkCallback(callback);
         mCm.unregisterNetworkCallback(defaultCallback);
@@ -2334,7 +1841,7 @@
     }
 
     @Test
-    public void testNetworkGoesIntoBackgroundAfterLinger() {
+    public void testNetworkGoesIntoBackgroundAfterLinger() throws Exception {
         setAlwaysOnNetworks(true);
         NetworkRequest request = new NetworkRequest.Builder()
                 .clearCapabilities()
@@ -2345,8 +1852,8 @@
         TestNetworkCallback defaultCallback = new TestNetworkCallback();
         mCm.registerDefaultNetworkCallback(defaultCallback);
 
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
 
         mCellNetworkAgent.connect(true);
         callback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
@@ -2356,7 +1863,7 @@
         mWiFiNetworkAgent.connect(true);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
 
         // File a request for cellular, then release it.
@@ -2365,7 +1872,7 @@
         NetworkCallback noopCallback = new NetworkCallback();
         mCm.requestNetwork(cellRequest, noopCallback);
         mCm.unregisterNetworkCallback(noopCallback);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
 
         // Let linger run its course.
         callback.assertNoCallback();
@@ -2379,7 +1886,7 @@
     }
 
     @Test
-    public void testExplicitlySelected() {
+    public void testExplicitlySelected() throws Exception {
         NetworkRequest request = new NetworkRequest.Builder()
                 .clearCapabilities().addCapability(NET_CAPABILITY_INTERNET)
                 .build();
@@ -2387,13 +1894,13 @@
         mCm.registerNetworkCallback(request, callback);
 
         // Bring up validated cell.
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
         callback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
 
         // Bring up unvalidated wifi with explicitlySelected=true.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
-        mWiFiNetworkAgent.explicitlySelected(false);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
+        mWiFiNetworkAgent.explicitlySelected(true, false);
         mWiFiNetworkAgent.connect(false);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
 
@@ -2409,47 +1916,69 @@
         // If the user chooses yes on the "No Internet access, stay connected?" dialog, we switch to
         // wifi even though it's unvalidated.
         mCm.setAcceptUnvalidated(mWiFiNetworkAgent.getNetwork(), true, false);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
 
         // Disconnect wifi, and then reconnect, again with explicitlySelected=true.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
-        mWiFiNetworkAgent.explicitlySelected(false);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
+        mWiFiNetworkAgent.explicitlySelected(true, false);
         mWiFiNetworkAgent.connect(false);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
 
         // If the user chooses no on the "No Internet access, stay connected?" dialog, we ask the
         // network to disconnect.
         mCm.setAcceptUnvalidated(mWiFiNetworkAgent.getNetwork(), false, false);
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // Reconnect, again with explicitlySelected=true, but this time validate.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
-        mWiFiNetworkAgent.explicitlySelected(false);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
+        mWiFiNetworkAgent.explicitlySelected(true, false);
         mWiFiNetworkAgent.connect(true);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
 
         // BUG: the network will no longer linger, even though it's validated and outscored.
         // TODO: fix this.
-        mEthernetNetworkAgent = new MockNetworkAgent(TRANSPORT_ETHERNET);
+        mEthernetNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_ETHERNET);
         mEthernetNetworkAgent.connect(true);
         callback.expectAvailableThenValidatedCallbacks(mEthernetNetworkAgent);
         assertEquals(mEthernetNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         callback.assertNoCallback();
 
+        // Disconnect wifi, and then reconnect as if the user had selected "yes, don't ask again"
+        // (i.e., with explicitlySelected=true and acceptUnvalidated=true). Expect to switch to
+        // wifi immediately.
+        mWiFiNetworkAgent.disconnect();
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
+        mWiFiNetworkAgent.explicitlySelected(true, true);
+        mWiFiNetworkAgent.connect(false);
+        callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mEthernetNetworkAgent);
+        assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
+        mEthernetNetworkAgent.disconnect();
+        callback.expectCallback(CallbackRecord.LOST, mEthernetNetworkAgent);
+
+        // Disconnect and reconnect with explicitlySelected=false and acceptUnvalidated=true.
+        // Check that the network is not scored specially and that the device prefers cell data.
+        mWiFiNetworkAgent.disconnect();
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
+        mWiFiNetworkAgent.explicitlySelected(false, true);
+        mWiFiNetworkAgent.connect(false);
+        callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
+        assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
+
         // Clean up.
         mWiFiNetworkAgent.disconnect();
         mCellNetworkAgent.disconnect();
-        mEthernetNetworkAgent.disconnect();
 
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
-        callback.expectCallback(CallbackState.LOST, mEthernetNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
     }
 
     private int[] makeIntArray(final int size, final int value) {
@@ -2500,7 +2029,7 @@
         assertTrue(testFactory.getMyStartRequested());
 
         // Now bring in a higher scored network.
-        MockNetworkAgent testAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        TestNetworkAgentWrapper testAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         // Rather than create a validated network which complicates things by registering it's
         // own NetworkRequest during startup, just bump up the score to cancel out the
         // unvalidated penalty.
@@ -2585,27 +2114,26 @@
                 .build();
 
         Class<IllegalArgumentException> expected = IllegalArgumentException.class;
-        assertException(() -> { mCm.requestNetwork(request1, new NetworkCallback()); }, expected);
-        assertException(() -> { mCm.requestNetwork(request1, pendingIntent); }, expected);
-        assertException(() -> { mCm.requestNetwork(request2, new NetworkCallback()); }, expected);
-        assertException(() -> { mCm.requestNetwork(request2, pendingIntent); }, expected);
+        assertThrows(expected, () -> mCm.requestNetwork(request1, new NetworkCallback()));
+        assertThrows(expected, () -> mCm.requestNetwork(request1, pendingIntent));
+        assertThrows(expected, () -> mCm.requestNetwork(request2, new NetworkCallback()));
+        assertThrows(expected, () -> mCm.requestNetwork(request2, pendingIntent));
     }
 
     @Test
     public void testMMSonWiFi() throws Exception {
         // Test bringing up cellular without MMS NetworkRequest gets reaped
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.addCapability(NET_CAPABILITY_MMS);
-        ConditionVariable cv = mCellNetworkAgent.getDisconnectedCV();
         mCellNetworkAgent.connectWithoutInternet();
-        waitFor(cv);
+        mCellNetworkAgent.expectDisconnected();
         waitForIdle();
         assertEmpty(mCm.getAllNetworks());
         verifyNoNetwork();
 
         // Test bringing up validated WiFi.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
-        cv = waitForConnectivityBroadcasts(1);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
+        final ConditionVariable cv = waitForConnectivityBroadcasts(1);
         mWiFiNetworkAgent.connect(true);
         waitFor(cv);
         verifyActiveNetwork(TRANSPORT_WIFI);
@@ -2617,23 +2145,22 @@
         mCm.requestNetwork(builder.build(), networkCallback);
 
         // Test bringing up unvalidated cellular with MMS
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.addCapability(NET_CAPABILITY_MMS);
         mCellNetworkAgent.connectWithoutInternet();
         networkCallback.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
         verifyActiveNetwork(TRANSPORT_WIFI);
 
         // Test releasing NetworkRequest disconnects cellular with MMS
-        cv = mCellNetworkAgent.getDisconnectedCV();
         mCm.unregisterNetworkCallback(networkCallback);
-        waitFor(cv);
+        mCellNetworkAgent.expectDisconnected();
         verifyActiveNetwork(TRANSPORT_WIFI);
     }
 
     @Test
     public void testMMSonCell() throws Exception {
         // Test bringing up cellular without MMS
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         ConditionVariable cv = waitForConnectivityBroadcasts(1);
         mCellNetworkAgent.connect(false);
         waitFor(cv);
@@ -2646,21 +2173,21 @@
         mCm.requestNetwork(builder.build(), networkCallback);
 
         // Test bringing up MMS cellular network
-        MockNetworkAgent mmsNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        TestNetworkAgentWrapper
+                mmsNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mmsNetworkAgent.addCapability(NET_CAPABILITY_MMS);
         mmsNetworkAgent.connectWithoutInternet();
         networkCallback.expectAvailableCallbacksUnvalidated(mmsNetworkAgent);
         verifyActiveNetwork(TRANSPORT_CELLULAR);
 
         // Test releasing MMS NetworkRequest does not disconnect main cellular NetworkAgent
-        cv = mmsNetworkAgent.getDisconnectedCV();
         mCm.unregisterNetworkCallback(networkCallback);
-        waitFor(cv);
+        mmsNetworkAgent.expectDisconnected();
         verifyActiveNetwork(TRANSPORT_CELLULAR);
     }
 
     @Test
-    public void testPartialConnectivity() {
+    public void testPartialConnectivity() throws Exception {
         // Register network callback.
         NetworkRequest request = new NetworkRequest.Builder()
                 .clearCapabilities().addCapability(NET_CAPABILITY_INTERNET)
@@ -2669,12 +2196,12 @@
         mCm.registerNetworkCallback(request, callback);
 
         // Bring up validated mobile data.
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
         callback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
 
         // Bring up wifi with partial connectivity.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connectWithPartialConnectivity();
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_PARTIAL_CONNECTIVITY, mWiFiNetworkAgent);
@@ -2693,15 +2220,12 @@
         // If user accepts partial connectivity network,
         // NetworkMonitor#setAcceptPartialConnectivity() should be called too.
         waitForIdle();
-        try {
-            verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
-        } catch (RemoteException e) {
-            fail(e.getMessage());
-        }
+        verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
+
         // Need a trigger point to let NetworkMonitor tell ConnectivityService that network is
         // validated.
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), true);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         NetworkCapabilities nc = callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED,
                 mWiFiNetworkAgent);
         assertTrue(nc.hasCapability(NET_CAPABILITY_PARTIAL_CONNECTIVITY));
@@ -2709,8 +2233,8 @@
 
         // Disconnect and reconnect wifi with partial connectivity again.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connectWithPartialConnectivity();
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_PARTIAL_CONNECTIVITY, mWiFiNetworkAgent);
@@ -2721,69 +2245,79 @@
         // If the user chooses no, disconnect wifi immediately.
         mCm.setAcceptPartialConnectivity(mWiFiNetworkAgent.getNetwork(), false/* accept */,
                 false /* always */);
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // If user accepted partial connectivity before, and device reconnects to that network
         // again, but now the network has full connectivity. The network shouldn't contain
         // NET_CAPABILITY_PARTIAL_CONNECTIVITY.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         // acceptUnvalidated is also used as setting for accepting partial networks.
-        mWiFiNetworkAgent.explicitlySelected(true /* acceptUnvalidated */);
+        mWiFiNetworkAgent.explicitlySelected(true /* explicitlySelected */,
+                true /* acceptUnvalidated */);
         mWiFiNetworkAgent.connect(true);
+
         // If user accepted partial connectivity network before,
         // NetworkMonitor#setAcceptPartialConnectivity() will be called in
         // ConnectivityService#updateNetworkInfo().
         waitForIdle();
-        try {
-            verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
-        } catch (RemoteException e) {
-            fail(e.getMessage());
-        }
+        verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         nc = callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         assertFalse(nc.hasCapability(NET_CAPABILITY_PARTIAL_CONNECTIVITY));
+
         // Wifi should be the default network.
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
-        // If user accepted partial connectivity before, and now the device reconnects to the
-        // partial connectivity network. The network should be valid and contain
-        // NET_CAPABILITY_PARTIAL_CONNECTIVITY.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
-        mWiFiNetworkAgent.explicitlySelected(true /* acceptUnvalidated */);
-        // Current design cannot send multi-testResult from NetworkMonitor to ConnectivityService.
-        // So, if user accepts partial connectivity, NetworkMonitor will send PARTIAL_CONNECTIVITY
-        // to ConnectivityService first then send VALID. Once NetworkMonitor support
-        // multi-testResult, this test case also need to be changed to meet the new design.
+        // The user accepted partial connectivity and selected "don't ask again". Now the user
+        // reconnects to the partial connectivity network. Switch to wifi as soon as partial
+        // connectivity is detected.
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
+        mWiFiNetworkAgent.explicitlySelected(true /* explicitlySelected */,
+                true /* acceptUnvalidated */);
         mWiFiNetworkAgent.connectWithPartialConnectivity();
         // If user accepted partial connectivity network before,
         // NetworkMonitor#setAcceptPartialConnectivity() will be called in
         // ConnectivityService#updateNetworkInfo().
         waitForIdle();
-        try {
-            verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
-        } catch (RemoteException e) {
-            fail(e.getMessage());
-        }
+        verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
-        // TODO: If the user accepted partial connectivity, we shouldn't switch to wifi until
-        // NetworkMonitor detects partial connectivity
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         callback.expectCapabilitiesWith(NET_CAPABILITY_PARTIAL_CONNECTIVITY, mWiFiNetworkAgent);
         mWiFiNetworkAgent.setNetworkValid();
+
         // Need a trigger point to let NetworkMonitor tell ConnectivityService that network is
         // validated.
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), true);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+
+        // If the user accepted partial connectivity, and the device auto-reconnects to the partial
+        // connectivity network, it should contain both PARTIAL_CONNECTIVITY and VALIDATED.
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
+        mWiFiNetworkAgent.explicitlySelected(false /* explicitlySelected */,
+                true /* acceptUnvalidated */);
+
+        // NetworkMonitor will immediately (once the HTTPS probe fails...) report the network as
+        // valid, because ConnectivityService calls setAcceptPartialConnectivity before it calls
+        // notifyNetworkConnected.
+        mWiFiNetworkAgent.connectWithPartialValidConnectivity();
+        waitForIdle();
+        verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
+        callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
+        callback.expectCapabilitiesWith(
+                NET_CAPABILITY_PARTIAL_CONNECTIVITY | NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
+        mWiFiNetworkAgent.disconnect();
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
     }
 
     @Test
-    public void testCaptivePortalOnPartialConnectivity() throws RemoteException {
+    public void testCaptivePortalOnPartialConnectivity() throws Exception {
         final TestNetworkCallback captivePortalCallback = new TestNetworkCallback();
         final NetworkRequest captivePortalRequest = new NetworkRequest.Builder()
                 .addCapability(NET_CAPABILITY_CAPTIVE_PORTAL).build();
@@ -2796,7 +2330,7 @@
 
         // Bring up a network with a captive portal.
         // Expect onAvailable callback of listen for NET_CAPABILITY_CAPTIVE_PORTAL.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         String redirectUrl = "http://android.com/path";
         mWiFiNetworkAgent.connectWithCaptivePortal(redirectUrl);
         captivePortalCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
@@ -2821,7 +2355,7 @@
                 false /* always */);
         waitForIdle();
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), true);
-        captivePortalCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        captivePortalCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         validatedCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
         NetworkCapabilities nc =
                 validatedCallback.expectCapabilitiesWith(NET_CAPABILITY_PARTIAL_CONNECTIVITY,
@@ -2832,7 +2366,7 @@
     }
 
     @Test
-    public void testCaptivePortal() {
+    public void testCaptivePortal() throws Exception {
         final TestNetworkCallback captivePortalCallback = new TestNetworkCallback();
         final NetworkRequest captivePortalRequest = new NetworkRequest.Builder()
                 .addCapability(NET_CAPABILITY_CAPTIVE_PORTAL).build();
@@ -2845,7 +2379,7 @@
 
         // Bring up a network with a captive portal.
         // Expect onAvailable callback of listen for NET_CAPABILITY_CAPTIVE_PORTAL.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         String firstRedirectUrl = "http://example.com/firstPath";
         mWiFiNetworkAgent.connectWithCaptivePortal(firstRedirectUrl);
         captivePortalCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
@@ -2854,11 +2388,11 @@
         // Take down network.
         // Expect onLost callback.
         mWiFiNetworkAgent.disconnect();
-        captivePortalCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        captivePortalCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // Bring up a network with a captive portal.
         // Expect onAvailable callback of listen for NET_CAPABILITY_CAPTIVE_PORTAL.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         String secondRedirectUrl = "http://example.com/secondPath";
         mWiFiNetworkAgent.connectWithCaptivePortal(secondRedirectUrl);
         captivePortalCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
@@ -2868,7 +2402,7 @@
         // Expect onLost callback because network no longer provides NET_CAPABILITY_CAPTIVE_PORTAL.
         mWiFiNetworkAgent.setNetworkValid();
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), true);
-        captivePortalCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        captivePortalCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // Expect NET_CAPABILITY_VALIDATED onAvailable callback.
         validatedCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
@@ -2880,11 +2414,11 @@
         // Expect NET_CAPABILITY_VALIDATED onLost callback.
         mWiFiNetworkAgent.setNetworkInvalid();
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), false);
-        validatedCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        validatedCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
     }
 
     @Test
-    public void testCaptivePortalApp() throws RemoteException {
+    public void testCaptivePortalApp() throws Exception {
         final TestNetworkCallback captivePortalCallback = new TestNetworkCallback();
         final NetworkRequest captivePortalRequest = new NetworkRequest.Builder()
                 .addCapability(NET_CAPABILITY_CAPTIVE_PORTAL).build();
@@ -2896,7 +2430,7 @@
         mCm.registerNetworkCallback(validatedRequest, validatedCallback);
 
         // Bring up wifi.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
         validatedCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
         Network wifiNetwork = mWiFiNetworkAgent.getNetwork();
@@ -2912,7 +2446,7 @@
         mWiFiNetworkAgent.setNetworkPortal("http://example.com");
         mCm.reportNetworkConnectivity(wifiNetwork, false);
         captivePortalCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        validatedCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        validatedCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // Check that startCaptivePortalApp sends the expected command to NetworkMonitor.
         mCm.startCaptivePortalApp(wifiNetwork);
@@ -2933,7 +2467,7 @@
         mWiFiNetworkAgent.setNetworkValid();
         mWiFiNetworkAgent.mNetworkMonitor.forceReevaluation(Process.myUid());
         validatedCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
-        captivePortalCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        captivePortalCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         verify(mNotificationManager, times(1)).notifyAsUser(anyString(),
                 eq(NotificationType.LOGGED_IN.eventId), any(), eq(UserHandle.ALL));
 
@@ -2942,7 +2476,7 @@
     }
 
     @Test
-    public void testAvoidOrIgnoreCaptivePortals() {
+    public void testAvoidOrIgnoreCaptivePortals() throws Exception {
         final TestNetworkCallback captivePortalCallback = new TestNetworkCallback();
         final NetworkRequest captivePortalRequest = new NetworkRequest.Builder()
                 .addCapability(NET_CAPABILITY_CAPTIVE_PORTAL).build();
@@ -2956,14 +2490,12 @@
         setCaptivePortalMode(Settings.Global.CAPTIVE_PORTAL_MODE_AVOID);
         // Bring up a network with a captive portal.
         // Expect it to fail to connect and not result in any callbacks.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         String firstRedirectUrl = "http://example.com/firstPath";
 
-        ConditionVariable disconnectCv = mWiFiNetworkAgent.getDisconnectedCV();
-        ConditionVariable avoidCv = mWiFiNetworkAgent.getPreventReconnectReceived();
         mWiFiNetworkAgent.connectWithCaptivePortal(firstRedirectUrl);
-        waitFor(disconnectCv);
-        waitFor(avoidCv);
+        mWiFiNetworkAgent.expectDisconnected();
+        mWiFiNetworkAgent.expectPreventReconnectReceived();
 
         assertNoCallbacks(captivePortalCallback, validatedCallback);
     }
@@ -2982,7 +2514,7 @@
      * does work.
      */
     @Test
-    public void testNetworkSpecifier() {
+    public void testNetworkSpecifier() throws Exception {
         // A NetworkSpecifier subclass that matches all networks but must not be visible to apps.
         class ConfidentialMatchAllNetworkSpecifier extends NetworkSpecifier implements
                 Parcelable {
@@ -3062,7 +2594,7 @@
         LocalStringNetworkSpecifier nsFoo = new LocalStringNetworkSpecifier("foo");
         LocalStringNetworkSpecifier nsBar = new LocalStringNetworkSpecifier("bar");
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
         cEmpty1.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         cEmpty2.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
@@ -3073,24 +2605,24 @@
         mWiFiNetworkAgent.setNetworkSpecifier(nsFoo);
         cFoo.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         for (TestNetworkCallback c: emptyCallbacks) {
-            c.expectCapabilitiesLike((caps) -> caps.getNetworkSpecifier().equals(nsFoo),
-                    mWiFiNetworkAgent);
+            c.expectCapabilitiesThat(mWiFiNetworkAgent,
+                    (caps) -> caps.getNetworkSpecifier().equals(nsFoo));
         }
-        cFoo.expectCapabilitiesLike((caps) -> caps.getNetworkSpecifier().equals(nsFoo),
-                mWiFiNetworkAgent);
+        cFoo.expectCapabilitiesThat(mWiFiNetworkAgent,
+                (caps) -> caps.getNetworkSpecifier().equals(nsFoo));
         assertEquals(nsFoo,
                 mCm.getNetworkCapabilities(mWiFiNetworkAgent.getNetwork()).getNetworkSpecifier());
         cFoo.assertNoCallback();
 
         mWiFiNetworkAgent.setNetworkSpecifier(nsBar);
-        cFoo.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        cFoo.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         cBar.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         for (TestNetworkCallback c: emptyCallbacks) {
-            c.expectCapabilitiesLike((caps) -> caps.getNetworkSpecifier().equals(nsBar),
-                    mWiFiNetworkAgent);
+            c.expectCapabilitiesThat(mWiFiNetworkAgent,
+                    (caps) -> caps.getNetworkSpecifier().equals(nsBar));
         }
-        cBar.expectCapabilitiesLike((caps) -> caps.getNetworkSpecifier().equals(nsBar),
-                mWiFiNetworkAgent);
+        cBar.expectCapabilitiesThat(mWiFiNetworkAgent,
+                (caps) -> caps.getNetworkSpecifier().equals(nsBar));
         assertEquals(nsBar,
                 mCm.getNetworkCapabilities(mWiFiNetworkAgent.getNetwork()).getNetworkSpecifier());
         cBar.assertNoCallback();
@@ -3098,23 +2630,23 @@
         mWiFiNetworkAgent.setNetworkSpecifier(new ConfidentialMatchAllNetworkSpecifier());
         cFoo.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         for (TestNetworkCallback c : emptyCallbacks) {
-            c.expectCapabilitiesLike((caps) -> caps.getNetworkSpecifier() == null,
-                    mWiFiNetworkAgent);
+            c.expectCapabilitiesThat(mWiFiNetworkAgent,
+                    (caps) -> caps.getNetworkSpecifier() == null);
         }
-        cFoo.expectCapabilitiesLike((caps) -> caps.getNetworkSpecifier() == null,
-                mWiFiNetworkAgent);
-        cBar.expectCapabilitiesLike((caps) -> caps.getNetworkSpecifier() == null,
-                mWiFiNetworkAgent);
+        cFoo.expectCapabilitiesThat(mWiFiNetworkAgent,
+                (caps) -> caps.getNetworkSpecifier() == null);
+        cBar.expectCapabilitiesThat(mWiFiNetworkAgent,
+                (caps) -> caps.getNetworkSpecifier() == null);
         assertNull(
                 mCm.getNetworkCapabilities(mWiFiNetworkAgent.getNetwork()).getNetworkSpecifier());
         cFoo.assertNoCallback();
         cBar.assertNoCallback();
 
         mWiFiNetworkAgent.setNetworkSpecifier(null);
-        cFoo.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        cBar.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        cFoo.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        cBar.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         for (TestNetworkCallback c: emptyCallbacks) {
-            c.expectCallback(CallbackState.NETWORK_CAPABILITIES, mWiFiNetworkAgent);
+            c.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED, mWiFiNetworkAgent);
         }
 
         assertNoCallbacks(cEmpty1, cEmpty2, cEmpty3, cEmpty4, cFoo, cBar);
@@ -3122,24 +2654,18 @@
 
     @Test
     public void testInvalidNetworkSpecifier() {
-        try {
+        assertThrows(IllegalArgumentException.class, () -> {
             NetworkRequest.Builder builder = new NetworkRequest.Builder();
             builder.setNetworkSpecifier(new MatchAllNetworkSpecifier());
-            fail("NetworkRequest builder with MatchAllNetworkSpecifier");
-        } catch (IllegalArgumentException expected) {
-            // expected
-        }
+        });
 
-        try {
+        assertThrows(IllegalArgumentException.class, () -> {
             NetworkCapabilities networkCapabilities = new NetworkCapabilities();
             networkCapabilities.addTransportType(TRANSPORT_WIFI)
                     .setNetworkSpecifier(new MatchAllNetworkSpecifier());
             mService.requestNetwork(networkCapabilities, null, 0, null,
                     ConnectivityManager.TYPE_WIFI);
-            fail("ConnectivityService requestNetwork with MatchAllNetworkSpecifier");
-        } catch (IllegalArgumentException expected) {
-            // expected
-        }
+        });
 
         class NonParcelableSpecifier extends NetworkSpecifier {
             public boolean satisfiedBy(NetworkSpecifier other) { return false; }
@@ -3148,24 +2674,22 @@
             @Override public int describeContents() { return 0; }
             @Override public void writeToParcel(Parcel p, int flags) {}
         }
-        NetworkRequest.Builder builder;
 
-        builder = new NetworkRequest.Builder().addTransportType(TRANSPORT_ETHERNET);
-        try {
+        final NetworkRequest.Builder builder =
+                new NetworkRequest.Builder().addTransportType(TRANSPORT_ETHERNET);
+        assertThrows(ClassCastException.class, () -> {
             builder.setNetworkSpecifier(new NonParcelableSpecifier());
             Parcel parcelW = Parcel.obtain();
             builder.build().writeToParcel(parcelW, 0);
-            fail("Parceling a non-parcelable specifier did not throw an exception");
-        } catch (Exception e) {
-            // expected
-        }
+        });
 
-        builder = new NetworkRequest.Builder().addTransportType(TRANSPORT_ETHERNET);
-        builder.setNetworkSpecifier(new ParcelableSpecifier());
-        NetworkRequest nr = builder.build();
+        final NetworkRequest nr =
+                new NetworkRequest.Builder().addTransportType(TRANSPORT_ETHERNET)
+                .setNetworkSpecifier(new ParcelableSpecifier())
+                .build();
         assertNotNull(nr);
 
-        try {
+        assertThrows(BadParcelableException.class, () -> {
             Parcel parcelW = Parcel.obtain();
             nr.writeToParcel(parcelW, 0);
             byte[] bytes = parcelW.marshall();
@@ -3175,14 +2699,11 @@
             parcelR.unmarshall(bytes, 0, bytes.length);
             parcelR.setDataPosition(0);
             NetworkRequest rereadNr = NetworkRequest.CREATOR.createFromParcel(parcelR);
-            fail("Unparceling a non-framework NetworkSpecifier did not throw an exception");
-        } catch (Exception e) {
-            // expected
-        }
+        });
     }
 
     @Test
-    public void testNetworkSpecifierUidSpoofSecurityException() {
+    public void testNetworkSpecifierUidSpoofSecurityException() throws Exception {
         class UidAwareNetworkSpecifier extends NetworkSpecifier implements Parcelable {
             @Override
             public boolean satisfiedBy(NetworkSpecifier other) {
@@ -3200,19 +2721,16 @@
             public void writeToParcel(Parcel dest, int flags) {}
         }
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
 
         UidAwareNetworkSpecifier networkSpecifier = new UidAwareNetworkSpecifier();
         NetworkRequest networkRequest = newWifiRequestBuilder().setNetworkSpecifier(
                 networkSpecifier).build();
         TestNetworkCallback networkCallback = new TestNetworkCallback();
-        try {
+        assertThrows(SecurityException.class, () -> {
             mCm.requestNetwork(networkRequest, networkCallback);
-            fail("Network request with spoofed UID did not throw a SecurityException");
-        } catch (SecurityException e) {
-            // expected
-        }
+        });
     }
 
     @Test
@@ -3224,36 +2742,20 @@
                 .build();
         // Registering a NetworkCallback with signal strength but w/o NETWORK_SIGNAL_STRENGTH_WAKEUP
         // permission should get SecurityException.
-        try {
-            mCm.registerNetworkCallback(r, new NetworkCallback());
-            fail("Expected SecurityException filing a callback with signal strength");
-        } catch (SecurityException expected) {
-            // expected
-        }
+        assertThrows(SecurityException.class, () ->
+                mCm.registerNetworkCallback(r, new NetworkCallback()));
 
-        try {
-            mCm.registerNetworkCallback(r, PendingIntent.getService(
-                    mServiceContext, 0, new Intent(), 0));
-            fail("Expected SecurityException filing a callback with signal strength");
-        } catch (SecurityException expected) {
-            // expected
-        }
+        assertThrows(SecurityException.class, () ->
+                mCm.registerNetworkCallback(r, PendingIntent.getService(
+                        mServiceContext, 0, new Intent(), 0)));
 
         // Requesting a Network with signal strength should get IllegalArgumentException.
-        try {
-            mCm.requestNetwork(r, new NetworkCallback());
-            fail("Expected IllegalArgumentException filing a request with signal strength");
-        } catch (IllegalArgumentException expected) {
-            // expected
-        }
+        assertThrows(IllegalArgumentException.class, () ->
+                mCm.requestNetwork(r, new NetworkCallback()));
 
-        try {
-            mCm.requestNetwork(r, PendingIntent.getService(
-                    mServiceContext, 0, new Intent(), 0));
-            fail("Expected IllegalArgumentException filing a request with signal strength");
-        } catch (IllegalArgumentException expected) {
-            // expected
-        }
+        assertThrows(IllegalArgumentException.class, () ->
+                mCm.requestNetwork(r, PendingIntent.getService(
+                        mServiceContext, 0, new Intent(), 0)));
     }
 
     @Test
@@ -3272,14 +2774,14 @@
         cellNetworkCallback.assertNoCallback();
 
         // Bring up cell and expect CALLBACK_AVAILABLE.
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         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 = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
         cellNetworkCallback.assertNoCallback();
         defaultNetworkCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
@@ -3287,12 +2789,12 @@
 
         // Bring down cell. Expect no default network callback, since it wasn't the default.
         mCellNetworkAgent.disconnect();
-        cellNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.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 = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
         cellNetworkCallback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
         defaultNetworkCallback.assertNoCallback();
@@ -3302,16 +2804,17 @@
         // followed by AVAILABLE cell.
         mWiFiNetworkAgent.disconnect();
         cellNetworkCallback.assertNoCallback();
-        defaultNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        defaultNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         defaultNetworkCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         mCellNetworkAgent.disconnect();
-        cellNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
-        defaultNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
+        defaultNetworkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         waitForIdle();
         assertEquals(null, mCm.getActiveNetwork());
 
         final int uid = Process.myUid();
-        final MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        final TestNetworkAgentWrapper
+                vpnNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN);
         final ArraySet<UidRange> ranges = new ArraySet<>();
         ranges.add(new UidRange(uid, uid));
         mMockVpn.setNetworkAgent(vpnNetworkAgent);
@@ -3322,7 +2825,7 @@
         assertEquals(defaultNetworkCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         vpnNetworkAgent.disconnect();
-        defaultNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        defaultNetworkCallback.expectCallback(CallbackRecord.LOST, vpnNetworkAgent);
         waitForIdle();
         assertEquals(null, mCm.getActiveNetwork());
     }
@@ -3336,7 +2839,7 @@
         mCm.requestNetwork(cellRequest, cellNetworkCallback);
 
         // Bring up the mobile network.
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
 
         // We should get onAvailable(), onCapabilitiesChanged(), and
@@ -3350,14 +2853,15 @@
         lp.setInterfaceName("foonet_data0");
         mCellNetworkAgent.sendLinkProperties(lp);
         // We should get onLinkPropertiesChanged().
-        cellNetworkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED,
+                mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
 
         // Suspend the network.
         mCellNetworkAgent.suspend();
         cellNetworkCallback.expectCapabilitiesWithout(NET_CAPABILITY_NOT_SUSPENDED,
                 mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.SUSPENDED, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.SUSPENDED, mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
 
         // Register a garden variety default network request.
@@ -3372,7 +2876,7 @@
         mCellNetworkAgent.resume();
         cellNetworkCallback.expectCapabilitiesWith(NET_CAPABILITY_NOT_SUSPENDED,
                 mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.RESUMED, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.RESUMED, mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
 
         dfltNetworkCallback = new TestNetworkCallback();
@@ -3405,7 +2909,7 @@
         waitForIdle();
     }
 
-    private boolean isForegroundNetwork(MockNetworkAgent network) {
+    private boolean isForegroundNetwork(TestNetworkAgentWrapper network) {
         NetworkCapabilities nc = mCm.getNetworkCapabilities(network.getNetwork());
         assertNotNull(nc);
         return nc.hasCapability(NET_CAPABILITY_FOREGROUND);
@@ -3424,21 +2928,21 @@
         mCm.registerNetworkCallback(request, callback);
         mCm.registerNetworkCallback(fgRequest, fgCallback);
 
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
         callback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
         fgCallback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
         assertTrue(isForegroundNetwork(mCellNetworkAgent));
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
 
         // When wifi connects, cell lingers.
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         fgCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        fgCallback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        fgCallback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         fgCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         assertTrue(isForegroundNetwork(mCellNetworkAgent));
         assertTrue(isForegroundNetwork(mWiFiNetworkAgent));
@@ -3446,7 +2950,7 @@
         // When lingering is complete, cell is still there but is now in the background.
         waitForIdle();
         int timeoutMs = TEST_LINGER_DELAY_MS + TEST_LINGER_DELAY_MS / 4;
-        fgCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent, timeoutMs);
+        fgCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent, timeoutMs);
         // Expect a network capabilities update sans FOREGROUND.
         callback.expectCapabilitiesWithout(NET_CAPABILITY_FOREGROUND, mCellNetworkAgent);
         assertFalse(isForegroundNetwork(mCellNetworkAgent));
@@ -3472,7 +2976,7 @@
         // Release the request. The network immediately goes into the background, since it was not
         // lingering.
         mCm.unregisterNetworkCallback(cellCallback);
-        fgCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        fgCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         // Expect a network capabilities update sans FOREGROUND.
         callback.expectCapabilitiesWithout(NET_CAPABILITY_FOREGROUND, mCellNetworkAgent);
         assertFalse(isForegroundNetwork(mCellNetworkAgent));
@@ -3480,8 +2984,8 @@
 
         // Disconnect wifi and check that cell is foreground again.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        fgCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        fgCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         fgCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         assertTrue(isForegroundNetwork(mCellNetworkAgent));
 
@@ -3517,20 +3021,20 @@
             };
         }
 
-        assertTimeLimit("Registering callbacks", REGISTER_TIME_LIMIT_MS, () -> {
+        assertRunsInAtMost("Registering callbacks", REGISTER_TIME_LIMIT_MS, () -> {
             for (NetworkCallback cb : callbacks) {
                 mCm.registerNetworkCallback(request, cb);
             }
         });
 
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         // Don't request that the network validate, because otherwise connect() will block until
         // the network gets NET_CAPABILITY_VALIDATED, after all the callbacks below have fired,
         // and we won't actually measure anything.
         mCellNetworkAgent.connect(false);
 
         long onAvailableDispatchingDuration = durationOf(() -> {
-            awaitLatch(availableLatch, 10 * CONNECT_TIME_LIMIT_MS);
+            await(availableLatch, 10 * CONNECT_TIME_LIMIT_MS);
         });
         Log.d(TAG, String.format("Dispatched %d of %d onAvailable callbacks in %dms",
                 NUM_REQUESTS - availableLatch.getCount(), NUM_REQUESTS,
@@ -3540,12 +3044,12 @@
                 onAvailableDispatchingDuration <= CONNECT_TIME_LIMIT_MS);
 
         // Give wifi a high enough score that we'll linger cell when wifi comes up.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.adjustScore(40);
         mWiFiNetworkAgent.connect(false);
 
         long onLostDispatchingDuration = durationOf(() -> {
-            awaitLatch(losingLatch, 10 * SWITCH_TIME_LIMIT_MS);
+            await(losingLatch, 10 * SWITCH_TIME_LIMIT_MS);
         });
         Log.d(TAG, String.format("Dispatched %d of %d onLosing callbacks in %dms",
                 NUM_REQUESTS - losingLatch.getCount(), NUM_REQUESTS, onLostDispatchingDuration));
@@ -3553,33 +3057,13 @@
                 NUM_REQUESTS, onLostDispatchingDuration, SWITCH_TIME_LIMIT_MS),
                 onLostDispatchingDuration <= SWITCH_TIME_LIMIT_MS);
 
-        assertTimeLimit("Unregistering callbacks", UNREGISTER_TIME_LIMIT_MS, () -> {
+        assertRunsInAtMost("Unregistering callbacks", UNREGISTER_TIME_LIMIT_MS, () -> {
             for (NetworkCallback cb : callbacks) {
                 mCm.unregisterNetworkCallback(cb);
             }
         });
     }
 
-    private long durationOf(Runnable fn) {
-        long startTime = SystemClock.elapsedRealtime();
-        fn.run();
-        return SystemClock.elapsedRealtime() - startTime;
-    }
-
-    private void assertTimeLimit(String descr, long timeLimit, Runnable fn) {
-        long timeTaken = durationOf(fn);
-        String msg = String.format("%s: took %dms, limit was %dms", descr, timeTaken, timeLimit);
-        Log.d(TAG, msg);
-        assertTrue(msg, timeTaken <= timeLimit);
-    }
-
-    private boolean awaitLatch(CountDownLatch l, long timeoutMs) {
-        try {
-            return l.await(timeoutMs, TimeUnit.MILLISECONDS);
-        } catch (InterruptedException e) {}
-        return false;
-    }
-
     @Test
     public void testMobileDataAlwaysOn() throws Exception {
         final TestNetworkCallback cellNetworkCallback = new TestNetworkCallback();
@@ -3603,7 +3087,7 @@
         assertTrue(testFactory.getMyStartRequested());
 
         // Bring up wifi. The factory stops looking for a network.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         // Score 60 - 40 penalty for not validated yet, then 60 when it validates
         testFactory.expectAddRequestsWithScores(20, 60);
         mWiFiNetworkAgent.connect(true);
@@ -3620,7 +3104,7 @@
 
         // Bring up cell data and check that the factory stops looking.
         assertLength(1, mCm.getAllNetworks());
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         testFactory.expectAddRequestsWithScores(10, 50);  // Unvalidated, then validated
         mCellNetworkAgent.connect(true);
         cellNetworkCallback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
@@ -3638,7 +3122,7 @@
         testFactory.waitForNetworkRequests(1);
 
         // ...  and cell data to be torn down.
-        cellNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         assertLength(1, mCm.getAllNetworks());
 
         testFactory.unregister();
@@ -3649,48 +3133,46 @@
     @Test
     public void testAvoidBadWifiSetting() throws Exception {
         final ContentResolver cr = mServiceContext.getContentResolver();
-        final WrappedMultinetworkPolicyTracker tracker = mService.getMultinetworkPolicyTracker();
         final String settingName = Settings.Global.NETWORK_AVOID_BAD_WIFI;
 
-        tracker.configRestrictsAvoidBadWifi = false;
+        mPolicyTracker.mConfigRestrictsAvoidBadWifi = false;
         String[] values = new String[] {null, "0", "1"};
         for (int i = 0; i < values.length; i++) {
             Settings.Global.putInt(cr, settingName, 1);
-            tracker.reevaluate();
+            mPolicyTracker.reevaluate();
             waitForIdle();
             String msg = String.format("config=false, setting=%s", values[i]);
             assertTrue(mService.avoidBadWifi());
-            assertFalse(msg, tracker.shouldNotifyWifiUnvalidated());
+            assertFalse(msg, mPolicyTracker.shouldNotifyWifiUnvalidated());
         }
 
-        tracker.configRestrictsAvoidBadWifi = true;
+        mPolicyTracker.mConfigRestrictsAvoidBadWifi = true;
 
         Settings.Global.putInt(cr, settingName, 0);
-        tracker.reevaluate();
+        mPolicyTracker.reevaluate();
         waitForIdle();
         assertFalse(mService.avoidBadWifi());
-        assertFalse(tracker.shouldNotifyWifiUnvalidated());
+        assertFalse(mPolicyTracker.shouldNotifyWifiUnvalidated());
 
         Settings.Global.putInt(cr, settingName, 1);
-        tracker.reevaluate();
+        mPolicyTracker.reevaluate();
         waitForIdle();
         assertTrue(mService.avoidBadWifi());
-        assertFalse(tracker.shouldNotifyWifiUnvalidated());
+        assertFalse(mPolicyTracker.shouldNotifyWifiUnvalidated());
 
         Settings.Global.putString(cr, settingName, null);
-        tracker.reevaluate();
+        mPolicyTracker.reevaluate();
         waitForIdle();
         assertFalse(mService.avoidBadWifi());
-        assertTrue(tracker.shouldNotifyWifiUnvalidated());
+        assertTrue(mPolicyTracker.shouldNotifyWifiUnvalidated());
     }
 
     @Test
     public void testAvoidBadWifi() throws Exception {
         final ContentResolver cr = mServiceContext.getContentResolver();
-        final WrappedMultinetworkPolicyTracker tracker = mService.getMultinetworkPolicyTracker();
 
         // Pretend we're on a carrier that restricts switching away from bad wifi.
-        tracker.configRestrictsAvoidBadWifi = true;
+        mPolicyTracker.mConfigRestrictsAvoidBadWifi = true;
 
         // File a request for cell to ensure it doesn't go down.
         final TestNetworkCallback cellNetworkCallback = new TestNetworkCallback();
@@ -3709,17 +3191,17 @@
         mCm.registerNetworkCallback(validatedWifiRequest, validatedWifiCallback);
 
         Settings.Global.putInt(cr, Settings.Global.NETWORK_AVOID_BAD_WIFI, 0);
-        tracker.reevaluate();
+        mPolicyTracker.reevaluate();
 
         // Bring up validated cell.
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
         cellNetworkCallback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
         defaultCallback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
         Network cellNetwork = mCellNetworkAgent.getNetwork();
 
         // Bring up validated wifi.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
         validatedWifiCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
@@ -3729,7 +3211,7 @@
         mWiFiNetworkAgent.setNetworkInvalid();
         mCm.reportNetworkConnectivity(wifiNetwork, false);
         defaultCallback.expectCapabilitiesWithout(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
-        validatedWifiCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        validatedWifiCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // Because avoid bad wifi is off, we don't switch to cellular.
         defaultCallback.assertNoCallback();
@@ -3741,14 +3223,14 @@
 
         // Simulate switching to a carrier that does not restrict avoiding bad wifi, and expect
         // that we switch back to cell.
-        tracker.configRestrictsAvoidBadWifi = false;
-        tracker.reevaluate();
+        mPolicyTracker.mConfigRestrictsAvoidBadWifi = false;
+        mPolicyTracker.reevaluate();
         defaultCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         assertEquals(mCm.getActiveNetwork(), cellNetwork);
 
         // Switch back to a restrictive carrier.
-        tracker.configRestrictsAvoidBadWifi = true;
-        tracker.reevaluate();
+        mPolicyTracker.mConfigRestrictsAvoidBadWifi = true;
+        mPolicyTracker.reevaluate();
         defaultCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         assertEquals(mCm.getActiveNetwork(), wifiNetwork);
 
@@ -3763,7 +3245,7 @@
 
         // Disconnect and reconnect wifi to clear the one-time switch above.
         mWiFiNetworkAgent.disconnect();
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
         validatedWifiCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
@@ -3773,11 +3255,11 @@
         mWiFiNetworkAgent.setNetworkInvalid();
         mCm.reportNetworkConnectivity(wifiNetwork, false);
         defaultCallback.expectCapabilitiesWithout(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
-        validatedWifiCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        validatedWifiCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // Simulate the user selecting "switch" and checking the don't ask again checkbox.
         Settings.Global.putInt(cr, Settings.Global.NETWORK_AVOID_BAD_WIFI, 1);
-        tracker.reevaluate();
+        mPolicyTracker.reevaluate();
 
         // We now switch to cell.
         defaultCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
@@ -3790,17 +3272,17 @@
         // Simulate the user turning the cellular fallback setting off and then on.
         // We switch to wifi and then to cell.
         Settings.Global.putString(cr, Settings.Global.NETWORK_AVOID_BAD_WIFI, null);
-        tracker.reevaluate();
+        mPolicyTracker.reevaluate();
         defaultCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         assertEquals(mCm.getActiveNetwork(), wifiNetwork);
         Settings.Global.putInt(cr, Settings.Global.NETWORK_AVOID_BAD_WIFI, 1);
-        tracker.reevaluate();
+        mPolicyTracker.reevaluate();
         defaultCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         assertEquals(mCm.getActiveNetwork(), cellNetwork);
 
         // If cell goes down, we switch to wifi.
         mCellNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         validatedWifiCallback.assertNoCallback();
 
@@ -3812,14 +3294,13 @@
     @Test
     public void testMeteredMultipathPreferenceSetting() throws Exception {
         final ContentResolver cr = mServiceContext.getContentResolver();
-        final WrappedMultinetworkPolicyTracker tracker = mService.getMultinetworkPolicyTracker();
         final String settingName = Settings.Global.NETWORK_METERED_MULTIPATH_PREFERENCE;
 
         for (int config : Arrays.asList(0, 3, 2)) {
             for (String setting: Arrays.asList(null, "0", "2", "1")) {
-                tracker.configMeteredMultipathPreference = config;
+                mPolicyTracker.mConfigMeteredMultipathPreference = config;
                 Settings.Global.putString(cr, settingName, setting);
-                tracker.reevaluate();
+                mPolicyTracker.reevaluate();
                 waitForIdle();
 
                 final int expected = (setting != null) ? Integer.parseInt(setting) : config;
@@ -3834,13 +3315,13 @@
      * time-out period expires.
      */
     @Test
-    public void testSatisfiedNetworkRequestDoesNotTriggerOnUnavailable() {
+    public void testSatisfiedNetworkRequestDoesNotTriggerOnUnavailable() throws Exception {
         NetworkRequest nr = new NetworkRequest.Builder().addTransportType(
                 NetworkCapabilities.TRANSPORT_WIFI).build();
         final TestNetworkCallback networkCallback = new TestNetworkCallback();
         mCm.requestNetwork(nr, networkCallback, TEST_REQUEST_TIMEOUT_MS);
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
         networkCallback.expectAvailableCallbacks(mWiFiNetworkAgent, false, false, false,
                 TEST_CALLBACK_TIMEOUT_MS);
@@ -3854,18 +3335,18 @@
      * not trigger onUnavailable() once the time-out period expires.
      */
     @Test
-    public void testSatisfiedThenLostNetworkRequestDoesNotTriggerOnUnavailable() {
+    public void testSatisfiedThenLostNetworkRequestDoesNotTriggerOnUnavailable() throws Exception {
         NetworkRequest nr = new NetworkRequest.Builder().addTransportType(
                 NetworkCapabilities.TRANSPORT_WIFI).build();
         final TestNetworkCallback networkCallback = new TestNetworkCallback();
         mCm.requestNetwork(nr, networkCallback, TEST_REQUEST_TIMEOUT_MS);
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
         networkCallback.expectAvailableCallbacks(mWiFiNetworkAgent, false, false, false,
                 TEST_CALLBACK_TIMEOUT_MS);
         mWiFiNetworkAgent.disconnect();
-        networkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // Validate that UNAVAILABLE is not called
         networkCallback.assertNoCallback();
@@ -3877,7 +3358,7 @@
      * (somehow) satisfied - the callback isn't called later.
      */
     @Test
-    public void testTimedoutNetworkRequest() {
+    public void testTimedoutNetworkRequest() throws Exception {
         NetworkRequest nr = new NetworkRequest.Builder().addTransportType(
                 NetworkCapabilities.TRANSPORT_WIFI).build();
         final TestNetworkCallback networkCallback = new TestNetworkCallback();
@@ -3885,10 +3366,10 @@
         mCm.requestNetwork(nr, networkCallback, timeoutMs);
 
         // pass timeout and validate that UNAVAILABLE is called
-        networkCallback.expectCallback(CallbackState.UNAVAILABLE, null);
+        networkCallback.expectCallback(CallbackRecord.UNAVAILABLE, null);
 
         // create a network satisfying request - validate that request not triggered
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
         networkCallback.assertNoCallback();
     }
@@ -3898,7 +3379,7 @@
      * trigger the callback.
      */
     @Test
-    public void testNoCallbackAfterUnregisteredNetworkRequest() {
+    public void testNoCallbackAfterUnregisteredNetworkRequest() throws Exception {
         NetworkRequest nr = new NetworkRequest.Builder().addTransportType(
                 NetworkCapabilities.TRANSPORT_WIFI).build();
         final TestNetworkCallback networkCallback = new TestNetworkCallback();
@@ -3911,7 +3392,7 @@
         networkCallback.assertNoCallback();
 
         // create a network satisfying request - validate that request not triggered
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
         networkCallback.assertNoCallback();
     }
@@ -3976,7 +3457,7 @@
             // Simulate the factory releasing the request as unfulfillable and expect onUnavailable!
             testFactory.triggerUnfulfillable(requests.get(newRequestId));
 
-            networkCallback.expectCallback(CallbackState.UNAVAILABLE, null);
+            networkCallback.expectCallback(CallbackRecord.UNAVAILABLE, null);
             testFactory.waitForRequests();
 
             // unregister network callback - a no-op (since already freed by the
@@ -3990,7 +3471,7 @@
 
     private static class TestKeepaliveCallback extends PacketKeepaliveCallback {
 
-        public static enum CallbackType { ON_STARTED, ON_STOPPED, ON_ERROR };
+        public enum CallbackType { ON_STARTED, ON_STOPPED, ON_ERROR }
 
         private class CallbackValue {
             public CallbackType callbackType;
@@ -4038,25 +3519,19 @@
             mCallbacks.add(new CallbackValue(CallbackType.ON_ERROR, error));
         }
 
-        private void expectCallback(CallbackValue callbackValue) {
-            try {
-                assertEquals(
-                        callbackValue,
-                        mCallbacks.poll(TIMEOUT_MS, TimeUnit.MILLISECONDS));
-            } catch (InterruptedException e) {
-                fail(callbackValue.callbackType + " callback not seen after " + TIMEOUT_MS + " ms");
-            }
+        private void expectCallback(CallbackValue callbackValue) throws InterruptedException {
+            assertEquals(callbackValue, mCallbacks.poll(TIMEOUT_MS, TimeUnit.MILLISECONDS));
         }
 
-        public void expectStarted() {
+        public void expectStarted() throws Exception {
             expectCallback(new CallbackValue(CallbackType.ON_STARTED));
         }
 
-        public void expectStopped() {
+        public void expectStopped() throws Exception {
             expectCallback(new CallbackValue(CallbackType.ON_STOPPED));
         }
 
-        public void expectError(int error) {
+        public void expectError(int error) throws Exception {
             expectCallback(new CallbackValue(CallbackType.ON_ERROR, error));
         }
     }
@@ -4117,25 +3592,20 @@
             mCallbacks.add(new CallbackValue(CallbackType.ON_ERROR, error));
         }
 
-        private void expectCallback(CallbackValue callbackValue) {
-            try {
-                assertEquals(
-                        callbackValue,
-                        mCallbacks.poll(TIMEOUT_MS, TimeUnit.MILLISECONDS));
-            } catch (InterruptedException e) {
-                fail(callbackValue.callbackType + " callback not seen after " + TIMEOUT_MS + " ms");
-            }
+        private void expectCallback(CallbackValue callbackValue) throws InterruptedException {
+            assertEquals(callbackValue, mCallbacks.poll(TIMEOUT_MS, TimeUnit.MILLISECONDS));
+
         }
 
-        public void expectStarted() {
+        public void expectStarted() throws InterruptedException {
             expectCallback(new CallbackValue(CallbackType.ON_STARTED));
         }
 
-        public void expectStopped() {
+        public void expectStopped() throws InterruptedException {
             expectCallback(new CallbackValue(CallbackType.ON_STOPPED));
         }
 
-        public void expectError(int error) {
+        public void expectError(int error) throws InterruptedException {
             expectCallback(new CallbackValue(CallbackType.ON_ERROR, error));
         }
 
@@ -4146,13 +3616,13 @@
         }
     }
 
-    private Network connectKeepaliveNetwork(LinkProperties lp) {
+    private Network connectKeepaliveNetwork(LinkProperties lp) throws Exception {
         // Ensure the network is disconnected before we do anything.
         if (mWiFiNetworkAgent != null) {
             assertNull(mCm.getNetworkCapabilities(mWiFiNetworkAgent.getNetwork()));
         }
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         ConditionVariable cv = waitForConnectivityBroadcasts(1);
         mWiFiNetworkAgent.connect(true);
         waitFor(cv);
@@ -4216,10 +3686,10 @@
         callback.expectError(PacketKeepalive.ERROR_HARDWARE_UNSUPPORTED);
 
         // Check that a started keepalive can be stopped.
-        mWiFiNetworkAgent.setStartKeepaliveError(PacketKeepalive.SUCCESS);
+        mWiFiNetworkAgent.setStartKeepaliveEvent(PacketKeepalive.SUCCESS);
         ka = mCm.startNattKeepalive(myNet, validKaInterval, callback, myIPv4, 12345, dstIPv4);
         callback.expectStarted();
-        mWiFiNetworkAgent.setStopKeepaliveError(PacketKeepalive.SUCCESS);
+        mWiFiNetworkAgent.setStopKeepaliveEvent(PacketKeepalive.SUCCESS);
         ka.stop();
         callback.expectStopped();
 
@@ -4237,7 +3707,7 @@
         ka = mCm.startNattKeepalive(myNet, validKaInterval, callback, myIPv4, 12345, dstIPv4);
         callback.expectStarted();
         mWiFiNetworkAgent.disconnect();
-        waitFor(mWiFiNetworkAgent.getDisconnectedCV());
+        mWiFiNetworkAgent.expectDisconnected();
         callback.expectError(PacketKeepalive.ERROR_INVALID_NETWORK);
 
         // ... and that stopping it after that has no adverse effects.
@@ -4248,7 +3718,7 @@
 
         // Reconnect.
         myNet = connectKeepaliveNetwork(lp);
-        mWiFiNetworkAgent.setStartKeepaliveError(PacketKeepalive.SUCCESS);
+        mWiFiNetworkAgent.setStartKeepaliveEvent(PacketKeepalive.SUCCESS);
 
         // Check that keepalive slots start from 1 and increment. The first one gets slot 1.
         mWiFiNetworkAgent.setExpectedKeepaliveSlot(1);
@@ -4279,13 +3749,9 @@
         callback3.expectStopped();
     }
 
-    @FunctionalInterface
-    private interface ThrowingConsumer<T> {
-        void accept(T t) throws Exception;
-    }
-
     // Helper method to prepare the executor and run test
-    private void runTestWithSerialExecutors(ThrowingConsumer<Executor> functor) throws Exception {
+    private void runTestWithSerialExecutors(ExceptionUtils.ThrowingConsumer<Executor> functor)
+            throws Exception {
         final ExecutorService executorSingleThread = Executors.newSingleThreadExecutor();
         final Executor executorInline = (Runnable r) -> r.run();
         functor.accept(executorSingleThread);
@@ -4372,12 +3838,12 @@
         }
 
         // Check that a started keepalive can be stopped.
-        mWiFiNetworkAgent.setStartKeepaliveError(SocketKeepalive.SUCCESS);
+        mWiFiNetworkAgent.setStartKeepaliveEvent(SocketKeepalive.SUCCESS);
         try (SocketKeepalive ka = mCm.createSocketKeepalive(
                 myNet, testSocket, myIPv4, dstIPv4, executor, callback)) {
             ka.start(validKaInterval);
             callback.expectStarted();
-            mWiFiNetworkAgent.setStopKeepaliveError(SocketKeepalive.SUCCESS);
+            mWiFiNetworkAgent.setStopKeepaliveEvent(SocketKeepalive.SUCCESS);
             ka.stop();
             callback.expectStopped();
 
@@ -4417,7 +3883,7 @@
             ka.start(validKaInterval);
             callback.expectStarted();
             mWiFiNetworkAgent.disconnect();
-            waitFor(mWiFiNetworkAgent.getDisconnectedCV());
+            mWiFiNetworkAgent.expectDisconnected();
             callback.expectError(SocketKeepalive.ERROR_INVALID_NETWORK);
 
             // ... and that stopping it after that has no adverse effects.
@@ -4430,7 +3896,7 @@
 
         // Reconnect.
         myNet = connectKeepaliveNetwork(lp);
-        mWiFiNetworkAgent.setStartKeepaliveError(SocketKeepalive.SUCCESS);
+        mWiFiNetworkAgent.setStartKeepaliveEvent(SocketKeepalive.SUCCESS);
 
         // Check that keepalive slots start from 1 and increment. The first one gets slot 1.
         mWiFiNetworkAgent.setExpectedKeepaliveSlot(1);
@@ -4467,7 +3933,7 @@
         // assertFalse(isUdpPortInUse(srcPort2));
 
         mWiFiNetworkAgent.disconnect();
-        waitFor(mWiFiNetworkAgent.getDisconnectedCV());
+        mWiFiNetworkAgent.expectDisconnected();
         mWiFiNetworkAgent = null;
     }
 
@@ -4543,7 +4009,7 @@
         testSocketV6.close();
 
         mWiFiNetworkAgent.disconnect();
-        waitFor(mWiFiNetworkAgent.getDisconnectedCV());
+        mWiFiNetworkAgent.expectDisconnected();
         mWiFiNetworkAgent = null;
     }
 
@@ -4559,8 +4025,8 @@
         lp.addLinkAddress(new LinkAddress(myIPv4, 25));
         lp.addRoute(new RouteInfo(InetAddress.getByName("192.0.2.254")));
         Network myNet = connectKeepaliveNetwork(lp);
-        mWiFiNetworkAgent.setStartKeepaliveError(SocketKeepalive.SUCCESS);
-        mWiFiNetworkAgent.setStopKeepaliveError(SocketKeepalive.SUCCESS);
+        mWiFiNetworkAgent.setStartKeepaliveEvent(SocketKeepalive.SUCCESS);
+        mWiFiNetworkAgent.setStopKeepaliveEvent(SocketKeepalive.SUCCESS);
 
         TestSocketKeepaliveCallback callback = new TestSocketKeepaliveCallback(executor);
 
@@ -4596,14 +4062,14 @@
         // assertFalse(isUdpPortInUse(srcPort));
 
         mWiFiNetworkAgent.disconnect();
-        waitFor(mWiFiNetworkAgent.getDisconnectedCV());
+        mWiFiNetworkAgent.expectDisconnected();
         mWiFiNetworkAgent = null;
     }
 
     private static boolean isUdpPortInUse(int port) {
         try (DatagramSocket ignored = new DatagramSocket(port)) {
             return false;
-        } catch (IOException ignored) {
+        } catch (IOException alreadyInUse) {
             return true;
         }
     }
@@ -4615,23 +4081,19 @@
     }
 
     private static class TestNetworkPinner extends NetworkPinner {
-        public static boolean awaitPin(int timeoutMs) {
+        public static boolean awaitPin(int timeoutMs) throws InterruptedException {
             synchronized(sLock) {
                 if (sNetwork == null) {
-                    try {
-                        sLock.wait(timeoutMs);
-                    } catch (InterruptedException e) {}
+                    sLock.wait(timeoutMs);
                 }
                 return sNetwork != null;
             }
         }
 
-        public static boolean awaitUnpin(int timeoutMs) {
+        public static boolean awaitUnpin(int timeoutMs) throws InterruptedException {
             synchronized(sLock) {
                 if (sNetwork != null) {
-                    try {
-                        sLock.wait(timeoutMs);
-                    } catch (InterruptedException e) {}
+                    sLock.wait(timeoutMs);
                 }
                 return sNetwork == null;
             }
@@ -4654,7 +4116,7 @@
     }
 
     @Test
-    public void testNetworkPinner() {
+    public void testNetworkPinner() throws Exception {
         NetworkRequest wifiRequest = new NetworkRequest.Builder()
                 .addTransportType(TRANSPORT_WIFI)
                 .build();
@@ -4663,9 +4125,9 @@
         TestNetworkPinner.pin(mServiceContext, wifiRequest);
         assertNull(mCm.getBoundNetworkForProcess());
 
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
 
         // When wi-fi connects, expect to be pinned.
@@ -4678,7 +4140,7 @@
         assertNotPinnedToWifi();
 
         // Reconnecting does not cause the pin to come back.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
         assertFalse(TestNetworkPinner.awaitPin(100));
         assertNotPinnedToWifi();
@@ -4700,14 +4162,14 @@
 
         // Pinning takes effect even if the pinned network is the default when the pin is set...
         TestNetworkPinner.pin(mServiceContext, wifiRequest);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
         assertTrue(TestNetworkPinner.awaitPin(100));
         assertPinnedToWifiWithWifiDefault();
 
         // ... and is maintained even when that network is no longer the default.
         cv = waitForConnectivityBroadcasts(1);
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mCellNetworkAgent.connect(true);
         waitFor(cv);
         assertPinnedToWifiWithCellDefault();
@@ -4749,25 +4211,20 @@
         }
 
         // Test that the limit is enforced when MAX_REQUESTS simultaneous requests are added.
-        try {
-            mCm.requestNetwork(networkRequest, new NetworkCallback());
-            fail("Registering " + MAX_REQUESTS + " network requests did not throw exception");
-        } catch (TooManyRequestsException expected) {}
-        try {
-            mCm.registerNetworkCallback(networkRequest, new NetworkCallback());
-            fail("Registering " + MAX_REQUESTS + " network callbacks did not throw exception");
-        } catch (TooManyRequestsException expected) {}
-        try {
-            mCm.requestNetwork(networkRequest,
-                PendingIntent.getBroadcast(mContext, 0, new Intent("c"), 0));
-            fail("Registering " + MAX_REQUESTS + " PendingIntent requests did not throw exception");
-        } catch (TooManyRequestsException expected) {}
-        try {
-            mCm.registerNetworkCallback(networkRequest,
-                PendingIntent.getBroadcast(mContext, 0, new Intent("d"), 0));
-            fail("Registering " + MAX_REQUESTS
-                    + " PendingIntent callbacks did not throw exception");
-        } catch (TooManyRequestsException expected) {}
+        assertThrows(TooManyRequestsException.class, () ->
+                mCm.requestNetwork(networkRequest, new NetworkCallback())
+        );
+        assertThrows(TooManyRequestsException.class, () ->
+                mCm.registerNetworkCallback(networkRequest, new NetworkCallback())
+        );
+        assertThrows(TooManyRequestsException.class, () ->
+                mCm.requestNetwork(networkRequest,
+                        PendingIntent.getBroadcast(mContext, 0, new Intent("c"), 0))
+        );
+        assertThrows(TooManyRequestsException.class, () ->
+                mCm.registerNetworkCallback(networkRequest,
+                        PendingIntent.getBroadcast(mContext, 0, new Intent("d"), 0))
+        );
 
         for (Object o : registered) {
             if (o instanceof NetworkCallback) {
@@ -4811,11 +4268,11 @@
     }
 
     @Test
-    public void testNetworkInfoOfTypeNone() {
+    public void testNetworkInfoOfTypeNone() throws Exception {
         ConditionVariable broadcastCV = waitForConnectivityBroadcasts(1);
 
         verifyNoNetwork();
-        MockNetworkAgent wifiAware = new MockNetworkAgent(TRANSPORT_WIFI_AWARE);
+        TestNetworkAgentWrapper wifiAware = new TestNetworkAgentWrapper(TRANSPORT_WIFI_AWARE);
         assertNull(mCm.getActiveNetworkInfo());
 
         Network[] allNetworks = mCm.getAllNetworks();
@@ -4841,7 +4298,7 @@
 
         // Disconnect wifi aware network.
         wifiAware.disconnect();
-        callback.expectCallbackLike((info) -> info.state == CallbackState.LOST, TIMEOUT_MS);
+        callback.expectCallbackThat(TIMEOUT_MS, (info) -> info instanceof CallbackRecord.Lost);
         mCm.unregisterNetworkCallback(callback);
 
         verifyNoNetwork();
@@ -4858,21 +4315,21 @@
         assertNull(mCm.getLinkProperties(TYPE_NONE));
         assertFalse(mCm.isNetworkSupported(TYPE_NONE));
 
-        assertException(() -> { mCm.networkCapabilitiesForType(TYPE_NONE); },
-                IllegalArgumentException.class);
+        assertThrows(IllegalArgumentException.class,
+                () -> { mCm.networkCapabilitiesForType(TYPE_NONE); });
 
         Class<UnsupportedOperationException> unsupported = UnsupportedOperationException.class;
-        assertException(() -> { mCm.startUsingNetworkFeature(TYPE_WIFI, ""); }, unsupported);
-        assertException(() -> { mCm.stopUsingNetworkFeature(TYPE_WIFI, ""); }, unsupported);
+        assertThrows(unsupported, () -> { mCm.startUsingNetworkFeature(TYPE_WIFI, ""); });
+        assertThrows(unsupported, () -> { mCm.stopUsingNetworkFeature(TYPE_WIFI, ""); });
         // TODO: let test context have configuration application target sdk version
         // and test that pre-M requesting for TYPE_NONE sends back APN_REQUEST_FAILED
-        assertException(() -> { mCm.startUsingNetworkFeature(TYPE_NONE, ""); }, unsupported);
-        assertException(() -> { mCm.stopUsingNetworkFeature(TYPE_NONE, ""); }, unsupported);
-        assertException(() -> { mCm.requestRouteToHostAddress(TYPE_NONE, null); }, unsupported);
+        assertThrows(unsupported, () -> { mCm.startUsingNetworkFeature(TYPE_NONE, ""); });
+        assertThrows(unsupported, () -> { mCm.stopUsingNetworkFeature(TYPE_NONE, ""); });
+        assertThrows(unsupported, () -> { mCm.requestRouteToHostAddress(TYPE_NONE, null); });
     }
 
     @Test
-    public void testLinkPropertiesEnsuresDirectlyConnectedRoutes() {
+    public void testLinkPropertiesEnsuresDirectlyConnectedRoutes() throws Exception {
         final NetworkRequest networkRequest = new NetworkRequest.Builder()
                 .addTransportType(TRANSPORT_WIFI).build();
         final TestNetworkCallback networkCallback = new TestNetworkCallback();
@@ -4888,16 +4345,17 @@
 
         // Verify direct routes are added when network agent is first registered in
         // ConnectivityService.
-        MockNetworkAgent networkAgent = new MockNetworkAgent(TRANSPORT_WIFI, lp);
+        TestNetworkAgentWrapper networkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI, lp);
         networkAgent.connect(true);
-        networkCallback.expectCallback(CallbackState.AVAILABLE, networkAgent);
-        networkCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES, networkAgent);
-        CallbackInfo cbi = networkCallback.expectCallback(CallbackState.LINK_PROPERTIES,
+        networkCallback.expectCallback(CallbackRecord.AVAILABLE, networkAgent);
+        networkCallback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED, networkAgent);
+        CallbackRecord.LinkPropertiesChanged cbi =
+                networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED,
                 networkAgent);
-        networkCallback.expectCallback(CallbackState.BLOCKED_STATUS, networkAgent);
+        networkCallback.expectCallback(CallbackRecord.BLOCKED_STATUS, networkAgent);
         networkCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, networkAgent);
         networkCallback.assertNoCallback();
-        checkDirectlyConnectedRoutes(cbi.arg, Arrays.asList(myIpv4Address),
+        checkDirectlyConnectedRoutes(cbi.getLp(), Arrays.asList(myIpv4Address),
                 Arrays.asList(myIpv4DefaultRoute));
         checkDirectlyConnectedRoutes(mCm.getLinkProperties(networkAgent.getNetwork()),
                 Arrays.asList(myIpv4Address), Arrays.asList(myIpv4DefaultRoute));
@@ -4909,9 +4367,9 @@
         newLp.addLinkAddress(myIpv6Address1);
         newLp.addLinkAddress(myIpv6Address2);
         networkAgent.sendLinkProperties(newLp);
-        cbi = networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, networkAgent);
+        cbi = networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, networkAgent);
         networkCallback.assertNoCallback();
-        checkDirectlyConnectedRoutes(cbi.arg,
+        checkDirectlyConnectedRoutes(cbi.getLp(),
                 Arrays.asList(myIpv4Address, myIpv6Address1, myIpv6Address2),
                 Arrays.asList(myIpv4DefaultRoute));
         mCm.unregisterNetworkCallback(networkCallback);
@@ -4919,8 +4377,8 @@
 
     @Test
     public void testStatsIfacesChanged() throws Exception {
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
 
         Network[] onlyCell = new Network[] {mCellNetworkAgent.getNetwork()};
         Network[] onlyWifi = new Network[] {mWiFiNetworkAgent.getNetwork()};
@@ -4935,11 +4393,8 @@
         mCellNetworkAgent.sendLinkProperties(cellLp);
         waitForIdle();
         verify(mStatsService, atLeastOnce())
-                .forceUpdateIfaces(
-                        eq(onlyCell),
-                        any(NetworkState[].class),
-                        eq(MOBILE_IFNAME));
-        assertEquals(new VpnInfo[0], NetworkStatsFactory.getVpnInfos());
+                .forceUpdateIfaces(eq(onlyCell), any(NetworkState[].class), eq(MOBILE_IFNAME),
+                        eq(new VpnInfo[0]));
         reset(mStatsService);
 
         // Default network switch should update ifaces.
@@ -4948,65 +4403,47 @@
         waitForIdle();
         assertEquals(wifiLp, mService.getActiveLinkProperties());
         verify(mStatsService, atLeastOnce())
-                .forceUpdateIfaces(
-                        eq(onlyWifi),
-                        any(NetworkState[].class),
-                        eq(WIFI_IFNAME));
-        assertEquals(new VpnInfo[0], NetworkStatsFactory.getVpnInfos());
+                .forceUpdateIfaces(eq(onlyWifi), any(NetworkState[].class), eq(WIFI_IFNAME),
+                        eq(new VpnInfo[0]));
         reset(mStatsService);
 
         // Disconnect should update ifaces.
         mWiFiNetworkAgent.disconnect();
         waitForIdle();
         verify(mStatsService, atLeastOnce())
-                .forceUpdateIfaces(
-                        eq(onlyCell),
-                        any(NetworkState[].class),
-                        eq(MOBILE_IFNAME));
-        assertEquals(new VpnInfo[0], NetworkStatsFactory.getVpnInfos());
+                .forceUpdateIfaces(eq(onlyCell), any(NetworkState[].class),
+                        eq(MOBILE_IFNAME), eq(new VpnInfo[0]));
         reset(mStatsService);
 
         // Metered change should update ifaces
         mCellNetworkAgent.addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_METERED);
         waitForIdle();
         verify(mStatsService, atLeastOnce())
-                .forceUpdateIfaces(
-                        eq(onlyCell),
-                        any(NetworkState[].class),
-                        eq(MOBILE_IFNAME));
-        assertEquals(new VpnInfo[0], NetworkStatsFactory.getVpnInfos());
+                .forceUpdateIfaces(eq(onlyCell), any(NetworkState[].class), eq(MOBILE_IFNAME),
+                        eq(new VpnInfo[0]));
         reset(mStatsService);
 
         mCellNetworkAgent.removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_METERED);
         waitForIdle();
         verify(mStatsService, atLeastOnce())
-                .forceUpdateIfaces(
-                        eq(onlyCell),
-                        any(NetworkState[].class),
-                        eq(MOBILE_IFNAME));
-        assertEquals(new VpnInfo[0], NetworkStatsFactory.getVpnInfos());
+                .forceUpdateIfaces(eq(onlyCell), any(NetworkState[].class), eq(MOBILE_IFNAME),
+                        eq(new VpnInfo[0]));
         reset(mStatsService);
 
         // Captive portal change shouldn't update ifaces
         mCellNetworkAgent.addCapability(NetworkCapabilities.NET_CAPABILITY_CAPTIVE_PORTAL);
         waitForIdle();
         verify(mStatsService, never())
-                .forceUpdateIfaces(
-                        eq(onlyCell),
-                        any(NetworkState[].class),
-                        eq(MOBILE_IFNAME));
-        assertEquals(new VpnInfo[0], NetworkStatsFactory.getVpnInfos());
+                .forceUpdateIfaces(eq(onlyCell), any(NetworkState[].class), eq(MOBILE_IFNAME),
+                        eq(new VpnInfo[0]));
         reset(mStatsService);
 
         // Roaming change should update ifaces
         mCellNetworkAgent.addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_ROAMING);
         waitForIdle();
         verify(mStatsService, atLeastOnce())
-                .forceUpdateIfaces(
-                        eq(onlyCell),
-                        any(NetworkState[].class),
-                        eq(MOBILE_IFNAME));
-        assertEquals(new VpnInfo[0], NetworkStatsFactory.getVpnInfos());
+                .forceUpdateIfaces(eq(onlyCell), any(NetworkState[].class), eq(MOBILE_IFNAME),
+                        eq(new VpnInfo[0]));
         reset(mStatsService);
     }
 
@@ -5017,7 +4454,7 @@
         // Clear any interactions that occur as a result of CS starting up.
         reset(mMockDnsResolver);
 
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         waitForIdle();
         verify(mMockDnsResolver, never()).setResolverConfiguration(any());
         verifyNoMoreInteractions(mMockDnsResolver);
@@ -5100,7 +4537,7 @@
                 .addTransportType(TRANSPORT_CELLULAR).build();
         mCm.requestNetwork(cellRequest, cellNetworkCallback);
 
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         waitForIdle();
         // CS tells netd about the empty DNS config for this network.
         verify(mMockDnsResolver, never()).setResolverConfiguration(any());
@@ -5130,21 +4567,21 @@
         ResolverParamsParcel resolvrParams = mResolverParamsParcelCaptor.getValue();
         assertEquals(2, resolvrParams.tlsServers.length);
         assertTrue(ArrayUtils.containsAll(resolvrParams.tlsServers,
-                new String[]{"2001:db8::1", "192.0.2.1"}));
+                new String[] { "2001:db8::1", "192.0.2.1" }));
         // Opportunistic mode.
         assertEquals(2, resolvrParams.tlsServers.length);
         assertTrue(ArrayUtils.containsAll(resolvrParams.tlsServers,
-                new String[]{"2001:db8::1", "192.0.2.1"}));
+                new String[] { "2001:db8::1", "192.0.2.1" }));
         reset(mMockDnsResolver);
-        cellNetworkCallback.expectCallback(CallbackState.AVAILABLE, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES,
+        cellNetworkCallback.expectCallback(CallbackRecord.AVAILABLE, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED,
                 mCellNetworkAgent);
-        CallbackInfo cbi = cellNetworkCallback.expectCallback(
-                CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.BLOCKED_STATUS, mCellNetworkAgent);
+        CallbackRecord.LinkPropertiesChanged cbi = cellNetworkCallback.expectCallback(
+                CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.BLOCKED_STATUS, mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
-        assertFalse(((LinkProperties)cbi.arg).isPrivateDnsActive());
-        assertNull(((LinkProperties)cbi.arg).getPrivateDnsServerName());
+        assertFalse(cbi.getLp().isPrivateDnsActive());
+        assertNull(cbi.getLp().getPrivateDnsServerName());
 
         setPrivateDnsSettings(PRIVATE_DNS_MODE_OFF, "ignored.example.com");
         verify(mMockDnsResolver, times(1)).setResolverConfiguration(
@@ -5152,7 +4589,7 @@
         resolvrParams = mResolverParamsParcelCaptor.getValue();
         assertEquals(2, resolvrParams.servers.length);
         assertTrue(ArrayUtils.containsAll(resolvrParams.servers,
-                new String[]{"2001:db8::1", "192.0.2.1"}));
+                new String[] { "2001:db8::1", "192.0.2.1" }));
         reset(mMockDnsResolver);
         cellNetworkCallback.assertNoCallback();
 
@@ -5162,21 +4599,21 @@
         resolvrParams = mResolverParamsParcelCaptor.getValue();
         assertEquals(2, resolvrParams.servers.length);
         assertTrue(ArrayUtils.containsAll(resolvrParams.servers,
-                new String[]{"2001:db8::1", "192.0.2.1"}));
+                new String[] { "2001:db8::1", "192.0.2.1" }));
         assertEquals(2, resolvrParams.tlsServers.length);
         assertTrue(ArrayUtils.containsAll(resolvrParams.tlsServers,
-                new String[]{"2001:db8::1", "192.0.2.1"}));
+                new String[] { "2001:db8::1", "192.0.2.1" }));
         reset(mMockDnsResolver);
         cellNetworkCallback.assertNoCallback();
 
         setPrivateDnsSettings(PRIVATE_DNS_MODE_PROVIDER_HOSTNAME, "strict.example.com");
         // Can't test dns configuration for strict mode without properly mocking
         // out the DNS lookups, but can test that LinkProperties is updated.
-        cbi = cellNetworkCallback.expectCallback(CallbackState.LINK_PROPERTIES,
+        cbi = cellNetworkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
-        assertTrue(((LinkProperties)cbi.arg).isPrivateDnsActive());
-        assertEquals("strict.example.com", ((LinkProperties)cbi.arg).getPrivateDnsServerName());
+        assertTrue(cbi.getLp().isPrivateDnsActive());
+        assertEquals("strict.example.com", cbi.getLp().getPrivateDnsServerName());
     }
 
     @Test
@@ -5189,23 +4626,23 @@
                 .addTransportType(TRANSPORT_CELLULAR).build();
         mCm.requestNetwork(cellRequest, cellNetworkCallback);
 
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         waitForIdle();
         LinkProperties lp = new LinkProperties();
         mCellNetworkAgent.sendLinkProperties(lp);
         mCellNetworkAgent.connect(false);
         waitForIdle();
-        cellNetworkCallback.expectCallback(CallbackState.AVAILABLE, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES,
+        cellNetworkCallback.expectCallback(CallbackRecord.AVAILABLE, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED,
                 mCellNetworkAgent);
-        CallbackInfo cbi = cellNetworkCallback.expectCallback(
-                CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.BLOCKED_STATUS, mCellNetworkAgent);
+        CallbackRecord.LinkPropertiesChanged cbi = cellNetworkCallback.expectCallback(
+                CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.BLOCKED_STATUS, mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
-        assertFalse(((LinkProperties)cbi.arg).isPrivateDnsActive());
-        assertNull(((LinkProperties)cbi.arg).getPrivateDnsServerName());
+        assertFalse(cbi.getLp().isPrivateDnsActive());
+        assertNull(cbi.getLp().getPrivateDnsServerName());
         Set<InetAddress> dnsServers = new HashSet<>();
-        checkDnsServers(cbi.arg, dnsServers);
+        checkDnsServers(cbi.getLp(), dnsServers);
 
         // Send a validation event for a server that is not part of the current
         // resolver config. The validation event should be ignored.
@@ -5217,13 +4654,13 @@
         LinkProperties lp2 = new LinkProperties(lp);
         lp2.addDnsServer(InetAddress.getByName("145.100.185.16"));
         mCellNetworkAgent.sendLinkProperties(lp2);
-        cbi = cellNetworkCallback.expectCallback(CallbackState.LINK_PROPERTIES,
+        cbi = cellNetworkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
-        assertFalse(((LinkProperties)cbi.arg).isPrivateDnsActive());
-        assertNull(((LinkProperties)cbi.arg).getPrivateDnsServerName());
+        assertFalse(cbi.getLp().isPrivateDnsActive());
+        assertNull(cbi.getLp().getPrivateDnsServerName());
         dnsServers.add(InetAddress.getByName("145.100.185.16"));
-        checkDnsServers(cbi.arg, dnsServers);
+        checkDnsServers(cbi.getLp(), dnsServers);
 
         // Send a validation event containing a hostname that is not part of
         // the current resolver config. The validation event should be ignored.
@@ -5241,39 +4678,39 @@
         // private dns fields should be sent.
         mService.mNetdEventCallback.onPrivateDnsValidationEvent(
                 mCellNetworkAgent.getNetwork().netId, "145.100.185.16", "", true);
-        cbi = cellNetworkCallback.expectCallback(CallbackState.LINK_PROPERTIES,
+        cbi = cellNetworkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
-        assertTrue(((LinkProperties)cbi.arg).isPrivateDnsActive());
-        assertNull(((LinkProperties)cbi.arg).getPrivateDnsServerName());
-        checkDnsServers(cbi.arg, dnsServers);
+        assertTrue(cbi.getLp().isPrivateDnsActive());
+        assertNull(cbi.getLp().getPrivateDnsServerName());
+        checkDnsServers(cbi.getLp(), dnsServers);
 
         // The private dns fields in LinkProperties should be preserved when
         // the network agent sends unrelated changes.
         LinkProperties lp3 = new LinkProperties(lp2);
         lp3.setMtu(1300);
         mCellNetworkAgent.sendLinkProperties(lp3);
-        cbi = cellNetworkCallback.expectCallback(CallbackState.LINK_PROPERTIES,
+        cbi = cellNetworkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
-        assertTrue(((LinkProperties)cbi.arg).isPrivateDnsActive());
-        assertNull(((LinkProperties)cbi.arg).getPrivateDnsServerName());
-        checkDnsServers(cbi.arg, dnsServers);
-        assertEquals(1300, ((LinkProperties)cbi.arg).getMtu());
+        assertTrue(cbi.getLp().isPrivateDnsActive());
+        assertNull(cbi.getLp().getPrivateDnsServerName());
+        checkDnsServers(cbi.getLp(), dnsServers);
+        assertEquals(1300, cbi.getLp().getMtu());
 
         // Removing the only validated server should affect the private dns
         // fields in LinkProperties.
         LinkProperties lp4 = new LinkProperties(lp3);
         lp4.removeDnsServer(InetAddress.getByName("145.100.185.16"));
         mCellNetworkAgent.sendLinkProperties(lp4);
-        cbi = cellNetworkCallback.expectCallback(CallbackState.LINK_PROPERTIES,
+        cbi = cellNetworkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
-        assertFalse(((LinkProperties)cbi.arg).isPrivateDnsActive());
-        assertNull(((LinkProperties)cbi.arg).getPrivateDnsServerName());
+        assertFalse(cbi.getLp().isPrivateDnsActive());
+        assertNull(cbi.getLp().getPrivateDnsServerName());
         dnsServers.remove(InetAddress.getByName("145.100.185.16"));
-        checkDnsServers(cbi.arg, dnsServers);
-        assertEquals(1300, ((LinkProperties)cbi.arg).getMtu());
+        checkDnsServers(cbi.getLp(), dnsServers);
+        assertEquals(1300, cbi.getLp().getMtu());
     }
 
     private void checkDirectlyConnectedRoutes(Object callbackObj,
@@ -5300,31 +4737,8 @@
         assertTrue(lp.getDnsServers().containsAll(dnsServers));
     }
 
-    private static <T> void assertEmpty(T[] ts) {
-        int length = ts.length;
-        assertEquals("expected empty array, but length was " + length, 0, length);
-    }
-
-    private static <T> void assertLength(int expected, T[] got) {
-        int length = got.length;
-        assertEquals(String.format("expected array of length %s, but length was %s for %s",
-                expected, length, Arrays.toString(got)), expected, length);
-    }
-
-    private static <T> void assertException(Runnable block, Class<T> expected) {
-        try {
-            block.run();
-            fail("Expected exception of type " + expected);
-        } catch (Exception got) {
-            if (!got.getClass().equals(expected)) {
-                fail("Expected exception of type " + expected + " but got " + got);
-            }
-            return;
-        }
-    }
-
     @Test
-    public void testVpnNetworkActive() {
+    public void testVpnNetworkActive() throws Exception {
         final int uid = Process.myUid();
 
         final TestNetworkCallback genericNetworkCallback = new TestNetworkCallback();
@@ -5347,7 +4761,7 @@
         mCm.registerDefaultNetworkCallback(defaultCallback);
         defaultCallback.assertNoCallback();
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
 
         genericNetworkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
@@ -5357,14 +4771,16 @@
         vpnNetworkCallback.assertNoCallback();
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
-        final MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        final TestNetworkAgentWrapper
+                vpnNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN);
         final ArraySet<UidRange> ranges = new ArraySet<>();
         ranges.add(new UidRange(uid, uid));
         mMockVpn.setNetworkAgent(vpnNetworkAgent);
         mMockVpn.setUids(ranges);
         // VPN networks do not satisfy the default request and are automatically validated
         // by NetworkMonitor
-        assertFalse(NetworkMonitorUtils.isValidationRequired(vpnNetworkAgent.mNetworkCapabilities));
+        assertFalse(NetworkMonitorUtils.isValidationRequired(
+                vpnNetworkAgent.getNetworkCapabilities()));
         vpnNetworkAgent.setNetworkValid();
 
         vpnNetworkAgent.connect(false);
@@ -5378,19 +4794,19 @@
         defaultCallback.expectAvailableCallbacksUnvalidated(vpnNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
-        genericNetworkCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES, vpnNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED, vpnNetworkAgent);
         genericNotVpnNetworkCallback.assertNoCallback();
-        vpnNetworkCallback.expectCapabilitiesLike(nc -> null == nc.getUids(), vpnNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES, vpnNetworkAgent);
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent, nc -> null == nc.getUids());
+        defaultCallback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED, vpnNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         ranges.clear();
         vpnNetworkAgent.setUids(ranges);
 
-        genericNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOST, vpnNetworkAgent);
         genericNotVpnNetworkCallback.assertNoCallback();
         wifiNetworkCallback.assertNoCallback();
-        vpnNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        vpnNetworkCallback.expectCallback(CallbackRecord.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
@@ -5398,7 +4814,7 @@
         // 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);
+        defaultCallback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED, vpnNetworkAgent);
 
         ranges.add(new UidRange(uid, uid));
         mMockVpn.setUids(ranges);
@@ -5410,23 +4826,23 @@
         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);
+        defaultCallback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED, vpnNetworkAgent);
 
         mWiFiNetworkAgent.disconnect();
 
-        genericNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        genericNotVpnNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        wifiNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        genericNotVpnNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        wifiNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         vpnNetworkCallback.assertNoCallback();
         defaultCallback.assertNoCallback();
 
         vpnNetworkAgent.disconnect();
 
-        genericNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOST, vpnNetworkAgent);
         genericNotVpnNetworkCallback.assertNoCallback();
         wifiNetworkCallback.assertNoCallback();
-        vpnNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        vpnNetworkCallback.expectCallback(CallbackRecord.LOST, vpnNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, vpnNetworkAgent);
         assertEquals(null, mCm.getActiveNetwork());
 
         mCm.unregisterNetworkCallback(genericNetworkCallback);
@@ -5436,19 +4852,20 @@
     }
 
     @Test
-    public void testVpnWithoutInternet() {
+    public void testVpnWithoutInternet() throws Exception {
         final int uid = Process.myUid();
 
         final TestNetworkCallback defaultCallback = new TestNetworkCallback();
         mCm.registerDefaultNetworkCallback(defaultCallback);
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
 
         defaultCallback.expectAvailableThenValidatedCallbacks(mWiFiNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
-        MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        TestNetworkAgentWrapper
+                vpnNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN);
         final ArraySet<UidRange> ranges = new ArraySet<>();
         ranges.add(new UidRange(uid, uid));
         mMockVpn.setNetworkAgent(vpnNetworkAgent);
@@ -5466,19 +4883,20 @@
     }
 
     @Test
-    public void testVpnWithInternet() {
+    public void testVpnWithInternet() throws Exception {
         final int uid = Process.myUid();
 
         final TestNetworkCallback defaultCallback = new TestNetworkCallback();
         mCm.registerDefaultNetworkCallback(defaultCallback);
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
 
         defaultCallback.expectAvailableThenValidatedCallbacks(mWiFiNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
-        MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        TestNetworkAgentWrapper
+                vpnNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN);
         final ArraySet<UidRange> ranges = new ArraySet<>();
         ranges.add(new UidRange(uid, uid));
         mMockVpn.setNetworkAgent(vpnNetworkAgent);
@@ -5490,7 +4908,7 @@
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         vpnNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, vpnNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
 
         mCm.unregisterNetworkCallback(defaultCallback);
@@ -5502,14 +4920,15 @@
         mCm.registerDefaultNetworkCallback(callback);
 
         // Bring up Ethernet.
-        mEthernetNetworkAgent = new MockNetworkAgent(TRANSPORT_ETHERNET);
+        mEthernetNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_ETHERNET);
         mEthernetNetworkAgent.connect(true);
         callback.expectAvailableThenValidatedCallbacks(mEthernetNetworkAgent);
         callback.assertNoCallback();
 
         // Bring up a VPN that has the INTERNET capability, initially unvalidated.
         final int uid = Process.myUid();
-        final MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        final TestNetworkAgentWrapper
+                vpnNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN);
         final ArraySet<UidRange> ranges = new ArraySet<>();
         ranges.add(new UidRange(uid, uid));
         mMockVpn.setNetworkAgent(vpnNetworkAgent);
@@ -5520,7 +4939,7 @@
         // Even though the VPN is unvalidated, it becomes the default network for our app.
         callback.expectAvailableCallbacksUnvalidated(vpnNetworkAgent);
         // TODO: this looks like a spurious callback.
-        callback.expectCallback(CallbackState.NETWORK_CAPABILITIES, vpnNetworkAgent);
+        callback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED, vpnNetworkAgent);
         callback.assertNoCallback();
 
         assertTrue(vpnNetworkAgent.getScore() > mEthernetNetworkAgent.getScore());
@@ -5531,9 +4950,10 @@
         assertFalse(nc.hasCapability(NET_CAPABILITY_VALIDATED));
         assertTrue(nc.hasCapability(NET_CAPABILITY_INTERNET));
 
-        assertFalse(NetworkMonitorUtils.isValidationRequired(vpnNetworkAgent.mNetworkCapabilities));
+        assertFalse(NetworkMonitorUtils.isValidationRequired(
+                vpnNetworkAgent.getNetworkCapabilities()));
         assertTrue(NetworkMonitorUtils.isPrivateDnsValidationRequired(
-                vpnNetworkAgent.mNetworkCapabilities));
+                vpnNetworkAgent.getNetworkCapabilities()));
 
         // Pretend that the VPN network validates.
         vpnNetworkAgent.setNetworkValid();
@@ -5544,12 +4964,12 @@
         callback.assertNoCallback();
 
         vpnNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, vpnNetworkAgent);
         callback.expectAvailableCallbacksValidated(mEthernetNetworkAgent);
     }
 
     @Test
-    public void testVpnSetUnderlyingNetworks() {
+    public void testVpnSetUnderlyingNetworks() throws Exception {
         final int uid = Process.myUid();
 
         final TestNetworkCallback vpnNetworkCallback = new TestNetworkCallback();
@@ -5561,7 +4981,8 @@
         mCm.registerNetworkCallback(vpnNetworkRequest, vpnNetworkCallback);
         vpnNetworkCallback.assertNoCallback();
 
-        final MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        final TestNetworkAgentWrapper
+                vpnNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN);
         final ArraySet<UidRange> ranges = new ArraySet<>();
         ranges.add(new UidRange(uid, uid));
         mMockVpn.setNetworkAgent(vpnNetworkAgent);
@@ -5578,76 +4999,76 @@
         assertFalse(nc.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // Connect cell and use it as an underlying network.
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
 
         mService.setUnderlyingNetworksForVpn(
                 new Network[] { mCellNetworkAgent.getNetwork() });
 
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
-                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(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)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
-                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // 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)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
-                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // Use Wifi but not cell. Note the VPN is now unmetered.
         mService.setUnderlyingNetworksForVpn(
                 new Network[] { mWiFiNetworkAgent.getNetwork() });
 
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && !caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
-                && caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // Use both again.
         mService.setUnderlyingNetworksForVpn(
                 new Network[] { mCellNetworkAgent.getNetwork(), mWiFiNetworkAgent.getNetwork() });
 
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
-                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // 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)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && !caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
-                && caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // Disconnect wifi too. No underlying networks means this is now metered.
         mWiFiNetworkAgent.disconnect();
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && !caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
-                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         mMockVpn.disconnect();
     }
 
     @Test
-    public void testNullUnderlyingNetworks() {
+    public void testNullUnderlyingNetworks() throws Exception {
         final int uid = Process.myUid();
 
         final TestNetworkCallback vpnNetworkCallback = new TestNetworkCallback();
@@ -5659,7 +5080,8 @@
         mCm.registerNetworkCallback(vpnNetworkRequest, vpnNetworkCallback);
         vpnNetworkCallback.assertNoCallback();
 
-        final MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        final TestNetworkAgentWrapper
+                vpnNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN);
         final ArraySet<UidRange> ranges = new ArraySet<>();
         ranges.add(new UidRange(uid, uid));
         mMockVpn.setNetworkAgent(vpnNetworkAgent);
@@ -5677,23 +5099,23 @@
         assertFalse(nc.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // Connect to Cell; Cell is the default network.
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
 
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
-                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // Connect to WiFi; WiFi is the new default.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.addCapability(NET_CAPABILITY_NOT_METERED);
         mWiFiNetworkAgent.connect(true);
 
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && !caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
-                && caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // Disconnect Cell. The default network did not change, so there shouldn't be any changes in
         // the capabilities.
@@ -5702,19 +5124,19 @@
         // Disconnect wifi too. Now we have no default network.
         mWiFiNetworkAgent.disconnect();
 
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && !caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
-                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         mMockVpn.disconnect();
     }
 
     @Test
-    public void testIsActiveNetworkMeteredOverWifi() {
+    public void testIsActiveNetworkMeteredOverWifi() throws Exception {
         // Returns true by default when no network is available.
         assertTrue(mCm.isActiveNetworkMetered());
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.addCapability(NET_CAPABILITY_NOT_METERED);
         mWiFiNetworkAgent.connect(true);
         waitForIdle();
@@ -5723,10 +5145,10 @@
     }
 
     @Test
-    public void testIsActiveNetworkMeteredOverCell() {
+    public void testIsActiveNetworkMeteredOverCell() throws Exception {
         // Returns true by default when no network is available.
         assertTrue(mCm.isActiveNetworkMetered());
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.removeCapability(NET_CAPABILITY_NOT_METERED);
         mCellNetworkAgent.connect(true);
         waitForIdle();
@@ -5735,17 +5157,18 @@
     }
 
     @Test
-    public void testIsActiveNetworkMeteredOverVpnTrackingPlatformDefault() {
+    public void testIsActiveNetworkMeteredOverVpnTrackingPlatformDefault() throws Exception {
         // Returns true by default when no network is available.
         assertTrue(mCm.isActiveNetworkMetered());
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.removeCapability(NET_CAPABILITY_NOT_METERED);
         mCellNetworkAgent.connect(true);
         waitForIdle();
         assertTrue(mCm.isActiveNetworkMetered());
 
         // Connect VPN network. By default it is using current default network (Cell).
-        MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        TestNetworkAgentWrapper
+                vpnNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN);
         final ArraySet<UidRange> ranges = new ArraySet<>();
         final int uid = Process.myUid();
         ranges.add(new UidRange(uid, uid));
@@ -5761,7 +5184,7 @@
         assertTrue(mCm.isActiveNetworkMetered());
 
         // Connect WiFi.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.addCapability(NET_CAPABILITY_NOT_METERED);
         mWiFiNetworkAgent.connect(true);
         waitForIdle();
@@ -5789,23 +5212,24 @@
     }
 
    @Test
-   public void testIsActiveNetworkMeteredOverVpnSpecifyingUnderlyingNetworks() {
+   public void testIsActiveNetworkMeteredOverVpnSpecifyingUnderlyingNetworks() throws Exception {
         // Returns true by default when no network is available.
         assertTrue(mCm.isActiveNetworkMetered());
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.removeCapability(NET_CAPABILITY_NOT_METERED);
         mCellNetworkAgent.connect(true);
         waitForIdle();
         assertTrue(mCm.isActiveNetworkMetered());
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.addCapability(NET_CAPABILITY_NOT_METERED);
         mWiFiNetworkAgent.connect(true);
         waitForIdle();
         assertFalse(mCm.isActiveNetworkMetered());
 
         // Connect VPN network.
-        MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        TestNetworkAgentWrapper
+                vpnNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN);
         final ArraySet<UidRange> ranges = new ArraySet<>();
         final int uid = Process.myUid();
         ranges.add(new UidRange(uid, uid));
@@ -5860,17 +5284,18 @@
     }
 
     @Test
-    public void testIsActiveNetworkMeteredOverAlwaysMeteredVpn() {
+    public void testIsActiveNetworkMeteredOverAlwaysMeteredVpn() throws Exception {
         // Returns true by default when no network is available.
         assertTrue(mCm.isActiveNetworkMetered());
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.addCapability(NET_CAPABILITY_NOT_METERED);
         mWiFiNetworkAgent.connect(true);
         waitForIdle();
         assertFalse(mCm.isActiveNetworkMetered());
 
         // Connect VPN network.
-        MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        TestNetworkAgentWrapper
+                vpnNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN);
         final ArraySet<UidRange> ranges = new ArraySet<>();
         final int uid = Process.myUid();
         ranges.add(new UidRange(uid, uid));
@@ -5907,28 +5332,28 @@
     }
 
     @Test
-    public void testNetworkBlockedStatus() {
+    public void testNetworkBlockedStatus() throws Exception {
         final TestNetworkCallback cellNetworkCallback = new TestNetworkCallback();
         final NetworkRequest cellRequest = new NetworkRequest.Builder()
                 .addTransportType(TRANSPORT_CELLULAR)
                 .build();
         mCm.registerNetworkCallback(cellRequest, cellNetworkCallback);
 
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
         cellNetworkCallback.expectAvailableThenValidatedCallbacks(mCellNetworkAgent);
 
-        mService.setUidRulesChanged(RULE_REJECT_ALL);
+        setUidRulesChanged(RULE_REJECT_ALL);
         cellNetworkCallback.expectBlockedStatusCallback(true, mCellNetworkAgent);
 
         // ConnectivityService should cache it not to invoke the callback again.
-        mService.setUidRulesChanged(RULE_REJECT_METERED);
+        setUidRulesChanged(RULE_REJECT_METERED);
         cellNetworkCallback.assertNoCallback();
 
-        mService.setUidRulesChanged(RULE_NONE);
+        setUidRulesChanged(RULE_NONE);
         cellNetworkCallback.expectBlockedStatusCallback(false, mCellNetworkAgent);
 
-        mService.setUidRulesChanged(RULE_REJECT_METERED);
+        setUidRulesChanged(RULE_REJECT_METERED);
         cellNetworkCallback.expectBlockedStatusCallback(true, mCellNetworkAgent);
 
         // Restrict the network based on UID rule and NOT_METERED capability change.
@@ -5939,18 +5364,18 @@
         cellNetworkCallback.expectCapabilitiesWithout(NET_CAPABILITY_NOT_METERED,
                 mCellNetworkAgent);
         cellNetworkCallback.expectBlockedStatusCallback(true, mCellNetworkAgent);
-        mService.setUidRulesChanged(RULE_ALLOW_METERED);
+        setUidRulesChanged(RULE_ALLOW_METERED);
         cellNetworkCallback.expectBlockedStatusCallback(false, mCellNetworkAgent);
 
-        mService.setUidRulesChanged(RULE_NONE);
+        setUidRulesChanged(RULE_NONE);
         cellNetworkCallback.assertNoCallback();
 
         // Restrict the network based on BackgroundRestricted.
-        mService.setRestrictBackgroundChanged(true);
+        setRestrictBackgroundChanged(true);
         cellNetworkCallback.expectBlockedStatusCallback(true, mCellNetworkAgent);
-        mService.setRestrictBackgroundChanged(true);
+        setRestrictBackgroundChanged(true);
         cellNetworkCallback.assertNoCallback();
-        mService.setRestrictBackgroundChanged(false);
+        setRestrictBackgroundChanged(false);
         cellNetworkCallback.expectBlockedStatusCallback(false, mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
 
@@ -5958,30 +5383,30 @@
     }
 
     @Test
-    public void testNetworkBlockedStatusBeforeAndAfterConnect() {
+    public void testNetworkBlockedStatusBeforeAndAfterConnect() throws Exception {
         final TestNetworkCallback defaultCallback = new TestNetworkCallback();
         mCm.registerDefaultNetworkCallback(defaultCallback);
 
         // No Networkcallbacks invoked before any network is active.
-        mService.setUidRulesChanged(RULE_REJECT_ALL);
-        mService.setUidRulesChanged(RULE_NONE);
-        mService.setUidRulesChanged(RULE_REJECT_METERED);
+        setUidRulesChanged(RULE_REJECT_ALL);
+        setUidRulesChanged(RULE_NONE);
+        setUidRulesChanged(RULE_REJECT_METERED);
         defaultCallback.assertNoCallback();
 
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
         defaultCallback.expectAvailableCallbacksUnvalidatedAndBlocked(mCellNetworkAgent);
         defaultCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mCellNetworkAgent);
 
         // Allow to use the network after switching to NOT_METERED network.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.addCapability(NET_CAPABILITY_NOT_METERED);
         mWiFiNetworkAgent.connect(true);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
 
         // Switch to METERED network. Restrict the use of the network.
         mWiFiNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidatedAndBlocked(mCellNetworkAgent);
 
         // Network becomes NOT_METERED.
@@ -5990,12 +5415,12 @@
         defaultCallback.expectBlockedStatusCallback(false, mCellNetworkAgent);
 
         // Verify there's no Networkcallbacks invoked after data saver on/off.
-        mService.setRestrictBackgroundChanged(true);
-        mService.setRestrictBackgroundChanged(false);
+        setRestrictBackgroundChanged(true);
+        setRestrictBackgroundChanged(false);
         defaultCallback.assertNoCallback();
 
         mCellNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         defaultCallback.assertNoCallback();
 
         mCm.unregisterNetworkCallback(defaultCallback);
@@ -6027,7 +5452,7 @@
     }
 
     @Test
-    public void testStackedLinkProperties() throws UnknownHostException, RemoteException {
+    public void testStackedLinkProperties() throws Exception {
         final LinkAddress myIpv4 = new LinkAddress("1.2.3.4/24");
         final LinkAddress myIpv6 = new LinkAddress("2001:db8:1::1/64");
         final String kNat64PrefixString = "2001:db8:64:64:64:64::";
@@ -6041,7 +5466,7 @@
         mCm.registerNetworkCallback(networkRequest, networkCallback);
 
         // Prepare ipv6 only link properties.
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         final int cellNetId = mCellNetworkAgent.getNetwork().netId;
         final LinkProperties cellLp = new LinkProperties();
         cellLp.setInterfaceName(MOBILE_IFNAME);
@@ -6071,7 +5496,7 @@
         // the NAT64 prefix was removed because one was never discovered.
         cellLp.addLinkAddress(myIpv4);
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         verify(mMockDnsResolver, times(1)).stopPrefix64Discovery(cellNetId);
         verify(mMockDnsResolver, atLeastOnce()).setResolverConfiguration(any());
 
@@ -6084,23 +5509,23 @@
         cellLp.removeLinkAddress(myIpv4);
         cellLp.removeRoute(new RouteInfo(myIpv4, null, MOBILE_IFNAME));
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         verify(mMockDnsResolver, times(1)).startPrefix64Discovery(cellNetId);
 
         // When NAT64 prefix discovery succeeds, LinkProperties are updated and clatd is started.
-        Nat464Xlat clat = mService.getNat464Xlat(mCellNetworkAgent);
+        Nat464Xlat clat = getNat464Xlat(mCellNetworkAgent);
         assertNull(mCm.getLinkProperties(mCellNetworkAgent.getNetwork()).getNat64Prefix());
         mService.mNetdEventCallback.onNat64PrefixEvent(cellNetId, true /* added */,
                 kNat64PrefixString, 96);
-        LinkProperties lpBeforeClat = (LinkProperties) networkCallback.expectCallback(
-                CallbackState.LINK_PROPERTIES, mCellNetworkAgent).arg;
+        LinkProperties lpBeforeClat = networkCallback.expectCallback(
+                CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent).getLp();
         assertEquals(0, lpBeforeClat.getStackedLinks().size());
         assertEquals(kNat64Prefix, lpBeforeClat.getNat64Prefix());
         verify(mMockNetd, times(1)).clatdStart(MOBILE_IFNAME, kNat64Prefix.toString());
 
         // Clat iface comes up. Expect stacked link to be added.
         clat.interfaceLinkStateChanged(CLAT_PREFIX + MOBILE_IFNAME, true);
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         List<LinkProperties> stackedLps = mCm.getLinkProperties(mCellNetworkAgent.getNetwork())
                 .getStackedLinks();
         assertEquals(makeClatLinkProperties(myIpv4), stackedLps.get(0));
@@ -6108,7 +5533,7 @@
         // Change trivial linkproperties and see if stacked link is preserved.
         cellLp.addDnsServer(InetAddress.getByName("8.8.8.8"));
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
 
         List<LinkProperties> stackedLpsAfterChange =
                 mCm.getLinkProperties(mCellNetworkAgent.getNetwork()).getStackedLinks();
@@ -6126,12 +5551,12 @@
         cellLp.addLinkAddress(myIpv4);
         cellLp.addRoute(new RouteInfo(myIpv4, null, MOBILE_IFNAME));
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         verify(mMockNetd, times(1)).clatdStop(MOBILE_IFNAME);
         verify(mMockDnsResolver, times(1)).stopPrefix64Discovery(cellNetId);
 
         // As soon as stop is called, the linkproperties lose the stacked interface.
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         LinkProperties actualLpAfterIpv4 = mCm.getLinkProperties(mCellNetworkAgent.getNetwork());
         LinkProperties expected = new LinkProperties(cellLp);
         expected.setNat64Prefix(kNat64Prefix);
@@ -6150,54 +5575,52 @@
         // Stopping prefix discovery causes netd to tell us that the NAT64 prefix is gone.
         mService.mNetdEventCallback.onNat64PrefixEvent(cellNetId, false /* added */,
                 kNat64PrefixString, 96);
-        networkCallback.expectLinkPropertiesLike((lp) -> lp.getNat64Prefix() == null,
-                mCellNetworkAgent);
+        networkCallback.expectLinkPropertiesThat(mCellNetworkAgent,
+                (lp) -> lp.getNat64Prefix() == null);
 
         // Remove IPv4 address and expect prefix discovery and clatd to be started again.
         cellLp.removeLinkAddress(myIpv4);
         cellLp.removeRoute(new RouteInfo(myIpv4, null, MOBILE_IFNAME));
         cellLp.removeDnsServer(InetAddress.getByName("8.8.8.8"));
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         verify(mMockDnsResolver, times(1)).startPrefix64Discovery(cellNetId);
         mService.mNetdEventCallback.onNat64PrefixEvent(cellNetId, true /* added */,
                 kNat64PrefixString, 96);
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         verify(mMockNetd, times(1)).clatdStart(MOBILE_IFNAME, kNat64Prefix.toString());
 
 
         // Clat iface comes up. Expect stacked link to be added.
         clat.interfaceLinkStateChanged(CLAT_PREFIX + MOBILE_IFNAME, true);
-        networkCallback.expectLinkPropertiesLike(
-                (lp) -> lp.getStackedLinks().size() == 1 && lp.getNat64Prefix() != null,
-                mCellNetworkAgent);
+        networkCallback.expectLinkPropertiesThat(mCellNetworkAgent,
+                (lp) -> lp.getStackedLinks().size() == 1 && lp.getNat64Prefix() != null);
 
         // NAT64 prefix is removed. Expect that clat is stopped.
         mService.mNetdEventCallback.onNat64PrefixEvent(cellNetId, false /* added */,
                 kNat64PrefixString, 96);
-        networkCallback.expectLinkPropertiesLike(
-                (lp) -> lp.getStackedLinks().size() == 0 && lp.getNat64Prefix() == null,
-                mCellNetworkAgent);
+        networkCallback.expectLinkPropertiesThat(mCellNetworkAgent,
+                (lp) -> lp.getStackedLinks().size() == 0 && lp.getNat64Prefix() == null);
         verify(mMockNetd, times(1)).clatdStop(MOBILE_IFNAME);
-        networkCallback.expectLinkPropertiesLike((lp) -> lp.getStackedLinks().size() == 0,
-                mCellNetworkAgent);
+        networkCallback.expectLinkPropertiesThat(mCellNetworkAgent,
+                (lp) -> lp.getStackedLinks().size() == 0);
 
         // Clean up.
         mCellNetworkAgent.disconnect();
-        networkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         networkCallback.assertNoCallback();
         mCm.unregisterNetworkCallback(networkCallback);
     }
 
     @Test
-    public void testDataActivityTracking() throws RemoteException {
+    public void testDataActivityTracking() throws Exception {
         final TestNetworkCallback networkCallback = new TestNetworkCallback();
         final NetworkRequest networkRequest = new NetworkRequest.Builder()
                 .addCapability(NET_CAPABILITY_INTERNET)
                 .build();
         mCm.registerNetworkCallback(networkRequest, networkCallback);
 
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         final LinkProperties cellLp = new LinkProperties();
         cellLp.setInterfaceName(MOBILE_IFNAME);
         mCellNetworkAgent.sendLinkProperties(cellLp);
@@ -6207,7 +5630,7 @@
         verify(mNetworkManagementService, times(1)).addIdleTimer(eq(MOBILE_IFNAME), anyInt(),
                 eq(ConnectivityManager.TYPE_MOBILE));
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         final LinkProperties wifiLp = new LinkProperties();
         wifiLp.setInterfaceName(WIFI_IFNAME);
         mWiFiNetworkAgent.sendLinkProperties(wifiLp);
@@ -6216,7 +5639,7 @@
         reset(mNetworkManagementService);
         mWiFiNetworkAgent.connect(true);
         networkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        networkCallback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         networkCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         verify(mNetworkManagementService, times(1)).addIdleTimer(eq(WIFI_IFNAME), anyInt(),
                 eq(ConnectivityManager.TYPE_WIFI));
@@ -6225,26 +5648,26 @@
         // Disconnect wifi and switch back to cell
         reset(mNetworkManagementService);
         mWiFiNetworkAgent.disconnect();
-        networkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         assertNoCallbacks(networkCallback);
         verify(mNetworkManagementService, times(1)).removeIdleTimer(eq(WIFI_IFNAME));
         verify(mNetworkManagementService, times(1)).addIdleTimer(eq(MOBILE_IFNAME), anyInt(),
                 eq(ConnectivityManager.TYPE_MOBILE));
 
         // reconnect wifi
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         wifiLp.setInterfaceName(WIFI_IFNAME);
         mWiFiNetworkAgent.sendLinkProperties(wifiLp);
         mWiFiNetworkAgent.connect(true);
         networkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        networkCallback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         networkCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
 
         // Disconnect cell
         reset(mNetworkManagementService);
         reset(mMockNetd);
         mCellNetworkAgent.disconnect();
-        networkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         // LOST callback is triggered earlier than removing idle timer. Broadcast should also be
         // sent as network being switched. Ensure rule removal for cell will not be triggered
         // unexpectedly before network being removed.
@@ -6265,24 +5688,20 @@
         mCm.unregisterNetworkCallback(networkCallback);
     }
 
-    private void verifyTcpBufferSizeChange(String tcpBufferSizes) {
+    private void verifyTcpBufferSizeChange(String tcpBufferSizes) throws Exception {
         String[] values = tcpBufferSizes.split(",");
         String rmemValues = String.join(" ", values[0], values[1], values[2]);
         String wmemValues = String.join(" ", values[3], values[4], values[5]);
         waitForIdle();
-        try {
-            verify(mMockNetd, atLeastOnce()).setTcpRWmemorySize(rmemValues, wmemValues);
-        } catch (RemoteException e) {
-            fail("mMockNetd should never throw RemoteException");
-        }
+        verify(mMockNetd, atLeastOnce()).setTcpRWmemorySize(rmemValues, wmemValues);
         reset(mMockNetd);
     }
 
     @Test
-    public void testTcpBufferReset() {
+    public void testTcpBufferReset() throws Exception {
         final String testTcpBufferSizes = "1,2,3,4,5,6";
 
-        mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         reset(mMockNetd);
         // Switching default network updates TCP buffer sizes.
         mCellNetworkAgent.connect(false);
@@ -6296,18 +5715,18 @@
     }
 
     @Test
-    public void testGetGlobalProxyForNetwork() {
+    public void testGetGlobalProxyForNetwork() throws Exception {
         final ProxyInfo testProxyInfo = ProxyInfo.buildDirectProxy("test", 8888);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         final Network wifiNetwork = mWiFiNetworkAgent.getNetwork();
         when(mService.mProxyTracker.getGlobalProxy()).thenReturn(testProxyInfo);
         assertEquals(testProxyInfo, mService.getProxyForNetwork(wifiNetwork));
     }
 
     @Test
-    public void testGetProxyForActiveNetwork() {
+    public void testGetProxyForActiveNetwork() throws Exception {
         final ProxyInfo testProxyInfo = ProxyInfo.buildDirectProxy("test", 8888);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
         waitForIdle();
         assertNull(mService.getProxyForNetwork(null));
@@ -6322,18 +5741,19 @@
     }
 
     @Test
-    public void testGetProxyForVPN() {
+    public void testGetProxyForVPN() throws Exception {
         final ProxyInfo testProxyInfo = ProxyInfo.buildDirectProxy("test", 8888);
 
         // Set up a WiFi network with no proxy
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
         waitForIdle();
         assertNull(mService.getProxyForNetwork(null));
 
         // Set up a VPN network with a proxy
         final int uid = Process.myUid();
-        final MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        final TestNetworkAgentWrapper
+                vpnNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN);
         final ArraySet<UidRange> ranges = new ArraySet<>();
         ranges.add(new UidRange(uid, uid));
         mMockVpn.setUids(ranges);
@@ -6382,7 +5802,7 @@
         lp.addRoute(new RouteInfo(new IpPrefix(Inet4Address.ANY, 0), null));
         // The uid range needs to cover the test app so the network is visible to it.
         final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(VPN_USER));
-        final MockNetworkAgent vpnNetworkAgent = establishVpn(lp, VPN_UID, vpnRange);
+        final TestNetworkAgentWrapper vpnNetworkAgent = establishVpn(lp, VPN_UID, vpnRange);
 
         // Connected VPN should have interface rules set up. There are two expected invocations,
         // one during VPN uid update, one during VPN LinkProperties update
@@ -6408,7 +5828,8 @@
         lp.addRoute(new RouteInfo(new IpPrefix(Inet4Address.ANY, 0), null));
         // The uid range needs to cover the test app so the network is visible to it.
         final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(VPN_USER));
-        final MockNetworkAgent vpnNetworkAgent = establishVpn(lp, Process.SYSTEM_UID, vpnRange);
+        final TestNetworkAgentWrapper vpnNetworkAgent = establishVpn(
+                lp, Process.SYSTEM_UID, vpnRange);
 
         // Legacy VPN should not have interface rules set up
         verify(mMockNetd, never()).firewallAddUidInterfaceRules(any(), any());
@@ -6423,7 +5844,8 @@
         lp.addRoute(new RouteInfo(new IpPrefix(Inet6Address.ANY, 0), RTN_UNREACHABLE));
         // The uid range needs to cover the test app so the network is visible to it.
         final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(VPN_USER));
-        final MockNetworkAgent vpnNetworkAgent = establishVpn(lp, Process.SYSTEM_UID, vpnRange);
+        final TestNetworkAgentWrapper vpnNetworkAgent = establishVpn(
+                lp, Process.SYSTEM_UID, vpnRange);
 
         // IPv6 unreachable route should not be misinterpreted as a default route
         verify(mMockNetd, never()).firewallAddUidInterfaceRules(any(), any());
@@ -6436,7 +5858,7 @@
         lp.addRoute(new RouteInfo(new IpPrefix(Inet4Address.ANY, 0), null));
         // The uid range needs to cover the test app so the network is visible to it.
         final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(VPN_USER));
-        final MockNetworkAgent vpnNetworkAgent = establishVpn(lp, VPN_UID, vpnRange);
+        final TestNetworkAgentWrapper vpnNetworkAgent = establishVpn(lp, VPN_UID, vpnRange);
 
         // Connected VPN should have interface rules set up. There are two expected invocations,
         // one during VPN uid update, one during VPN LinkProperties update
@@ -6485,7 +5907,7 @@
         lp.addRoute(new RouteInfo(new IpPrefix(Inet6Address.ANY, 0), null));
         // The uid range needs to cover the test app so the network is visible to it.
         final UidRange vpnRange = UidRange.createForUser(VPN_USER);
-        final MockNetworkAgent vpnNetworkAgent = establishVpn(lp, VPN_UID,
+        final TestNetworkAgentWrapper vpnNetworkAgent = establishVpn(lp, VPN_UID,
                 Collections.singleton(vpnRange));
 
         reset(mMockNetd);
@@ -6507,9 +5929,10 @@
     }
 
 
-    private MockNetworkAgent establishVpn(LinkProperties lp, int establishingUid,
-            Set<UidRange> vpnRange) {
-        final MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN, lp);
+    private TestNetworkAgentWrapper establishVpn(LinkProperties lp, int establishingUid,
+            Set<UidRange> vpnRange) throws Exception {
+        final TestNetworkAgentWrapper
+                vpnNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN, lp);
         vpnNetworkAgent.getNetworkCapabilities().setEstablishingVpnAppUid(establishingUid);
         mMockVpn.setNetworkAgent(vpnNetworkAgent);
         mMockVpn.connect();
@@ -6519,14 +5942,6 @@
         return vpnNetworkAgent;
     }
 
-    private void assertContainsExactly(int[] actual, int... expected) {
-        int[] sortedActual = Arrays.copyOf(actual, actual.length);
-        int[] sortedExpected = Arrays.copyOf(expected, expected.length);
-        Arrays.sort(sortedActual);
-        Arrays.sort(sortedExpected);
-        assertArrayEquals(sortedExpected, sortedActual);
-    }
-
     private static PackageInfo buildPackageInfo(boolean hasSystemPermission, int uid) {
         final PackageInfo packageInfo = new PackageInfo();
         packageInfo.requestedPermissions = new String[0];
diff --git a/tests/net/java/com/android/server/connectivity/NetdEventListenerServiceTest.java b/tests/net/java/com/android/server/connectivity/NetdEventListenerServiceTest.java
index e4117b8..aef9386 100644
--- a/tests/net/java/com/android/server/connectivity/NetdEventListenerServiceTest.java
+++ b/tests/net/java/com/android/server/connectivity/NetdEventListenerServiceTest.java
@@ -19,8 +19,9 @@
 import static android.net.metrics.INetdEventListener.EVENT_GETADDRINFO;
 import static android.net.metrics.INetdEventListener.EVENT_GETHOSTBYNAME;
 
+import static com.android.testutils.MiscAssertsKt.assertStringContains;
+
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -111,15 +112,15 @@
         String[] events2 = remove(listNetdEvent(), baseline);
         int expectedLength2 = uids.length + 1; // +1 for the WakeupStats line
         assertEquals(expectedLength2, events2.length);
-        assertContains(events2[0], "WakeupStats");
-        assertContains(events2[0], "wlan0");
-        assertContains(events2[0], "0x800");
-        assertContains(events2[0], "0x86dd");
+        assertStringContains(events2[0], "WakeupStats");
+        assertStringContains(events2[0], "wlan0");
+        assertStringContains(events2[0], "0x800");
+        assertStringContains(events2[0], "0x86dd");
         for (int i = 0; i < uids.length; i++) {
             String got = events2[i+1];
-            assertContains(got, "WakeupEvent");
-            assertContains(got, "wlan0");
-            assertContains(got, "uid: " + uids[i]);
+            assertStringContains(got, "WakeupEvent");
+            assertStringContains(got, "wlan0");
+            assertStringContains(got, "uid: " + uids[i]);
         }
 
         int uid = 20000;
@@ -131,13 +132,13 @@
         String[] events3 = remove(listNetdEvent(), baseline);
         int expectedLength3 = BUFFER_LENGTH + 1; // +1 for the WakeupStats line
         assertEquals(expectedLength3, events3.length);
-        assertContains(events2[0], "WakeupStats");
-        assertContains(events2[0], "wlan0");
+        assertStringContains(events2[0], "WakeupStats");
+        assertStringContains(events2[0], "wlan0");
         for (int i = 1; i < expectedLength3; i++) {
             String got = events3[i];
-            assertContains(got, "WakeupEvent");
-            assertContains(got, "wlan0");
-            assertContains(got, "uid: " + uid);
+            assertStringContains(got, "WakeupEvent");
+            assertStringContains(got, "wlan0");
+            assertStringContains(got, "uid: " + uid);
         }
 
         uid = 45678;
@@ -145,9 +146,9 @@
 
         String[] events4 = remove(listNetdEvent(), baseline);
         String lastEvent = events4[events4.length - 1];
-        assertContains(lastEvent, "WakeupEvent");
-        assertContains(lastEvent, "wlan0");
-        assertContains(lastEvent, "uid: " + uid);
+        assertStringContains(lastEvent, "WakeupEvent");
+        assertStringContains(lastEvent, "wlan0");
+        assertStringContains(lastEvent, "uid: " + uid);
     }
 
     @Test
@@ -529,10 +530,6 @@
         return buffer.toString().split("\\n");
     }
 
-    static void assertContains(String got, String want) {
-        assertTrue(got + " did not contain \"" + want + "\"", got.contains(want));
-    }
-
     static <T> T[] remove(T[] array, T[] filtered) {
         List<T> c = Arrays.asList(filtered);
         int next = 0;
diff --git a/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java b/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java
index 9b4f49c..8f90f13 100644
--- a/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java
+++ b/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java
@@ -29,6 +29,7 @@
 
 import static com.android.server.net.NetworkStatsCollection.multiplySafe;
 
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.fail;
@@ -43,7 +44,6 @@
 import android.os.UserHandle;
 import android.telephony.SubscriptionPlan;
 import android.telephony.TelephonyManager;
-import android.test.MoreAsserts;
 import android.text.format.DateUtils;
 import android.util.RecurrenceRule;
 
@@ -240,11 +240,11 @@
                 60 * MINUTE_IN_MILLIS, entry);
 
         // Verify the set of relevant UIDs for each access level.
-        MoreAsserts.assertEquals(new int[] { myUid },
+        assertArrayEquals(new int[] { myUid },
                 collection.getRelevantUids(NetworkStatsAccess.Level.DEFAULT));
-        MoreAsserts.assertEquals(new int[] { Process.SYSTEM_UID, myUid, otherUidInSameUser },
+        assertArrayEquals(new int[] { Process.SYSTEM_UID, myUid, otherUidInSameUser },
                 collection.getRelevantUids(NetworkStatsAccess.Level.USER));
-        MoreAsserts.assertEquals(
+        assertArrayEquals(
                 new int[] { Process.SYSTEM_UID, myUid, otherUidInSameUser, uidInDifferentUser },
                 collection.getRelevantUids(NetworkStatsAccess.Level.DEVICE));
 
diff --git a/tests/net/java/com/android/server/net/NetworkStatsFactoryTest.java b/tests/net/java/com/android/server/net/NetworkStatsFactoryTest.java
index 7329474..a21f509 100644
--- a/tests/net/java/com/android/server/net/NetworkStatsFactoryTest.java
+++ b/tests/net/java/com/android/server/net/NetworkStatsFactoryTest.java
@@ -79,8 +79,7 @@
         // related to networkStatsFactory is compiled to a minimal native library and loaded here.
         System.loadLibrary("networkstatsfactorytestjni");
         mFactory = new NetworkStatsFactory(mTestProc, false);
-        NetworkStatsFactory.updateVpnInfos(new VpnInfo[0]);
-        NetworkStatsFactory.clearStackedIfaces();
+        mFactory.updateVpnInfos(new VpnInfo[0]);
     }
 
     @After
@@ -107,7 +106,7 @@
     @Test
     public void vpnRewriteTrafficThroughItself() throws Exception {
         VpnInfo[] vpnInfos = new VpnInfo[] {createVpnInfo(new String[] {TEST_IFACE})};
-        NetworkStatsFactory.updateVpnInfos(vpnInfos);
+        mFactory.updateVpnInfos(vpnInfos);
 
         // create some traffic (assume 10 bytes of MTU for VPN interface and 1 byte encryption
         // overhead per packet):
@@ -136,8 +135,8 @@
     @Test
     public void vpnWithClat() throws Exception {
         VpnInfo[] vpnInfos = new VpnInfo[] {createVpnInfo(new String[] {CLAT_PREFIX + TEST_IFACE})};
-        NetworkStatsFactory.updateVpnInfos(vpnInfos);
-        NetworkStatsFactory.noteStackedIface(CLAT_PREFIX + TEST_IFACE, TEST_IFACE);
+        mFactory.updateVpnInfos(vpnInfos);
+        mFactory.noteStackedIface(CLAT_PREFIX + TEST_IFACE, TEST_IFACE);
 
         // create some traffic (assume 10 bytes of MTU for VPN interface and 1 byte encryption
         // overhead per packet):
@@ -169,7 +168,7 @@
     @Test
     public void vpnWithOneUnderlyingIface() throws Exception {
         VpnInfo[] vpnInfos = new VpnInfo[] {createVpnInfo(new String[] {TEST_IFACE})};
-        NetworkStatsFactory.updateVpnInfos(vpnInfos);
+        mFactory.updateVpnInfos(vpnInfos);
 
         // create some traffic (assume 10 bytes of MTU for VPN interface and 1 byte encryption
         // overhead per packet):
@@ -193,7 +192,7 @@
     public void vpnWithOneUnderlyingIfaceAndOwnTraffic() throws Exception {
         // WiFi network is connected and VPN is using WiFi (which has TEST_IFACE).
         VpnInfo[] vpnInfos = new VpnInfo[] {createVpnInfo(new String[] {TEST_IFACE})};
-        NetworkStatsFactory.updateVpnInfos(vpnInfos);
+        mFactory.updateVpnInfos(vpnInfos);
 
         // create some traffic (assume 10 bytes of MTU for VPN interface and 1 byte encryption
         // overhead per packet):
@@ -221,7 +220,7 @@
     public void vpnWithOneUnderlyingIface_withCompression() throws Exception {
         // WiFi network is connected and VPN is using WiFi (which has TEST_IFACE).
         VpnInfo[] vpnInfos = new VpnInfo[] {createVpnInfo(new String[] {TEST_IFACE})};
-        NetworkStatsFactory.updateVpnInfos(vpnInfos);
+        mFactory.updateVpnInfos(vpnInfos);
 
         // create some traffic (assume 10 bytes of MTU for VPN interface and 1 byte encryption
         // overhead per packet):
@@ -244,7 +243,7 @@
         // Cell (which has TEST_IFACE2) and has declared both of them in its underlying network set.
         // Additionally, VPN is duplicating traffic across both WiFi and Cell.
         VpnInfo[] vpnInfos = new VpnInfo[] {createVpnInfo(new String[] {TEST_IFACE, TEST_IFACE2})};
-        NetworkStatsFactory.updateVpnInfos(vpnInfos);
+        mFactory.updateVpnInfos(vpnInfos);
 
         // create some traffic (assume 10 bytes of MTU for VPN interface and 1 byte encryption
         // overhead per packet):
@@ -270,7 +269,7 @@
         // Cell (which has TEST_IFACE2) and has declared both of them in its underlying network set.
         // Additionally, VPN is arbitrarily splitting traffic across WiFi and Cell.
         VpnInfo[] vpnInfos = new VpnInfo[] {createVpnInfo(new String[] {TEST_IFACE, TEST_IFACE2})};
-        NetworkStatsFactory.updateVpnInfos(vpnInfos);
+        mFactory.updateVpnInfos(vpnInfos);
 
         // create some traffic (assume 10 bytes of MTU for VPN interface and 1 byte encryption
         // overhead per packet):
@@ -297,7 +296,7 @@
         // Cell (which has TEST_IFACE2) and has declared both of them in its underlying network set.
         // Additionally, VPN is arbitrarily splitting compressed traffic across WiFi and Cell.
         VpnInfo[] vpnInfos = new VpnInfo[] {createVpnInfo(new String[] {TEST_IFACE, TEST_IFACE2})};
-        NetworkStatsFactory.updateVpnInfos(vpnInfos);
+        mFactory.updateVpnInfos(vpnInfos);
 
         // create some traffic (assume 10 bytes of MTU for VPN interface:
         // 1000 bytes (100 packets) were sent/received by UID_RED over VPN.
@@ -319,7 +318,7 @@
         // WiFi and Cell networks are connected and VPN is using Cell (which has TEST_IFACE2),
         // but has declared only WiFi (TEST_IFACE) in its underlying network set.
         VpnInfo[] vpnInfos = new VpnInfo[] {createVpnInfo(new String[] {TEST_IFACE})};
-        NetworkStatsFactory.updateVpnInfos(vpnInfos);
+        mFactory.updateVpnInfos(vpnInfos);
 
         // create some traffic (assume 10 bytes of MTU for VPN interface and 1 byte encryption
         // overhead per packet):
@@ -383,7 +382,7 @@
 
     @Test
     public void testDoubleClatAccountingSimple() throws Exception {
-        NetworkStatsFactory.noteStackedIface("v4-wlan0", "wlan0");
+        mFactory.noteStackedIface("v4-wlan0", "wlan0");
 
         // xt_qtaguid_with_clat_simple is a synthetic file that simulates
         //  - 213 received 464xlat packets of size 200 bytes
@@ -398,7 +397,7 @@
 
     @Test
     public void testDoubleClatAccounting() throws Exception {
-        NetworkStatsFactory.noteStackedIface("v4-wlan0", "wlan0");
+        mFactory.noteStackedIface("v4-wlan0", "wlan0");
 
         NetworkStats stats = parseDetailedStats(R.raw.xt_qtaguid_with_clat);
         assertEquals(42, stats.size());
@@ -419,8 +418,6 @@
         assertStatsEntry(stats, "lo", 0, SET_DEFAULT, 0x0, 1288L, 1288L);
 
         assertNoStatsEntry(stats, "wlan0", 1029, SET_DEFAULT, 0x0);
-
-        NetworkStatsFactory.clearStackedIfaces();
     }
 
     @Test
@@ -435,7 +432,7 @@
         long rootRxBytesAfter = 1398634L;
         assertEquals("UID 0 traffic should be ~0", 4623, rootRxBytesAfter - rootRxBytesBefore);
 
-        NetworkStatsFactory.noteStackedIface("v4-wlan0", "wlan0");
+        mFactory.noteStackedIface("v4-wlan0", "wlan0");
         NetworkStats stats;
 
         // Stats snapshot before the download
@@ -447,8 +444,6 @@
         stats = parseDetailedStats(R.raw.xt_qtaguid_with_clat_100mb_download_after);
         assertStatsEntry(stats, "v4-wlan0", 10106, SET_FOREGROUND, 0x0, appRxBytesAfter, 7867488L);
         assertStatsEntry(stats, "wlan0", 0, SET_DEFAULT, 0x0, rootRxBytesAfter, 0L);
-
-        NetworkStatsFactory.clearStackedIfaces();
     }
 
     /**
diff --git a/tests/net/java/com/android/server/net/NetworkStatsObserversTest.java b/tests/net/java/com/android/server/net/NetworkStatsObserversTest.java
index f21a7dd..c0f9dc1 100644
--- a/tests/net/java/com/android/server/net/NetworkStatsObserversTest.java
+++ b/tests/net/java/com/android/server/net/NetworkStatsObserversTest.java
@@ -28,8 +28,6 @@
 import static android.net.TrafficStats.MB_IN_BYTES;
 import static android.text.format.DateUtils.MINUTE_IN_MILLIS;
 
-import static com.android.internal.util.TestUtils.waitForIdleHandler;
-
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Matchers.any;
@@ -55,6 +53,7 @@
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.server.net.NetworkStatsServiceTest.LatchedHandler;
+import com.android.testutils.HandlerUtilsKt;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -441,7 +440,7 @@
     }
 
     private void waitForObserverToIdle() {
-        waitForIdleHandler(mObserverHandlerThread, WAIT_TIMEOUT_MS);
-        waitForIdleHandler(mHandler, WAIT_TIMEOUT_MS);
+        HandlerUtilsKt.waitForIdle(mObserverHandlerThread, WAIT_TIMEOUT_MS);
+        HandlerUtilsKt.waitForIdle(mHandler, WAIT_TIMEOUT_MS);
     }
 }
diff --git a/tests/net/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/net/java/com/android/server/net/NetworkStatsServiceTest.java
index 956b2a7..1d29a82 100644
--- a/tests/net/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/net/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -37,6 +37,7 @@
 import static android.net.NetworkStats.SET_FOREGROUND;
 import static android.net.NetworkStats.STATS_PER_IFACE;
 import static android.net.NetworkStats.STATS_PER_UID;
+import static android.net.NetworkStats.TAG_ALL;
 import static android.net.NetworkStats.TAG_NONE;
 import static android.net.NetworkStats.UID_ALL;
 import static android.net.NetworkStatsHistory.FIELD_ALL;
@@ -50,7 +51,6 @@
 import static android.text.format.DateUtils.MINUTE_IN_MILLIS;
 import static android.text.format.DateUtils.WEEK_IN_MILLIS;
 
-import static com.android.internal.util.TestUtils.waitForIdleHandler;
 import static com.android.server.net.NetworkStatsService.ACTION_NETWORK_STATS_POLL;
 
 import static org.junit.Assert.assertEquals;
@@ -96,10 +96,12 @@
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.internal.net.VpnInfo;
 import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.test.BroadcastInterceptingContext;
 import com.android.server.net.NetworkStatsService.NetworkStatsSettings;
 import com.android.server.net.NetworkStatsService.NetworkStatsSettings.Config;
+import com.android.testutils.HandlerUtilsKt;
 
 import libcore.io.IoUtils;
 
@@ -154,6 +156,7 @@
     private File mStatsDir;
 
     private @Mock INetworkManagementService mNetManager;
+    private @Mock NetworkStatsFactory mStatsFactory;
     private @Mock NetworkStatsSettings mSettings;
     private @Mock IBinder mBinder;
     private @Mock AlarmManager mAlarmManager;
@@ -189,8 +192,8 @@
 
         mService = new NetworkStatsService(
                 mServiceContext, mNetManager, mAlarmManager, wakeLock, mClock,
-                TelephonyManager.getDefault(), mSettings, new NetworkStatsObservers(),
-                mStatsDir, getBaseDir(mStatsDir));
+                TelephonyManager.getDefault(), mSettings, mStatsFactory,
+                new NetworkStatsObservers(),  mStatsDir, getBaseDir(mStatsDir));
         mHandlerThread = new HandlerThread("HandlerThread");
         mHandlerThread.start();
         Handler.Callback callback = new NetworkStatsService.HandlerCallback(mService);
@@ -205,7 +208,7 @@
 
         mService.systemReady();
         // Verify that system ready fetches realtime stats
-        verify(mNetManager).getNetworkStatsUidDetail(UID_ALL, INTERFACES_ALL);
+        verify(mStatsFactory).readNetworkStatsDetail(UID_ALL, INTERFACES_ALL, TAG_ALL);
 
         mSession = mService.openSession();
         assertNotNull("openSession() failed", mSession);
@@ -239,9 +242,8 @@
         NetworkState[] states = new NetworkState[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states), new VpnInfo[0]);
 
         // verify service has empty history for wifi
         assertNetworkTotal(sTemplateWifi, 0L, 0L, 0L, 0L, 0);
@@ -283,9 +285,8 @@
         NetworkState[] states = new NetworkState[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states), new VpnInfo[0]);
 
         // verify service has empty history for wifi
         assertNetworkTotal(sTemplateWifi, 0L, 0L, 0L, 0L, 0);
@@ -357,9 +358,8 @@
         NetworkState[] states = new NetworkState[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states), new VpnInfo[0]);
 
         // modify some number on wifi, and trigger poll event
         incrementCurrentTime(2 * HOUR_IN_MILLIS);
@@ -398,9 +398,8 @@
         NetworkState[] states = new NetworkState[] {buildMobile3gState(IMSI_1)};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_MOBILE, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_MOBILE, states, getActiveIface(states), new VpnInfo[0]);
 
         // create some traffic on first network
         incrementCurrentTime(HOUR_IN_MILLIS);
@@ -433,9 +432,8 @@
                 .addValues(TEST_IFACE, UID_RED, SET_DEFAULT, TAG_NONE, 1536L, 12L, 512L, 4L, 0L)
                 .addValues(TEST_IFACE, UID_RED, SET_DEFAULT, 0xF00D, 512L, 4L, 512L, 4L, 0L)
                 .addValues(TEST_IFACE, UID_BLUE, SET_DEFAULT, TAG_NONE, 512L, 4L, 0L, 0L, 0L));
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_MOBILE, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_MOBILE, states, getActiveIface(states), new VpnInfo[0]);
         forcePollAndWaitForIdle();
 
 
@@ -473,9 +471,8 @@
         NetworkState[] states = new NetworkState[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states), new VpnInfo[0]);
 
         // create some traffic
         incrementCurrentTime(HOUR_IN_MILLIS);
@@ -531,9 +528,8 @@
         NetworkState[] states = new NetworkState[] {buildMobile3gState(IMSI_1)};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_MOBILE, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_MOBILE, states, getActiveIface(states), new VpnInfo[0]);
 
         // create some traffic
         incrementCurrentTime(HOUR_IN_MILLIS);
@@ -558,9 +554,8 @@
         expectNetworkStatsUidDetail(new NetworkStats(getElapsedRealtime(), 1)
                 .addValues(TEST_IFACE, UID_RED, SET_DEFAULT, TAG_NONE, 1024L, 8L, 1024L, 8L, 0L)
                 .addValues(TEST_IFACE, UID_RED, SET_DEFAULT, 0xF00D, 512L, 4L, 512L, 4L, 0L));
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_MOBILE, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_MOBILE, states, getActiveIface(states), new VpnInfo[0]);
         forcePollAndWaitForIdle();
 
 
@@ -588,9 +583,8 @@
         NetworkState[] states = new NetworkState[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states), new VpnInfo[0]);
 
         // create some traffic for two apps
         incrementCurrentTime(HOUR_IN_MILLIS);
@@ -646,9 +640,8 @@
         NetworkState[] states = new NetworkState[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states), new VpnInfo[0]);
 
         NetworkStats.Entry entry1 = new NetworkStats.Entry(
                 TEST_IFACE, UID_RED, SET_DEFAULT, TAG_NONE, 50L, 5L, 50L, 5L, 0L);
@@ -690,9 +683,8 @@
 
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states), new VpnInfo[0]);
 
         NetworkStats.Entry uidStats = new NetworkStats.Entry(
                 TEST_IFACE, UID_BLUE, SET_DEFAULT, 0xF00D, 1024L, 8L, 512L, 4L, 0L);
@@ -704,10 +696,13 @@
                 "otherif", UID_BLUE, SET_DEFAULT, 0xF00D, 1024L, 8L, 512L, 4L, 0L);
 
         final String[] ifaceFilter = new String[] { TEST_IFACE };
+        final String[] augmentedIfaceFilter = new String[] { stackedIface, TEST_IFACE };
         incrementCurrentTime(HOUR_IN_MILLIS);
         expectDefaultSettings();
         expectNetworkStatsSummary(buildEmptyStats());
-        when(mNetManager.getNetworkStatsUidDetail(eq(UID_ALL), any()))
+        when(mStatsFactory.augmentWithStackedInterfaces(eq(ifaceFilter)))
+                .thenReturn(augmentedIfaceFilter);
+        when(mStatsFactory.readNetworkStatsDetail(eq(UID_ALL), any(), eq(TAG_ALL)))
                 .thenReturn(new NetworkStats(getElapsedRealtime(), 1)
                         .addValues(uidStats));
         when(mNetManager.getNetworkStatsTethering(STATS_PER_UID))
@@ -717,15 +712,17 @@
 
         NetworkStats stats = mService.getDetailedUidStats(ifaceFilter);
 
-        // mNetManager#getNetworkStatsUidDetail(UID_ALL, INTERFACES_ALL) has following invocations:
+        // mStatsFactory#readNetworkStatsDetail() has the following invocations:
         // 1) NetworkStatsService#systemReady from #setUp.
         // 2) mService#forceUpdateIfaces in the test above.
         //
         // Additionally, we should have one call from the above call to mService#getDetailedUidStats
-        // with the augmented ifaceFilter
-        verify(mNetManager, times(2)).getNetworkStatsUidDetail(UID_ALL, INTERFACES_ALL);
-        verify(mNetManager, times(1)).getNetworkStatsUidDetail(
-                eq(UID_ALL), eq(NetworkStatsFactory.augmentWithStackedInterfaces(ifaceFilter)));
+        // with the augmented ifaceFilter.
+        verify(mStatsFactory, times(2)).readNetworkStatsDetail(UID_ALL, INTERFACES_ALL, TAG_ALL);
+        verify(mStatsFactory, times(1)).readNetworkStatsDetail(
+                eq(UID_ALL),
+                eq(augmentedIfaceFilter),
+                eq(TAG_ALL));
         assertTrue(ArrayUtils.contains(stats.getUniqueIfaces(), TEST_IFACE));
         assertTrue(ArrayUtils.contains(stats.getUniqueIfaces(), stackedIface));
         assertEquals(2, stats.size());
@@ -740,9 +737,8 @@
         NetworkState[] states = new NetworkState[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states), new VpnInfo[0]);
 
         // create some initial traffic
         incrementCurrentTime(HOUR_IN_MILLIS);
@@ -797,9 +793,8 @@
         NetworkState[] states = new NetworkState[] {buildWifiState(true /* isMetered */)};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states), new VpnInfo[0]);
 
         // create some initial traffic
         incrementCurrentTime(HOUR_IN_MILLIS);
@@ -837,9 +832,8 @@
             new NetworkState[] {buildMobile3gState(IMSI_1, true /* isRoaming */)};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_MOBILE, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_MOBILE, states, getActiveIface(states), new VpnInfo[0]);
 
         // Create some traffic
         incrementCurrentTime(HOUR_IN_MILLIS);
@@ -875,9 +869,8 @@
         NetworkState[] states = new NetworkState[] {buildMobile3gState(IMSI_1)};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_MOBILE, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_MOBILE, states, getActiveIface(states), new VpnInfo[0]);
 
         // create some tethering traffic
         incrementCurrentTime(HOUR_IN_MILLIS);
@@ -916,9 +909,8 @@
         NetworkState[] states = new NetworkState[] {buildWifiState()};
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
-        expectBandwidthControlCheck();
 
-        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states));
+        mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states), new VpnInfo[0]);
 
         // verify service has empty history for wifi
         assertNetworkTotal(sTemplateWifi, 0L, 0L, 0L, 0L, 0);
@@ -936,8 +928,6 @@
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
 
-
-
         // Register and verify request and that binder was called
         DataUsageRequest request =
                 mService.registerUsageCallback(mServiceContext.getOpPackageName(), inputRequest,
@@ -949,14 +939,11 @@
 
         // Send dummy message to make sure that any previous message has been handled
         mHandler.sendMessage(mHandler.obtainMessage(-1));
-        waitForIdleHandler(mHandler, WAIT_TIMEOUT);
-
-
+        HandlerUtilsKt.waitForIdle(mHandler, WAIT_TIMEOUT);
 
         // Make sure that the caller binder gets connected
         verify(mBinder).linkToDeath(any(IBinder.DeathRecipient.class), anyInt());
 
-
         // modify some number on wifi, and trigger poll event
         // not enough traffic to call data usage callback
         incrementCurrentTime(HOUR_IN_MILLIS);
@@ -1062,7 +1049,6 @@
 
     private void expectSystemReady() throws Exception {
         expectNetworkStatsSummary(buildEmptyStats());
-        expectBandwidthControlCheck();
     }
 
     private String getActiveIface(NetworkState... states) throws Exception {
@@ -1084,11 +1070,11 @@
     }
 
     private void expectNetworkStatsSummaryDev(NetworkStats summary) throws Exception {
-        when(mNetManager.getNetworkStatsSummaryDev()).thenReturn(summary);
+        when(mStatsFactory.readNetworkStatsSummaryDev()).thenReturn(summary);
     }
 
     private void expectNetworkStatsSummaryXt(NetworkStats summary) throws Exception {
-        when(mNetManager.getNetworkStatsSummaryXt()).thenReturn(summary);
+        when(mStatsFactory.readNetworkStatsSummaryXt()).thenReturn(summary);
     }
 
     private void expectNetworkStatsTethering(int how, NetworkStats stats)
@@ -1102,7 +1088,8 @@
 
     private void expectNetworkStatsUidDetail(NetworkStats detail, NetworkStats tetherStats)
             throws Exception {
-        when(mNetManager.getNetworkStatsUidDetail(UID_ALL, INTERFACES_ALL)).thenReturn(detail);
+        when(mStatsFactory.readNetworkStatsDetail(UID_ALL, INTERFACES_ALL, TAG_ALL))
+                .thenReturn(detail);
 
         // also include tethering details, since they are folded into UID
         when(mNetManager.getNetworkStatsTethering(STATS_PER_UID)).thenReturn(tetherStats);
@@ -1130,10 +1117,6 @@
         when(mSettings.getUidTagPersistBytes(anyLong())).thenReturn(MB_IN_BYTES);
     }
 
-    private void expectBandwidthControlCheck() throws Exception {
-        when(mNetManager.isBandwidthControlEnabled()).thenReturn(true);
-    }
-
     private void assertStatsFilesExist(boolean exist) {
         final File basePath = new File(mStatsDir, "netstats");
         if (exist) {
@@ -1228,7 +1211,7 @@
         mServiceContext.sendBroadcast(new Intent(ACTION_NETWORK_STATS_POLL));
         // Send dummy message to make sure that any previous message has been handled
         mHandler.sendMessage(mHandler.obtainMessage(-1));
-        waitForIdleHandler(mHandler, WAIT_TIMEOUT);
+        HandlerUtilsKt.waitForIdle(mHandler, WAIT_TIMEOUT);
     }
 
     static class LatchedHandler extends Handler {
diff --git a/tests/net/smoketest/Android.bp b/tests/net/smoketest/Android.bp
index ef1ad2c..84ae2b5 100644
--- a/tests/net/smoketest/Android.bp
+++ b/tests/net/smoketest/Android.bp
@@ -14,4 +14,9 @@
     defaults: ["FrameworksNetTests-jni-defaults"],
     srcs: ["java/SmokeTest.java"],
     test_suites: ["device-tests"],
-}
+    static_libs: [
+        "androidx.test.rules",
+        "mockito-target-minus-junit4",
+        "services.core",
+    ],
+}
\ No newline at end of file