diff --git a/core/java/android/net/ConnectivityManager.java b/core/java/android/net/ConnectivityManager.java
index c96d19d..9e8acd0 100644
--- a/core/java/android/net/ConnectivityManager.java
+++ b/core/java/android/net/ConnectivityManager.java
@@ -1466,9 +1466,7 @@
         // Map from type to transports.
         final int NOT_FOUND = -1;
         final int transport = sLegacyTypeToTransport.get(type, NOT_FOUND);
-        if (transport == NOT_FOUND) {
-            throw new IllegalArgumentException("unknown legacy type: " + type);
-        }
+        Preconditions.checkArgument(transport != NOT_FOUND, "unknown legacy type: " + type);
         nc.addTransportType(transport);
 
         // Map from type to capabilities.
@@ -1814,9 +1812,7 @@
      */
     public void removeDefaultNetworkActiveListener(OnNetworkActiveListener l) {
         INetworkActivityListener rl = mNetworkActivityListeners.get(l);
-        if (rl == null) {
-            throw new IllegalArgumentException("Listener not registered: " + l);
-        }
+        Preconditions.checkArgument(rl != null, "Listener was not registered.");
         try {
             getNetworkManagementService().unregisterNetworkActivityListener(rl);
         } catch (RemoteException e) {
@@ -1873,9 +1869,8 @@
 
     /** {@hide} */
     public static final void enforceTetherChangePermission(Context context, String callingPkg) {
-        if (null == context || null == callingPkg) {
-            throw new IllegalArgumentException("arguments should not be null");
-        }
+        Preconditions.checkNotNull(context, "Context cannot be null");
+        Preconditions.checkNotNull(callingPkg, "callingPkg cannot be null");
 
         if (context.getResources().getStringArray(
                 com.android.internal.R.array.config_mobile_hotspot_provision_app).length == 2) {
@@ -2773,84 +2768,67 @@
         }
 
         CallbackHandler(Handler handler) {
-            this(handler.getLooper());
+            this(Preconditions.checkNotNull(handler, "Handler cannot be null.").getLooper());
         }
 
         @Override
         public void handleMessage(Message message) {
-            NetworkRequest request = getObject(message, NetworkRequest.class);
-            Network network = getObject(message, Network.class);
+            if (message.what == EXPIRE_LEGACY_REQUEST) {
+                expireRequest((NetworkCapabilities) message.obj, message.arg1);
+                return;
+            }
+
+            final NetworkRequest request = getObject(message, NetworkRequest.class);
+            final Network network = getObject(message, Network.class);
+            final NetworkCallback callback;
+            synchronized (sCallbacks) {
+                callback = sCallbacks.get(request);
+            }
             if (DBG) {
                 Log.d(TAG, getCallbackName(message.what) + " for network " + network);
             }
+            if (callback == null) {
+                Log.w(TAG, "callback not found for " + getCallbackName(message.what) + " message");
+                return;
+            }
+
             switch (message.what) {
                 case CALLBACK_PRECHECK: {
-                    NetworkCallback callback = getCallback(message);
-                    if (callback != null) {
-                        callback.onPreCheck(network);
-                    }
+                    callback.onPreCheck(network);
                     break;
                 }
                 case CALLBACK_AVAILABLE: {
-                    NetworkCallback callback = getCallback(message);
-                    if (callback != null) {
-                        callback.onAvailable(network);
-                    }
+                    callback.onAvailable(network);
                     break;
                 }
                 case CALLBACK_LOSING: {
-                    NetworkCallback callback = getCallback(message);
-                    if (callback != null) {
-                        callback.onLosing(network, message.arg1);
-                    }
+                    callback.onLosing(network, message.arg1);
                     break;
                 }
                 case CALLBACK_LOST: {
-                    NetworkCallback callback = getCallback(message);
-                    if (callback != null) {
-                        callback.onLost(network);
-                    }
+                    callback.onLost(network);
                     break;
                 }
                 case CALLBACK_UNAVAIL: {
-                    NetworkCallback callback = getCallback(message);
-                    if (callback != null) {
-                        callback.onUnavailable();
-                    }
+                    callback.onUnavailable();
                     break;
                 }
                 case CALLBACK_CAP_CHANGED: {
-                    NetworkCallback callback = getCallback(message);
-                    if (callback != null) {
-                        NetworkCapabilities cap = getObject(message, NetworkCapabilities.class);
-                        callback.onCapabilitiesChanged(network, cap);
-                    }
+                    NetworkCapabilities cap = getObject(message, NetworkCapabilities.class);
+                    callback.onCapabilitiesChanged(network, cap);
                     break;
                 }
                 case CALLBACK_IP_CHANGED: {
-                    NetworkCallback callback = getCallback(message);
-                    if (callback != null) {
-                        LinkProperties lp = getObject(message, LinkProperties.class);
-                        callback.onLinkPropertiesChanged(network, lp);
-                    }
+                    LinkProperties lp = getObject(message, LinkProperties.class);
+                    callback.onLinkPropertiesChanged(network, lp);
                     break;
                 }
                 case CALLBACK_SUSPENDED: {
-                    NetworkCallback callback = getCallback(message);
-                    if (callback != null) {
-                        callback.onNetworkSuspended(network);
-                    }
+                    callback.onNetworkSuspended(network);
                     break;
                 }
                 case CALLBACK_RESUMED: {
-                    NetworkCallback callback = getCallback(message);
-                    if (callback != null) {
-                        callback.onNetworkResumed(network);
-                    }
-                    break;
-                }
-                case EXPIRE_LEGACY_REQUEST: {
-                    expireRequest((NetworkCapabilities)message.obj, message.arg1);
+                    callback.onNetworkResumed(network);
                     break;
                 }
             }
@@ -2859,18 +2837,6 @@
         private <T> T getObject(Message msg, Class<T> c) {
             return (T) msg.getData().getParcelable(c.getSimpleName());
         }
-
-        private NetworkCallback getCallback(Message msg) {
-            final NetworkRequest req = getObject(msg, NetworkRequest.class);
-            final NetworkCallback callback;
-            synchronized(sCallbacks) {
-                callback = sCallbacks.get(req);
-            }
-            if (callback == null) {
-                Log.w(TAG, "callback not found for " + getCallbackName(msg.what) + " message");
-            }
-            return callback;
-        }
     }
 
     private CallbackHandler getDefaultHandler() {
@@ -2890,7 +2856,7 @@
 
     private NetworkRequest sendRequestForNetwork(NetworkCapabilities need, NetworkCallback callback,
             int timeoutMs, int action, int legacyType, CallbackHandler handler) {
-        Preconditions.checkArgument(callback != null, "null NetworkCallback");
+        checkCallbackNotNull(callback);
         Preconditions.checkArgument(action == REQUEST || need != null, "null NetworkCapabilities");
         final NetworkRequest request;
         try {
@@ -3042,14 +3008,11 @@
      */
     public void requestNetwork(NetworkRequest request, NetworkCallback networkCallback,
             int timeoutMs) {
-        if (timeoutMs <= 0) {
-            throw new IllegalArgumentException("Non-positive timeoutMs: " + timeoutMs);
-        }
+        checkTimeout(timeoutMs);
         int legacyType = inferLegacyTypeForNetworkCapabilities(request.networkCapabilities);
         requestNetwork(request, networkCallback, timeoutMs, legacyType, getDefaultHandler());
     }
 
-
     /**
      * Request a network to satisfy a set of {@link android.net.NetworkCapabilities}, limited
      * by a timeout.
@@ -3079,9 +3042,7 @@
      */
     public void requestNetwork(NetworkRequest request, NetworkCallback networkCallback,
             Handler handler, int timeoutMs) {
-        if (timeoutMs <= 0) {
-            throw new IllegalArgumentException("Non-positive timeoutMs");
-        }
+        checkTimeout(timeoutMs);
         int legacyType = inferLegacyTypeForNetworkCapabilities(request.networkCapabilities);
         CallbackHandler cbHandler = new CallbackHandler(handler);
         requestNetwork(request, networkCallback, timeoutMs, legacyType, cbHandler);
@@ -3153,7 +3114,7 @@
      *         {@link NetworkCapabilities#NET_CAPABILITY_CAPTIVE_PORTAL}.
      */
     public void requestNetwork(NetworkRequest request, PendingIntent operation) {
-        checkPendingIntent(operation);
+        checkPendingIntentNotNull(operation);
         try {
             mService.pendingRequestForNetwork(request.networkCapabilities, operation);
         } catch (RemoteException e) {
@@ -3176,7 +3137,7 @@
      *                  corresponding NetworkRequest you'd like to remove. Cannot be null.
      */
     public void releaseNetworkRequest(PendingIntent operation) {
-        checkPendingIntent(operation);
+        checkPendingIntentNotNull(operation);
         try {
             mService.releasePendingNetworkRequest(operation);
         } catch (RemoteException e) {
@@ -3184,10 +3145,16 @@
         }
     }
 
-    private void checkPendingIntent(PendingIntent intent) {
-        if (intent == null) {
-            throw new IllegalArgumentException("PendingIntent cannot be null.");
-        }
+    private static void checkPendingIntentNotNull(PendingIntent intent) {
+        Preconditions.checkNotNull(intent, "PendingIntent cannot be null.");
+    }
+
+    private static void checkCallbackNotNull(NetworkCallback callback) {
+        Preconditions.checkNotNull(callback, "null NetworkCallback");
+    }
+
+    private static void checkTimeout(int timeoutMs) {
+        Preconditions.checkArgumentPositive(timeoutMs, "timeoutMs must be strictly positive.");
     }
 
     /**
@@ -3257,7 +3224,7 @@
      *                  comes from {@link PendingIntent#getBroadcast}. Cannot be null.
      */
     public void registerNetworkCallback(NetworkRequest request, PendingIntent operation) {
-        checkPendingIntent(operation);
+        checkPendingIntentNotNull(operation);
         try {
             mService.pendingListenForNetwork(request.networkCapabilities, operation);
         } catch (RemoteException e) {
@@ -3301,8 +3268,9 @@
         // capabilities, this request is guaranteed, at all times, to be
         // satisfied by the same network, if any, that satisfies the default
         // request, i.e., the system default network.
+        NetworkCapabilities nullCapabilities = null;
         CallbackHandler cbHandler = new CallbackHandler(handler);
-        sendRequestForNetwork(null, networkCallback, 0, REQUEST, TYPE_NONE, cbHandler);
+        sendRequestForNetwork(nullCapabilities, networkCallback, 0, REQUEST, TYPE_NONE, cbHandler);
     }
 
     /**
@@ -3339,7 +3307,7 @@
      * @param networkCallback The {@link NetworkCallback} used when making the request.
      */
     public void unregisterNetworkCallback(NetworkCallback networkCallback) {
-        Preconditions.checkArgument(networkCallback != null, "null NetworkCallback");
+        checkCallbackNotNull(networkCallback);
         final List<NetworkRequest> reqs = new ArrayList<>();
         // Find all requests associated to this callback and stop callback triggers immediately.
         // Callback is reusable immediately. http://b/20701525, http://b/35921499.
@@ -3375,6 +3343,7 @@
      *                  Cannot be null.
      */
     public void unregisterNetworkCallback(PendingIntent operation) {
+        checkPendingIntentNotNull(operation);
         releaseNetworkRequest(operation);
     }
 
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index bbbf626..00b0f98 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -115,7 +115,7 @@
  * Tests for {@link ConnectivityService}.
  *
  * Build, install and run with:
- *  runtest frameworks-services -c com.android.server.ConnectivityServiceTest
+ *  runtest frameworks-net -c com.android.server.ConnectivityServiceTest
  */
 public class ConnectivityServiceTest extends AndroidTestCase {
     private static final String TAG = "ConnectivityServiceTest";
@@ -223,13 +223,32 @@
         }
     }
 
+    public void waitForIdle(int timeoutMs) {
+        waitForIdleHandler(mService.mHandlerThread, timeoutMs);
+        waitForIdle(mCellNetworkAgent, timeoutMs);
+        waitForIdle(mWiFiNetworkAgent, timeoutMs);
+        waitForIdle(mEthernetNetworkAgent, timeoutMs);
+        waitForIdleHandler(mService.mHandlerThread, timeoutMs);
+    }
+
+    public void waitForIdle(MockNetworkAgent agent, int timeoutMs) {
+        if (agent == null) {
+            return;
+        }
+        waitForIdleHandler(agent.mHandlerThread, timeoutMs);
+    }
+
+    private void waitForIdle() {
+        waitForIdle(TIMEOUT_MS);
+    }
+
     @SmallTest
     public void testWaitForIdle() {
         final int attempts = 50;  // Causes the test to take about 200ms on bullhead-eng.
 
         // Tests that waitForIdle returns immediately if the service is already idle.
         for (int i = 0; i < attempts; i++) {
-            mService.waitForIdle();
+            waitForIdle();
         }
 
         // Bring up a network that we can use to send messages to ConnectivityService.
@@ -243,7 +262,7 @@
         // Tests that calling waitForIdle waits for messages to be processed.
         for (int i = 0; i < attempts; i++) {
             mWiFiNetworkAgent.setSignalStrength(i);
-            mService.waitForIdle();
+            waitForIdle();
             assertEquals(i, mCm.getNetworkCapabilities(n).getSignalStrength());
         }
     }
@@ -344,18 +363,10 @@
             };
             // Waits for the NetworkAgent to be registered, which includes the creation of the
             // NetworkMonitor.
-            mService.waitForIdle();
+            waitForIdle();
             mWrappedNetworkMonitor = mService.getLastCreatedWrappedNetworkMonitor();
         }
 
-        public void waitForIdle(int timeoutMs) {
-            waitForIdleHandler(mHandlerThread, timeoutMs);
-        }
-
-        public void waitForIdle() {
-            waitForIdle(TIMEOUT_MS);
-        }
-
         public void adjustScore(int change) {
             mScore += change;
             mNetworkAgent.sendNetworkScore(mScore);
@@ -802,13 +813,22 @@
 
     public void tearDown() throws Exception {
         setMobileDataAlwaysOn(false);
-        if (mCellNetworkAgent != null) { mCellNetworkAgent.disconnect(); }
-        if (mWiFiNetworkAgent != null) { mWiFiNetworkAgent.disconnect(); }
-        mCellNetworkAgent = mWiFiNetworkAgent = null;
+        if (mCellNetworkAgent != null) {
+            mCellNetworkAgent.disconnect();
+            mCellNetworkAgent = null;
+        }
+        if (mWiFiNetworkAgent != null) {
+            mWiFiNetworkAgent.disconnect();
+            mWiFiNetworkAgent = null;
+        }
+        if (mEthernetNetworkAgent != null) {
+            mEthernetNetworkAgent.disconnect();
+            mEthernetNetworkAgent = null;
+        }
         super.tearDown();
     }
 
-    private int transportToLegacyType(int transport) {
+    private static int transportToLegacyType(int transport) {
         switch (transport) {
             case TRANSPORT_ETHERNET:
                 return TYPE_ETHERNET;
@@ -840,7 +860,8 @@
         }
         // Test getNetworkInfo(Network)
         assertNotNull(mCm.getNetworkInfo(mCm.getActiveNetwork()));
-        assertEquals(transportToLegacyType(transport), mCm.getNetworkInfo(mCm.getActiveNetwork()).getType());
+        assertEquals(transportToLegacyType(transport),
+                mCm.getNetworkInfo(mCm.getActiveNetwork()).getType());
         // Test getNetworkCapabilities(Network)
         assertNotNull(mCm.getNetworkCapabilities(mCm.getActiveNetwork()));
         assertTrue(mCm.getNetworkCapabilities(mCm.getActiveNetwork()).hasTransport(transport));
@@ -911,7 +932,7 @@
                 mCm.getAllNetworks()[1].equals(mCellNetworkAgent.getNetwork()));
         // Test cellular linger timeout.
         waitFor(mCellNetworkAgent.getDisconnectedCV());
-        mService.waitForIdle();
+        waitForIdle();
         assertEquals(1, mCm.getAllNetworks().length);
         verifyActiveNetwork(TRANSPORT_WIFI);
         assertEquals(1, mCm.getAllNetworks().length);
@@ -934,11 +955,11 @@
         // Test bringing up unvalidated cellular
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(false);
-        mService.waitForIdle();
+        waitForIdle();
         verifyActiveNetwork(TRANSPORT_WIFI);
         // Test cellular disconnect.
         mCellNetworkAgent.disconnect();
-        mService.waitForIdle();
+        waitForIdle();
         verifyActiveNetwork(TRANSPORT_WIFI);
         // Test bringing up validated cellular
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
@@ -1290,7 +1311,7 @@
         }
 
         void assertNoCallback() {
-            mService.waitForIdle();
+            waitForIdle();
             CallbackInfo c = mCallbacks.peek();
             assertNull("Unexpected callback: " + c, c);
         }
@@ -1331,7 +1352,7 @@
 
         // This should not trigger spurious onAvailable() callbacks, b/21762680.
         mCellNetworkAgent.adjustScore(-1);
-        mService.waitForIdle();
+        waitForIdle();
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
 
@@ -1369,7 +1390,7 @@
 
         // This should not trigger spurious onAvailable() callbacks, b/21762680.
         mCellNetworkAgent.adjustScore(-1);
-        mService.waitForIdle();
+        waitForIdle();
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
 
@@ -1481,7 +1502,7 @@
         defaultCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
 
         mCm.unregisterNetworkCallback(callback);
-        mService.waitForIdle();
+        waitForIdle();
 
         // Check that a network is only lingered or torn down if it would not satisfy a request even
         // if it validated.
@@ -1842,26 +1863,30 @@
         ConditionVariable cv = mCellNetworkAgent.getDisconnectedCV();
         mCellNetworkAgent.connectWithoutInternet();
         waitFor(cv);
-        mService.waitForIdle();
+        waitForIdle();
         assertEquals(0, mCm.getAllNetworks().length);
         verifyNoNetwork();
+
         // Test bringing up validated WiFi.
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         cv = waitForConnectivityBroadcasts(1);
         mWiFiNetworkAgent.connect(true);
         waitFor(cv);
         verifyActiveNetwork(TRANSPORT_WIFI);
+
         // Register MMS NetworkRequest
         NetworkRequest.Builder builder = new NetworkRequest.Builder();
         builder.addCapability(NetworkCapabilities.NET_CAPABILITY_MMS);
         final TestNetworkCallback networkCallback = new TestNetworkCallback();
         mCm.requestNetwork(builder.build(), networkCallback);
+
         // Test bringing up unvalidated cellular with MMS
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
         mCellNetworkAgent.addCapability(NET_CAPABILITY_MMS);
         mCellNetworkAgent.connectWithoutInternet();
         networkCallback.expectAvailableCallbacks(mCellNetworkAgent);
         verifyActiveNetwork(TRANSPORT_WIFI);
+
         // Test releasing NetworkRequest disconnects cellular with MMS
         cv = mCellNetworkAgent.getDisconnectedCV();
         mCm.unregisterNetworkCallback(networkCallback);
@@ -1877,17 +1902,20 @@
         mCellNetworkAgent.connect(false);
         waitFor(cv);
         verifyActiveNetwork(TRANSPORT_CELLULAR);
+
         // Register MMS NetworkRequest
         NetworkRequest.Builder builder = new NetworkRequest.Builder();
         builder.addCapability(NetworkCapabilities.NET_CAPABILITY_MMS);
         final TestNetworkCallback networkCallback = new TestNetworkCallback();
         mCm.requestNetwork(builder.build(), networkCallback);
+
         // Test bringing up MMS cellular network
         MockNetworkAgent mmsNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
         mmsNetworkAgent.addCapability(NET_CAPABILITY_MMS);
         mmsNetworkAgent.connectWithoutInternet();
         networkCallback.expectAvailableCallbacks(mmsNetworkAgent);
         verifyActiveNetwork(TRANSPORT_CELLULAR);
+
         // Test releasing MMS NetworkRequest does not disconnect main cellular NetworkAgent
         cv = mmsNetworkAgent.getDisconnectedCV();
         mCm.unregisterNetworkCallback(networkCallback);
@@ -2294,7 +2322,7 @@
         ContentResolver cr = mServiceContext.getContentResolver();
         Settings.Global.putInt(cr, Settings.Global.MOBILE_DATA_ALWAYS_ON, enable ? 1 : 0);
         mService.updateMobileDataAlwaysOn();
-        mService.waitForIdle();
+        waitForIdle();
     }
 
     private boolean isForegroundNetwork(MockNetworkAgent network) {
@@ -2336,7 +2364,7 @@
         assertTrue(isForegroundNetwork(mWiFiNetworkAgent));
 
         // When lingering is complete, cell is still there but is now in the background.
-        mService.waitForIdle();
+        waitForIdle();
         int timeoutMs = TEST_LINGER_DELAY_MS + TEST_LINGER_DELAY_MS / 4;
         fgCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent, timeoutMs);
         // Expect a network capabilities update sans FOREGROUND.
@@ -2392,6 +2420,11 @@
         //    and NUM_REQUESTS onAvailable callbacks to fire.
         // See how long it took.
         final int NUM_REQUESTS = 90;
+        final int REGISTER_TIME_LIMIT_MS = 180;
+        final int CONNECT_TIME_LIMIT_MS = 50;
+        final int SWITCH_TIME_LIMIT_MS = 50;
+        final int UNREGISTER_TIME_LIMIT_MS = 20;
+
         final NetworkRequest request = new NetworkRequest.Builder().clearCapabilities().build();
         final NetworkCallback[] callbacks = new NetworkCallback[NUM_REQUESTS];
         final CountDownLatch availableLatch = new CountDownLatch(NUM_REQUESTS);
@@ -2404,14 +2437,12 @@
             };
         }
 
-        final int REGISTER_TIME_LIMIT_MS = 180;
         assertTimeLimit("Registering callbacks", REGISTER_TIME_LIMIT_MS, () -> {
             for (NetworkCallback cb : callbacks) {
                 mCm.registerNetworkCallback(request, cb);
             }
         });
 
-        final int CONNECT_TIME_LIMIT_MS = 40;
         mCellNetworkAgent = new MockNetworkAgent(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,
@@ -2419,31 +2450,29 @@
         mCellNetworkAgent.connect(false);
 
         long onAvailableDispatchingDuration = durationOf(() -> {
-            if (!awaitLatch(availableLatch, CONNECT_TIME_LIMIT_MS)) {
-                fail(String.format("Only dispatched %d/%d onAvailable callbacks in %dms",
-                        NUM_REQUESTS - availableLatch.getCount(), NUM_REQUESTS,
-                        CONNECT_TIME_LIMIT_MS));
-            }
+            awaitLatch(availableLatch, 10 * CONNECT_TIME_LIMIT_MS);
         });
-        Log.d(TAG, String.format("Connect, %d callbacks: %dms, acceptable %dms",
-                NUM_REQUESTS, onAvailableDispatchingDuration, CONNECT_TIME_LIMIT_MS));
+        Log.d(TAG, String.format("Dispatched %d of %d onAvailable callbacks in %dms",
+                NUM_REQUESTS - availableLatch.getCount(), NUM_REQUESTS,
+                onAvailableDispatchingDuration));
+        assertTrue(String.format("Dispatching %d onAvailable callbacks in %dms, expected %dms",
+                NUM_REQUESTS, onAvailableDispatchingDuration, CONNECT_TIME_LIMIT_MS),
+                onAvailableDispatchingDuration <= CONNECT_TIME_LIMIT_MS);
 
-        final int SWITCH_TIME_LIMIT_MS = 40;
-        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         // Give wifi a high enough score that we'll linger cell when wifi comes up.
+        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.adjustScore(40);
         mWiFiNetworkAgent.connect(false);
 
         long onLostDispatchingDuration = durationOf(() -> {
-            if (!awaitLatch(losingLatch, SWITCH_TIME_LIMIT_MS)) {
-                fail(String.format("Only dispatched %d/%d onLosing callbacks in %dms",
-                        NUM_REQUESTS - losingLatch.getCount(), NUM_REQUESTS, SWITCH_TIME_LIMIT_MS));
-            }
+            awaitLatch(losingLatch, 10 * SWITCH_TIME_LIMIT_MS);
         });
-        Log.d(TAG, String.format("Linger, %d callbacks: %dms, acceptable %dms",
-                NUM_REQUESTS, onLostDispatchingDuration, SWITCH_TIME_LIMIT_MS));
+        Log.d(TAG, String.format("Dispatched %d of %d onLosing callbacks in %dms",
+                NUM_REQUESTS - losingLatch.getCount(), NUM_REQUESTS, onLostDispatchingDuration));
+        assertTrue(String.format("Dispatching %d onLosing callbacks in %dms, expected %dms",
+                NUM_REQUESTS, onLostDispatchingDuration, SWITCH_TIME_LIMIT_MS),
+                onLostDispatchingDuration <= SWITCH_TIME_LIMIT_MS);
 
-        final int UNREGISTER_TIME_LIMIT_MS = 10;
         assertTimeLimit("Unregistering callbacks", UNREGISTER_TIME_LIMIT_MS, () -> {
             for (NetworkCallback cb : callbacks) {
                 mCm.unregisterNetworkCallback(cb);
@@ -2466,9 +2495,7 @@
 
     private boolean awaitLatch(CountDownLatch l, long timeoutMs) {
         try {
-            if (l.await(timeoutMs, TimeUnit.MILLISECONDS)) {
-                return true;
-            }
+            return l.await(timeoutMs, TimeUnit.MILLISECONDS);
         } catch (InterruptedException e) {}
         return false;
     }
@@ -2520,7 +2547,7 @@
         assertFalse(testFactory.getMyStartRequested());  // Because the cell network outscores us.
 
         // Check that cell data stays up.
-        mService.waitForIdle();
+        waitForIdle();
         verifyActiveNetwork(TRANSPORT_WIFI);
         assertEquals(2, mCm.getAllNetworks().length);
 
@@ -2549,7 +2576,7 @@
         for (int i = 0; i < values.length; i++) {
             Settings.Global.putInt(cr, settingName, 1);
             tracker.reevaluate();
-            mService.waitForIdle();
+            waitForIdle();
             String msg = String.format("config=false, setting=%s", values[i]);
             assertTrue(mService.avoidBadWifi());
             assertFalse(msg, tracker.shouldNotifyWifiUnvalidated());
@@ -2559,19 +2586,19 @@
 
         Settings.Global.putInt(cr, settingName, 0);
         tracker.reevaluate();
-        mService.waitForIdle();
+        waitForIdle();
         assertFalse(mService.avoidBadWifi());
         assertFalse(tracker.shouldNotifyWifiUnvalidated());
 
         Settings.Global.putInt(cr, settingName, 1);
         tracker.reevaluate();
-        mService.waitForIdle();
+        waitForIdle();
         assertTrue(mService.avoidBadWifi());
         assertFalse(tracker.shouldNotifyWifiUnvalidated());
 
         Settings.Global.putString(cr, settingName, null);
         tracker.reevaluate();
-        mService.waitForIdle();
+        waitForIdle();
         assertFalse(mService.avoidBadWifi());
         assertTrue(tracker.shouldNotifyWifiUnvalidated());
     }
@@ -2714,7 +2741,7 @@
                 tracker.configMeteredMultipathPreference = config;
                 Settings.Global.putString(cr, settingName, setting);
                 tracker.reevaluate();
-                mService.waitForIdle();
+                waitForIdle();
 
                 final int expected = (setting != null) ? Integer.parseInt(setting) : config;
                 String msg = String.format("config=%d, setting=%s", config, setting);
@@ -2902,7 +2929,7 @@
         waitFor(cv);
         verifyActiveNetwork(TRANSPORT_WIFI);
         mWiFiNetworkAgent.sendLinkProperties(lp);
-        mService.waitForIdle();
+        waitForIdle();
         return mWiFiNetworkAgent.getNetwork();
     }
 
@@ -2981,7 +3008,7 @@
         callback.expectError(PacketKeepalive.ERROR_INVALID_NETWORK);
 
         // ... and that stopping it after that has no adverse effects.
-        mService.waitForIdle();
+        waitForIdle();
         final Network myNetAlias = myNet;
         assertNull(mCm.getNetworkCapabilities(myNetAlias));
         ka.stop();
@@ -2996,7 +3023,7 @@
         ka.stop();
         mWiFiNetworkAgent.disconnect();
         waitFor(mWiFiNetworkAgent.getDisconnectedCV());
-        mService.waitForIdle();
+        waitForIdle();
         callback.expectStopped();
 
         // Reconnect.
@@ -3197,7 +3224,7 @@
             mCm.unregisterNetworkCallback(pendingIntent);
         }
         pendingIntents.clear();
-        mService.waitForIdle(5000);
+        waitForIdle(5000);
 
         // Test that the limit is not hit when MAX_REQUESTS requests are added and removed.
         for (int i = 0; i < MAX_REQUESTS; i++) {
@@ -3205,20 +3232,20 @@
             mCm.requestNetwork(networkRequest, networkCallback);
             mCm.unregisterNetworkCallback(networkCallback);
         }
-        mService.waitForIdle();
+        waitForIdle();
         for (int i = 0; i < MAX_REQUESTS; i++) {
             NetworkCallback networkCallback = new NetworkCallback();
             mCm.registerNetworkCallback(networkRequest, networkCallback);
             mCm.unregisterNetworkCallback(networkCallback);
         }
-        mService.waitForIdle();
+        waitForIdle();
         for (int i = 0; i < MAX_REQUESTS; i++) {
             PendingIntent pendingIntent =
                     PendingIntent.getBroadcast(mContext, 0, new Intent("b" + i), 0);
             mCm.requestNetwork(networkRequest, pendingIntent);
             mCm.unregisterNetworkCallback(pendingIntent);
         }
-        mService.waitForIdle();
+        waitForIdle();
         for (int i = 0; i < MAX_REQUESTS; i++) {
             PendingIntent pendingIntent =
                     PendingIntent.getBroadcast(mContext, 0, new Intent("c" + i), 0);
diff --git a/tests/net/java/com/android/server/net/NetworkStatsAccessTest.java b/tests/net/java/com/android/server/net/NetworkStatsAccessTest.java
index e6b1c6d..23318c2 100644
--- a/tests/net/java/com/android/server/net/NetworkStatsAccessTest.java
+++ b/tests/net/java/com/android/server/net/NetworkStatsAccessTest.java
@@ -16,6 +16,7 @@
 
 package com.android.server.net;
 
+import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.when;
 
 import android.Manifest;
@@ -26,17 +27,21 @@
 import android.content.Context;
 import android.content.pm.PackageManager;
 import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
 import android.telephony.TelephonyManager;
 
 import com.android.server.LocalServices;
 
-import junit.framework.TestCase;
-
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+@RunWith(AndroidJUnit4.class)
 @SmallTest
-public class NetworkStatsAccessTest extends TestCase {
+public class NetworkStatsAccessTest {
     private static final String TEST_PKG = "com.example.test";
     private static final int TEST_UID = 12345;
 
@@ -48,9 +53,8 @@
     // Hold the real service so we can restore it when tearing down the test.
     private DevicePolicyManagerInternal mSystemDpmi;
 
-    @Override
+    @Before
     public void setUp() throws Exception {
-        super.setUp();
         MockitoAnnotations.initMocks(this);
 
         mSystemDpmi = LocalServices.getService(DevicePolicyManagerInternal.class);
@@ -61,13 +65,13 @@
         when(mContext.getSystemService(Context.APP_OPS_SERVICE)).thenReturn(mAppOps);
     }
 
-    @Override
+    @After
     public void tearDown() throws Exception {
         LocalServices.removeServiceForTest(DevicePolicyManagerInternal.class);
         LocalServices.addService(DevicePolicyManagerInternal.class, mSystemDpmi);
-        super.tearDown();
     }
 
+    @Test
     public void testCheckAccessLevel_hasCarrierPrivileges() throws Exception {
         setHasCarrierPrivileges(true);
         setIsDeviceOwner(false);
@@ -78,6 +82,7 @@
                 NetworkStatsAccess.checkAccessLevel(mContext, TEST_UID, TEST_PKG));
     }
 
+    @Test
     public void testCheckAccessLevel_isDeviceOwner() throws Exception {
         setHasCarrierPrivileges(false);
         setIsDeviceOwner(true);
@@ -88,6 +93,7 @@
                 NetworkStatsAccess.checkAccessLevel(mContext, TEST_UID, TEST_PKG));
     }
 
+    @Test
     public void testCheckAccessLevel_isProfileOwner() throws Exception {
         setHasCarrierPrivileges(false);
         setIsDeviceOwner(false);
@@ -98,36 +104,40 @@
                 NetworkStatsAccess.checkAccessLevel(mContext, TEST_UID, TEST_PKG));
     }
 
+    @Test
     public void testCheckAccessLevel_hasAppOpsBitAllowed() throws Exception {
         setHasCarrierPrivileges(false);
         setIsDeviceOwner(false);
         setIsProfileOwner(true);
         setHasAppOpsPermission(AppOpsManager.MODE_ALLOWED, false);
         setHasReadHistoryPermission(false);
-        assertEquals(NetworkStatsAccess.Level.USER,
+        assertEquals(NetworkStatsAccess.Level.DEVICESUMMARY,
                 NetworkStatsAccess.checkAccessLevel(mContext, TEST_UID, TEST_PKG));
     }
 
+    @Test
     public void testCheckAccessLevel_hasAppOpsBitDefault_grantedPermission() throws Exception {
         setHasCarrierPrivileges(false);
         setIsDeviceOwner(false);
         setIsProfileOwner(true);
         setHasAppOpsPermission(AppOpsManager.MODE_DEFAULT, true);
         setHasReadHistoryPermission(false);
-        assertEquals(NetworkStatsAccess.Level.USER,
+        assertEquals(NetworkStatsAccess.Level.DEVICESUMMARY,
                 NetworkStatsAccess.checkAccessLevel(mContext, TEST_UID, TEST_PKG));
     }
 
+    @Test
     public void testCheckAccessLevel_hasReadHistoryPermission() throws Exception {
         setHasCarrierPrivileges(false);
         setIsDeviceOwner(false);
         setIsProfileOwner(true);
         setHasAppOpsPermission(AppOpsManager.MODE_DEFAULT, false);
         setHasReadHistoryPermission(true);
-        assertEquals(NetworkStatsAccess.Level.USER,
+        assertEquals(NetworkStatsAccess.Level.DEVICESUMMARY,
                 NetworkStatsAccess.checkAccessLevel(mContext, TEST_UID, TEST_PKG));
     }
 
+    @Test
     public void testCheckAccessLevel_deniedAppOpsBit() throws Exception {
         setHasCarrierPrivileges(false);
         setIsDeviceOwner(false);
@@ -138,6 +148,7 @@
                 NetworkStatsAccess.checkAccessLevel(mContext, TEST_UID, TEST_PKG));
     }
 
+    @Test
     public void testCheckAccessLevel_deniedAppOpsBit_deniedPermission() throws Exception {
         setHasCarrierPrivileges(false);
         setIsDeviceOwner(false);
diff --git a/tests/net/java/com/android/server/net/NetworkStatsObserversTest.java b/tests/net/java/com/android/server/net/NetworkStatsObserversTest.java
index fe7840d..92dcdac 100644
--- a/tests/net/java/com/android/server/net/NetworkStatsObserversTest.java
+++ b/tests/net/java/com/android/server/net/NetworkStatsObserversTest.java
@@ -18,21 +18,23 @@
 
 import static android.net.ConnectivityManager.TYPE_MOBILE;
 import static android.net.ConnectivityManager.TYPE_WIFI;
-import static android.text.format.DateUtils.MINUTE_IN_MILLIS;
-import static org.mockito.Matchers.any;
-import static org.mockito.Matchers.anyInt;
-import static org.mockito.Matchers.isA;
-import static org.mockito.Mockito.when;
-
-import static android.net.NetworkStats.SET_DEFAULT;
 import static android.net.NetworkStats.METERED_NO;
 import static android.net.NetworkStats.ROAMING_NO;
+import static android.net.NetworkStats.SET_DEFAULT;
 import static android.net.NetworkStats.TAG_NONE;
 import static android.net.NetworkTemplate.buildTemplateMobileAll;
 import static android.net.NetworkTemplate.buildTemplateWifiWildcard;
 import static android.net.TrafficStats.MB_IN_BYTES;
 import static android.text.format.DateUtils.MINUTE_IN_MILLIS;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyInt;
+import static org.mockito.Matchers.isA;
+import static org.mockito.Mockito.when;
+
 import android.app.usage.NetworkStatsManager;
 import android.net.DataUsageRequest;
 import android.net.NetworkIdentity;
@@ -48,6 +50,7 @@
 import android.os.Process;
 import android.os.UserHandle;
 import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
 import android.telephony.TelephonyManager;
 import android.util.ArrayMap;
 
@@ -60,8 +63,9 @@
 import java.util.Objects;
 import java.util.List;
 
-import junit.framework.TestCase;
-
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
@@ -69,8 +73,9 @@
 /**
  * Tests for {@link NetworkStatsObservers}.
  */
+@RunWith(AndroidJUnit4.class)
 @SmallTest
-public class NetworkStatsObserversTest extends TestCase {
+public class NetworkStatsObserversTest {
     private static final String TEST_IFACE = "test0";
     private static final String TEST_IFACE2 = "test1";
     private static final long TEST_START = 1194220800000L;
@@ -88,7 +93,7 @@
     private static final int UID_GREEN = UserHandle.PER_USER_RANGE + 3;
     private static final int UID_ANOTHER_USER = 2 * UserHandle.PER_USER_RANGE + 4;
 
-    private static final long WAIT_TIMEOUT = 500;  // 1/2 sec
+    private static final long WAIT_TIMEOUT_MS = 500;
     private static final long THRESHOLD_BYTES = 2 * MB_IN_BYTES;
     private static final long BASE_BYTES = 7 * MB_IN_BYTES;
     private static final int INVALID_TYPE = -1;
@@ -101,7 +106,6 @@
     private Handler mObserverNoopHandler;
 
     private LatchedHandler mHandler;
-    private ConditionVariable mCv;
 
     private NetworkStatsObservers mStatsObservers;
     private Messenger mMessenger;
@@ -110,9 +114,8 @@
 
     @Mock private IBinder mockBinder;
 
-    @Override
+    @Before
     public void setUp() throws Exception {
-        super.setUp();
         MockitoAnnotations.initMocks(this);
 
         mObserverHandlerThread = new IdleableHandlerThread("HandlerThread");
@@ -125,14 +128,14 @@
             }
         };
 
-        mCv = new ConditionVariable();
-        mHandler = new LatchedHandler(Looper.getMainLooper(), mCv);
+        mHandler = new LatchedHandler(Looper.getMainLooper(), new ConditionVariable());
         mMessenger = new Messenger(mHandler);
 
         mActiveIfaces = new ArrayMap<>();
         mActiveUidIfaces = new ArrayMap<>();
     }
 
+    @Test
     public void testRegister_thresholdTooLow_setsDefaultThreshold() throws Exception {
         long thresholdTooLowBytes = 1L;
         DataUsageRequest inputRequest = new DataUsageRequest(
@@ -145,6 +148,7 @@
         assertEquals(THRESHOLD_BYTES, request.thresholdInBytes);
     }
 
+    @Test
     public void testRegister_highThreshold_accepted() throws Exception {
         long highThresholdBytes = 2 * THRESHOLD_BYTES;
         DataUsageRequest inputRequest = new DataUsageRequest(
@@ -157,6 +161,7 @@
         assertEquals(highThresholdBytes, request.thresholdInBytes);
     }
 
+    @Test
     public void testRegister_twoRequests_twoIds() throws Exception {
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateWifi, THRESHOLD_BYTES);
@@ -174,6 +179,7 @@
         assertEquals(THRESHOLD_BYTES, request2.thresholdInBytes);
     }
 
+    @Test
     public void testUnregister_unknownRequest_noop() throws Exception {
         DataUsageRequest unknownRequest = new DataUsageRequest(
                 123456 /* id */, sTemplateWifi, THRESHOLD_BYTES);
@@ -181,6 +187,7 @@
         mStatsObservers.unregister(unknownRequest, UID_RED);
     }
 
+    @Test
     public void testUnregister_knownRequest_releasesCaller() throws Exception {
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
@@ -198,6 +205,7 @@
         Mockito.verify(mockBinder).unlinkToDeath(any(IBinder.DeathRecipient.class), anyInt());
     }
 
+    @Test
     public void testUnregister_knownRequest_invalidUid_doesNotUnregister() throws Exception {
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
@@ -215,6 +223,7 @@
         Mockito.verifyZeroInteractions(mockBinder);
     }
 
+    @Test
     public void testUpdateStats_initialSample_doesNotNotify() throws Exception {
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
@@ -240,11 +249,9 @@
                 xtSnapshot, uidSnapshot, mActiveIfaces, mActiveUidIfaces,
                 VPN_INFO, TEST_START);
         waitForObserverToIdle();
-
-        assertTrue(mCv.block(WAIT_TIMEOUT));
-        assertEquals(INVALID_TYPE, mHandler.mLastMessageType);
     }
 
+    @Test
     public void testUpdateStats_belowThreshold_doesNotNotify() throws Exception {
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
@@ -276,12 +283,10 @@
                 xtSnapshot, uidSnapshot, mActiveIfaces, mActiveUidIfaces,
                 VPN_INFO, TEST_START);
         waitForObserverToIdle();
-
-        assertTrue(mCv.block(WAIT_TIMEOUT));
-        mCv.block(WAIT_TIMEOUT);
-        assertEquals(INVALID_TYPE, mHandler.mLastMessageType);
     }
 
+
+    @Test
     public void testUpdateStats_deviceAccess_notifies() throws Exception {
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
@@ -314,11 +319,10 @@
                 xtSnapshot, uidSnapshot, mActiveIfaces, mActiveUidIfaces,
                 VPN_INFO, TEST_START);
         waitForObserverToIdle();
-
-        assertTrue(mCv.block(WAIT_TIMEOUT));
         assertEquals(NetworkStatsManager.CALLBACK_LIMIT_REACHED, mHandler.mLastMessageType);
     }
 
+    @Test
     public void testUpdateStats_defaultAccess_notifiesSameUid() throws Exception {
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
@@ -352,11 +356,10 @@
                 xtSnapshot, uidSnapshot, mActiveIfaces, mActiveUidIfaces,
                 VPN_INFO, TEST_START);
         waitForObserverToIdle();
-
-        assertTrue(mCv.block(WAIT_TIMEOUT));
         assertEquals(NetworkStatsManager.CALLBACK_LIMIT_REACHED, mHandler.mLastMessageType);
     }
 
+    @Test
     public void testUpdateStats_defaultAccess_usageOtherUid_doesNotNotify() throws Exception {
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
@@ -390,11 +393,9 @@
                 xtSnapshot, uidSnapshot, mActiveIfaces, mActiveUidIfaces,
                 VPN_INFO, TEST_START);
         waitForObserverToIdle();
-
-        assertTrue(mCv.block(WAIT_TIMEOUT));
-        assertEquals(INVALID_TYPE, mHandler.mLastMessageType);
     }
 
+    @Test
     public void testUpdateStats_userAccess_usageSameUser_notifies() throws Exception {
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
@@ -428,11 +429,10 @@
                 xtSnapshot, uidSnapshot, mActiveIfaces, mActiveUidIfaces,
                 VPN_INFO, TEST_START);
         waitForObserverToIdle();
-
-        assertTrue(mCv.block(WAIT_TIMEOUT));
         assertEquals(NetworkStatsManager.CALLBACK_LIMIT_REACHED, mHandler.mLastMessageType);
     }
 
+    @Test
     public void testUpdateStats_userAccess_usageAnotherUser_doesNotNotify() throws Exception {
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
@@ -467,14 +467,22 @@
                 xtSnapshot, uidSnapshot, mActiveIfaces, mActiveUidIfaces,
                 VPN_INFO, TEST_START);
         waitForObserverToIdle();
-
-        assertTrue(mCv.block(WAIT_TIMEOUT));
-        assertEquals(INVALID_TYPE, mHandler.mLastMessageType);
     }
 
     private void waitForObserverToIdle() {
-        // Send dummy message to make sure that any previous message has been handled
-        mHandler.sendMessage(mHandler.obtainMessage(-1));
-        mObserverHandlerThread.waitForIdle(WAIT_TIMEOUT);
+        waitForIdleLooper(mObserverHandlerThread.getLooper(), WAIT_TIMEOUT_MS);
+        waitForIdleLooper(mHandler.getLooper(), WAIT_TIMEOUT_MS);
     }
+
+    // TODO: unify with ConnectivityService.waitForIdleHandler and
+    // NetworkServiceStatsTest.IdleableHandlerThread
+    private static void waitForIdleLooper(Looper looper, long timeoutMs) {
+        final ConditionVariable cv = new ConditionVariable();
+        final Handler handler = new Handler(looper);
+        handler.post(() -> cv.open());
+        if (!cv.block(timeoutMs)) {
+            fail("Looper did not become idle after " + timeoutMs + " ms");
+        }
+    }
+
 }