diff --git a/tests/net/Android.bp b/tests/net/Android.bp
index 5260b30..502aa97 100644
--- a/tests/net/Android.bp
+++ b/tests/net/Android.bp
@@ -56,6 +56,7 @@
         "androidx.test.rules",
         "FrameworksNetCommonTests",
         "frameworks-base-testutils",
+        "frameworks-net-integration-testutils",
         "framework-protos",
         "mockito-target-minus-junit4",
         "net-tests-utils",
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/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index e08cb16..a028a54 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,6 +67,7 @@
 import static android.net.NetworkPolicyManager.RULE_REJECT_METERED;
 import static android.net.RouteInfo.RTN_UNREACHABLE;
 
+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;
@@ -86,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;
@@ -93,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;
@@ -105,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;
@@ -125,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;
@@ -139,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;
@@ -168,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;
@@ -265,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";
@@ -276,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;
@@ -296,6 +300,7 @@
     @Mock PackageManager mPackageManager;
     @Mock UserManager mUserManager;
     @Mock NotificationManager mNotificationManager;
+    @Mock AlarmManager mAlarmManager;
 
     private ArgumentCaptor<ResolverParamsParcel> mResolverParamsParcelCaptor =
             ArgumentCaptor.forClass(ResolverParamsParcel.class);
@@ -368,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);
         }
 
@@ -397,25 +403,20 @@
         }
     }
 
-    public void waitForIdle(int timeoutMsAsInt) {
-        long timeoutMs = timeoutMsAsInt;
-        HandlerUtilsKt.waitForIdle(mService.mHandlerThread, timeoutMs);
-        waitForIdle(mCellNetworkAgent, timeoutMs);
-        waitForIdle(mWiFiNetworkAgent, timeoutMs);
-        waitForIdle(mEthernetNetworkAgent, timeoutMs);
-        HandlerUtilsKt.waitForIdle(mService.mHandlerThread, timeoutMs);
-        HandlerUtilsKt.waitForIdle(ConnectivityThread.get(), 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;
         }
-        HandlerUtilsKt.waitForIdle(agent.mHandlerThread, timeoutMs);
-    }
-
-    private void waitForIdle() {
-        waitForIdle(TIMEOUT_MS);
+        agent.waitForIdle(timeoutMs);
     }
 
     @Test
@@ -429,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();
@@ -449,7 +450,7 @@
     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();
@@ -469,7 +470,7 @@
         fail("expected race condition at least once in " + attempts + " attempts");
     }
 
-    private class MockNetworkAgent implements TestableNetworkCallback.HasNetwork {
+    private class TestNetworkAgentWrapper extends NetworkAgentWrapper {
         private static final int VALIDATION_RESULT_BASE = NETWORK_VALIDATION_PROBE_DNS
                 | NETWORK_VALIDATION_PROBE_HTTP
                 | NETWORK_VALIDATION_PROBE_HTTPS;
@@ -480,86 +481,35 @@
                 | 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) throws Exception {
+        TestNetworkAgentWrapper(int transport) throws Exception {
             this(transport, new LinkProperties());
         }
 
-        MockNetworkAgent(int transport, LinkProperties linkProperties) 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");
-            }
-            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(ignoreExceptions(this::onValidationRequested)).start();
                 return null;
@@ -576,56 +526,20 @@
                     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();
 
             mNmCallbacks.onNetworkMonitorCreated(mNetworkMonitor);
 
-            // Waits for the NetworkAgent to be registered, which includes the creation of the
-            // NetworkMonitor.
-            waitForIdle();
+            return na;
         }
 
         private void onValidationRequested() throws Exception {
@@ -645,55 +559,11 @@
             }
         }
 
-        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);
-            }
-        }
-
+        /**
+         * Connect without adding any internet capability.
+         */
         public void connectWithoutInternet() {
-            mNetworkInfo.setDetailedState(DetailedState.CONNECTED, null, null);
-            mNetworkAgent.sendNetworkInfo(mNetworkInfo);
+            super.connect();
         }
 
         /**
@@ -710,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();
                         }
                     }
@@ -763,47 +631,29 @@
             connect(false);
         }
 
-        public void suspend() {
-            mNetworkInfo.setDetailedState(DetailedState.SUSPENDED, null, null);
-            mNetworkAgent.sendNetworkInfo(mNetworkInfo);
+        void setNetworkValid() {
+            mNmValidationResult = VALIDATION_RESULT_VALID;
+            mNmValidationRedirectUrl = null;
         }
 
-        public void resume() {
-            mNetworkInfo.setDetailedState(DetailedState.CONNECTED, null, null);
-            mNetworkAgent.sendNetworkInfo(mNetworkInfo);
+        void setNetworkInvalid() {
+            mNmValidationResult = VALIDATION_RESULT_INVALID;
+            mNmValidationRedirectUrl = null;
         }
 
-        public void disconnect() {
-            mNetworkInfo.setDetailedState(DetailedState.DISCONNECTED, null, null);
-            mNetworkAgent.sendNetworkInfo(mNetworkInfo);
+        void setNetworkPortal(String redirectUrl) {
+            setNetworkInvalid();
+            mNmValidationRedirectUrl = redirectUrl;
         }
 
-        public Network getNetwork() {
-            return new Network(mNetworkAgent.netId);
+        void setNetworkPartial() {
+            mNmValidationResult = VALIDATION_RESULT_PARTIAL;
+            mNmValidationRedirectUrl = null;
         }
 
-        public ConditionVariable getPreventReconnectReceived() {
-            return mPreventReconnectReceived;
-        }
-
-        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() {
@@ -811,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);
         }
     }
 
@@ -995,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());
@@ -1070,187 +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) {
-            HandlerUtilsKt.waitForIdle(mHandlerThread, timeoutMs);
-        }
-
-        public void waitForIdle() {
-            waitForIdle(TIMEOUT_MS);
-        }
-
-        public void setUidRulesChanged(int uidRules) throws RemoteException {
-            mPolicyListener.onUidRulesChanged(Process.myUid(), uidRules);
-        }
-
-        public void setRestrictBackgroundChanged(boolean restrictBackground)
-                throws RemoteException {
-            mPolicyListener.onRestrictBackgroundChanged(restrictBackground);
+            return mConfigMeteredMultipathPreference;
         }
     }
 
@@ -1296,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);
@@ -1313,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
@@ -1322,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);
@@ -1338,6 +1106,9 @@
             mEthernetNetworkAgent = null;
         }
         FakeSettingsProvider.clearSettingsProvider();
+
+        mCsHandlerThread.quitSafely();
+        mAlarmManagerThread.quitSafely();
     }
 
     private void mockDefaultPackages() throws Exception {
@@ -1357,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());
@@ -1449,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.
@@ -1474,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);
@@ -1490,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);
@@ -1505,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);
@@ -1525,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);
@@ -1551,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);
@@ -1559,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);
@@ -1579,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);
@@ -1606,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);
@@ -1676,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);
@@ -1760,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);
@@ -1775,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);
@@ -1799,7 +1551,7 @@
         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);
@@ -1812,7 +1564,7 @@
         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(CallbackRecord.LOSING, mCellNetworkAgent);
@@ -1850,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);
@@ -1891,7 +1643,7 @@
         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;
@@ -1943,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);
@@ -1952,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);
@@ -1968,7 +1720,7 @@
 
         // 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);
@@ -1987,7 +1739,7 @@
 
         // 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.
@@ -2008,12 +1760,12 @@
         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);
@@ -2044,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);
@@ -2065,7 +1817,7 @@
         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(CallbackRecord.LOSING, mWiFiNetworkAgent);
@@ -2100,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);
@@ -2142,12 +1894,12 @@
         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 = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.explicitlySelected(true, false);
         mWiFiNetworkAgent.connect(false);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
@@ -2170,7 +1922,7 @@
         // Disconnect wifi, and then reconnect, again with explicitlySelected=true.
         mWiFiNetworkAgent.disconnect();
         callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.explicitlySelected(true, false);
         mWiFiNetworkAgent.connect(false);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
@@ -2181,7 +1933,7 @@
         callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // Reconnect, again with explicitlySelected=true, but this time validate.
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.explicitlySelected(true, false);
         mWiFiNetworkAgent.connect(true);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
@@ -2191,7 +1943,7 @@
 
         // 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());
@@ -2202,7 +1954,7 @@
         // wifi immediately.
         mWiFiNetworkAgent.disconnect();
         callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.explicitlySelected(true, true);
         mWiFiNetworkAgent.connect(false);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
@@ -2215,7 +1967,7 @@
         // Check that the network is not scored specially and that the device prefers cell data.
         mWiFiNetworkAgent.disconnect();
         callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.explicitlySelected(false, true);
         mWiFiNetworkAgent.connect(false);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
@@ -2277,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.
@@ -2371,18 +2123,17 @@
     @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);
@@ -2394,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);
@@ -2423,16 +2173,16 @@
         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);
     }
 
@@ -2446,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);
@@ -2484,7 +2234,7 @@
         // Disconnect and reconnect wifi with partial connectivity again.
         mWiFiNetworkAgent.disconnect();
         callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connectWithPartialConnectivity();
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_PARTIAL_CONNECTIVITY, mWiFiNetworkAgent);
@@ -2500,7 +2250,7 @@
         // 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 /* explicitlySelected */,
                 true /* acceptUnvalidated */);
@@ -2524,7 +2274,7 @@
         // 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 MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.explicitlySelected(true /* explicitlySelected */,
                 true /* acceptUnvalidated */);
         mWiFiNetworkAgent.connectWithPartialConnectivity();
@@ -2548,7 +2298,7 @@
 
         // 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 MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.explicitlySelected(false /* explicitlySelected */,
                 true /* acceptUnvalidated */);
 
@@ -2580,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);
@@ -2629,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);
@@ -2642,7 +2392,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 secondRedirectUrl = "http://example.com/secondPath";
         mWiFiNetworkAgent.connectWithCaptivePortal(secondRedirectUrl);
         captivePortalCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
@@ -2680,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();
@@ -2740,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);
     }
@@ -2846,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);
@@ -2973,7 +2721,7 @@
             public void writeToParcel(Parcel dest, int flags) {}
         }
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
 
         UidAwareNetworkSpecifier networkSpecifier = new UidAwareNetworkSpecifier();
@@ -3026,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);
@@ -3046,7 +2794,7 @@
         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();
@@ -3065,7 +2813,8 @@
         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);
@@ -3090,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
@@ -3160,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);
@@ -3179,13 +2928,13 @@
         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.
@@ -3278,7 +3027,7 @@
             }
         });
 
-        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.
@@ -3295,7 +3044,7 @@
                 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);
 
@@ -3338,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);
@@ -3355,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);
@@ -3384,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();
@@ -3444,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);
@@ -3476,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);
 
@@ -3498,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);
@@ -3512,7 +3259,7 @@
 
         // 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);
@@ -3525,11 +3272,11 @@
         // 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);
 
@@ -3547,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;
@@ -3575,7 +3321,7 @@
         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);
@@ -3595,7 +3341,7 @@
         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);
@@ -3623,7 +3369,7 @@
         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();
     }
@@ -3646,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();
     }
@@ -3876,7 +3622,7 @@
             assertNull(mCm.getNetworkCapabilities(mWiFiNetworkAgent.getNetwork()));
         }
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         ConditionVariable cv = waitForConnectivityBroadcasts(1);
         mWiFiNetworkAgent.connect(true);
         waitFor(cv);
@@ -3940,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();
 
@@ -3961,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.
@@ -3972,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);
@@ -4092,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();
 
@@ -4137,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.
@@ -4150,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);
@@ -4187,7 +3933,7 @@
         // assertFalse(isUdpPortInUse(srcPort2));
 
         mWiFiNetworkAgent.disconnect();
-        waitFor(mWiFiNetworkAgent.getDisconnectedCV());
+        mWiFiNetworkAgent.expectDisconnected();
         mWiFiNetworkAgent = null;
     }
 
@@ -4263,7 +4009,7 @@
         testSocketV6.close();
 
         mWiFiNetworkAgent.disconnect();
-        waitFor(mWiFiNetworkAgent.getDisconnectedCV());
+        mWiFiNetworkAgent.expectDisconnected();
         mWiFiNetworkAgent = null;
     }
 
@@ -4279,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);
 
@@ -4316,7 +4062,7 @@
         // assertFalse(isUdpPortInUse(srcPort));
 
         mWiFiNetworkAgent.disconnect();
-        waitFor(mWiFiNetworkAgent.getDisconnectedCV());
+        mWiFiNetworkAgent.expectDisconnected();
         mWiFiNetworkAgent = null;
     }
 
@@ -4379,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.
@@ -4394,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();
@@ -4416,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();
@@ -4526,7 +4272,7 @@
         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();
@@ -4599,7 +4345,7 @@
 
         // 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(CallbackRecord.AVAILABLE, networkAgent);
         networkCallback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED, networkAgent);
@@ -4631,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()};
@@ -4708,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);
@@ -4791,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());
@@ -4880,7 +4626,7 @@
                 .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);
@@ -5015,7 +4761,7 @@
         mCm.registerDefaultNetworkCallback(defaultCallback);
         defaultCallback.assertNoCallback();
 
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(false);
 
         genericNetworkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
@@ -5025,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);
@@ -5110,13 +4858,14 @@
         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);
@@ -5140,13 +4889,14 @@
         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);
@@ -5170,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);
@@ -5199,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();
@@ -5229,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);
@@ -5246,7 +4999,7 @@
         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(
@@ -5257,7 +5010,7 @@
                 && caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
                 && !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);
 
@@ -5327,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);
@@ -5345,7 +5099,7 @@
         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.expectCapabilitiesThat(vpnNetworkAgent,
@@ -5354,7 +5108,7 @@
                 && !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);
 
@@ -5382,7 +5136,7 @@
     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();
@@ -5394,7 +5148,7 @@
     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();
@@ -5406,14 +5160,15 @@
     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));
@@ -5429,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();
@@ -5460,20 +5215,21 @@
    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));
@@ -5531,14 +5287,15 @@
     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));
@@ -5582,21 +5339,21 @@
                 .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.
@@ -5607,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();
 
@@ -5631,18 +5388,18 @@
         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);
@@ -5658,8 +5415,8 @@
         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();
@@ -5709,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);
@@ -5756,7 +5513,7 @@
         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);
@@ -5863,7 +5620,7 @@
                 .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);
@@ -5873,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);
@@ -5898,7 +5655,7 @@
                 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);
@@ -5944,7 +5701,7 @@
     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);
@@ -5960,7 +5717,7 @@
     @Test
     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));
@@ -5969,7 +5726,7 @@
     @Test
     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));
@@ -5988,14 +5745,15 @@
         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);
@@ -6044,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
@@ -6070,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());
@@ -6085,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());
@@ -6098,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
@@ -6147,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);
@@ -6169,9 +5929,10 @@
     }
 
 
-    private MockNetworkAgent establishVpn(LinkProperties lp, int establishingUid,
+    private TestNetworkAgentWrapper establishVpn(LinkProperties lp, int establishingUid,
             Set<UidRange> vpnRange) throws Exception {
-        final MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN, lp);
+        final TestNetworkAgentWrapper
+                vpnNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN, lp);
         vpnNetworkAgent.getNetworkCapabilities().setEstablishingVpnAppUid(establishingUid);
         mMockVpn.setNetworkAgent(vpnNetworkAgent);
         mMockVpn.connect();