diff --git a/tests/cts/net/AndroidManifest.xml b/tests/cts/net/AndroidManifest.xml
index bca2d2c..001e294 100644
--- a/tests/cts/net/AndroidManifest.xml
+++ b/tests/cts/net/AndroidManifest.xml
@@ -20,6 +20,7 @@
 
     <uses-permission android:name="android.permission.ACCESS_NETWORK_STATE" />
     <uses-permission android:name="android.permission.ACCESS_WIFI_STATE" />
+    <uses-permission android:name="android.permission.ACCESS_COARSE_LOCATION" />
     <uses-permission android:name="android.permission.CHANGE_NETWORK_STATE" />
     <uses-permission android:name="android.permission.CHANGE_WIFI_STATE" />
     <uses-permission android:name="android.permission.DISABLE_KEYGUARD" />
diff --git a/tests/cts/net/jni/NativeMultinetworkJni.c b/tests/cts/net/jni/NativeMultinetworkJni.c
index 9a5ab9c..ad56b51 100644
--- a/tests/cts/net/jni/NativeMultinetworkJni.c
+++ b/tests/cts/net/jni/NativeMultinetworkJni.c
@@ -93,7 +93,8 @@
 void sockaddr_ntop(const struct sockaddr *sa, socklen_t salen, char *dst, const size_t size) {
     char addrstr[INET6_ADDRSTRLEN];
     char portstr[sizeof("65535")];
-    char buf[sizeof(addrstr) + sizeof(portstr) + sizeof("[]:")];
+    char buf[kSockaddrStrLen+1];
+
     int ret = getnameinfo(sa, salen,
                           addrstr, sizeof(addrstr),
                           portstr, sizeof(portstr),
@@ -106,7 +107,7 @@
         sprintf(buf, "???");
     }
 
-    strlcpy(dst, buf, (strlen(buf) < size - 1) ? strlen(buf) : size - 1);
+    strlcpy(dst, buf, size);
 }
 
 JNIEXPORT jint Java_android_net_cts_MultinetworkApiTest_runDatagramCheck(
@@ -122,9 +123,8 @@
     struct addrinfo *res = NULL;
     net_handle_t handle = (net_handle_t) nethandle;
 
-    // Quoth Ian Swett:
-    //     "QUIC always uses 80 and 443, but only 443 is used for secure(HTTPS) traffic."
-    int rval = android_getaddrinfofornetwork(handle, kHostname, "80", &kHints, &res);
+    static const char kPort[] = "443";
+    int rval = android_getaddrinfofornetwork(handle, kHostname, kPort, &kHints, &res);
     if (rval != 0) {
         ALOGD("android_getaddrinfofornetwork(%llu, %s) returned rval=%d errno=%d",
               handle, kHostname, rval, errno);
@@ -150,9 +150,9 @@
         return -errno;
     }
 
-    char addrstr[kSockaddrStrLen];
+    char addrstr[kSockaddrStrLen+1];
     sockaddr_ntop(res->ai_addr, res->ai_addrlen, addrstr, sizeof(addrstr));
-    ALOGD("Attempting connect() to %s...", addrstr);
+    ALOGD("Attempting connect() to %s ...", addrstr);
 
     rval = connect(fd, res->ai_addr, res->ai_addrlen);
     if (rval != 0) {
@@ -172,10 +172,12 @@
     ALOGD("... from %s", addrstr);
 
     // Don't let reads or writes block indefinitely.
-    const struct timeval timeo = { 5, 0 };  // 5 seconds
+    const struct timeval timeo = { 2, 0 };  // 2 seconds
     setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &timeo, sizeof(timeo));
     setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &timeo, sizeof(timeo));
 
+    // For reference see:
+    //     https://tools.ietf.org/html/draft-tsvwg-quic-protocol-01#section-6.1
     uint8_t quic_packet[] = {
         0x0c,                    // public flags: 64bit conn ID, 8bit sequence number
         0, 0, 0, 0, 0, 0, 0, 0,  // 64bit connection ID
@@ -186,19 +188,36 @@
 
     arc4random_buf(quic_packet + 1, 8);  // random connection ID
 
-    ssize_t sent = send(fd, quic_packet, sizeof(quic_packet), 0);
-    if (sent < (ssize_t)sizeof(quic_packet)) {
-        ALOGD("send(QUIC packet) returned sent=%zd, errno=%d", sent, errno);
-        close(fd);
-        return -errno;
-    }
-
     uint8_t response[1500];
-    ssize_t rcvd = recv(fd, response, sizeof(response), 0);
+    ssize_t sent, rcvd;
+    static const int MAX_RETRIES = 5;
+    int i, errnum = 0;
+
+    for (i = 0; i < MAX_RETRIES; i++) {
+        sent = send(fd, quic_packet, sizeof(quic_packet), 0);
+        if (sent < (ssize_t)sizeof(quic_packet)) {
+            errnum = errno;
+            ALOGD("send(QUIC packet) returned sent=%zd, errno=%d", sent, errnum);
+            close(fd);
+            return -errnum;
+        }
+
+        rcvd = recv(fd, response, sizeof(response), 0);
+        if (rcvd > 0) {
+            break;
+        } else {
+            errnum = errno;
+            ALOGD("[%d/%d] recv(QUIC response) returned rcvd=%zd, errno=%d",
+                  i + 1, MAX_RETRIES, rcvd, errnum);
+        }
+    }
     if (rcvd < sent) {
-        ALOGD("recv() returned rcvd=%zd, errno=%d", rcvd, errno);
+        ALOGD("QUIC UDP %s: sent=%zd but rcvd=%zd, errno=%d", kPort, sent, rcvd, errnum);
+        if (rcvd <= 0) {
+            ALOGD("Does this network block UDP port %s?", kPort);
+        }
         close(fd);
-        return -errno;
+        return -EPROTO;
     }
 
     int conn_id_cmp = memcmp(quic_packet + 1, response + 1, 8);
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index 7c70293..88dbd7c 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -16,6 +16,7 @@
 
 package android.net.cts;
 
+import android.app.PendingIntent;
 import android.content.BroadcastReceiver;
 import android.content.Context;
 import android.content.Intent;
@@ -23,13 +24,16 @@
 import android.content.pm.PackageManager;
 import android.net.ConnectivityManager;
 import android.net.Network;
+import android.net.NetworkCapabilities;
 import android.net.NetworkConfig;
 import android.net.NetworkInfo;
 import android.net.NetworkInfo.DetailedState;
 import android.net.NetworkInfo.State;
+import android.net.NetworkRequest;
 import android.net.wifi.WifiManager;
 import android.test.AndroidTestCase;
 import android.util.Log;
+import android.os.SystemProperties;
 
 import com.android.internal.telephony.PhoneConstants;
 
@@ -51,6 +55,10 @@
     public static final int TYPE_WIFI = ConnectivityManager.TYPE_WIFI;
     private static final int HOST_ADDRESS = 0x7f000001;// represent ip 127.0.0.1
 
+    // Action sent to ConnectivityActionReceiver when a network callback is sent via PendingIntent.
+    private static final String NETWORK_CALLBACK_ACTION =
+            "ConnectivityManagerTest.NetworkCallbackAction";
+
     // device could have only one interface: data, wifi.
     private static final int MIN_NUM_NETWORK_TYPES = 1;
 
@@ -59,7 +67,6 @@
     private PackageManager mPackageManager;
     private final HashMap<Integer, NetworkConfig> mNetworks =
             new HashMap<Integer, NetworkConfig>();
-    private final List<Integer>mProtectedNetworks = new ArrayList<Integer>();
 
     @Override
     protected void setUp() throws Exception {
@@ -71,20 +78,17 @@
         // Get com.android.internal.R.array.networkAttributes
         int resId = getContext().getResources().getIdentifier("networkAttributes", "array", "android");
         String[] naStrings = getContext().getResources().getStringArray(resId);
-
+        //TODO: What is the "correct" way to determine if this is a wifi only device?
+        boolean wifiOnly = SystemProperties.getBoolean("ro.radio.noril", false);
         for (String naString : naStrings) {
             try {
                 NetworkConfig n = new NetworkConfig(naString);
+                if (wifiOnly && ConnectivityManager.isNetworkTypeMobile(n.type)) {
+                    continue;
+                }
                 mNetworks.put(n.type, n);
             } catch (Exception e) {}
         }
-
-        // Get com.android.internal.R.array.config_protectedNetworks
-        resId = getContext().getResources().getIdentifier("config_protectedNetworks", "array", "android");
-        int[] protectedNetworks = getContext().getResources().getIntArray(resId);
-        for (int p : protectedNetworks) {
-            mProtectedNetworks.add(p);
-        }
     }
 
     public void testIsNetworkTypeValid() {
@@ -125,17 +129,18 @@
     public void testGetActiveNetworkInfo() {
         NetworkInfo ni = mCm.getActiveNetworkInfo();
 
-        assertTrue("You must have an active network connection to complete CTS", ni != null);
+        assertNotNull("You must have an active network connection to complete CTS", ni);
         assertTrue(ConnectivityManager.isNetworkTypeValid(ni.getType()));
         assertTrue(ni.getState() == State.CONNECTED);
     }
 
     public void testGetActiveNetwork() {
         Network network = mCm.getActiveNetwork();
-        assertTrue("You must have an active network connection to complete CTS", network != null);
+        assertNotNull("You must have an active network connection to complete CTS", network);
 
         NetworkInfo ni = mCm.getNetworkInfo(network);
-        assertTrue("Network returned from getActiveNetwork was invalid", ni != null);
+        assertNotNull("Network returned from getActiveNetwork was invalid", ni);
+
         // Similar to testGetActiveNetworkInfo above.
         assertTrue(ConnectivityManager.isNetworkTypeValid(ni.getType()));
         assertTrue(ni.getState() == State.CONNECTED);
@@ -177,6 +182,27 @@
         }
     }
 
+    private void assertStartUsingNetworkFeatureUnsupported(int networkType, String feature) {
+        try {
+            mCm.startUsingNetworkFeature(networkType, feature);
+            fail("startUsingNetworkFeature is no longer supported in the current API version");
+        } catch (UnsupportedOperationException expected) {}
+    }
+
+    private void assertStopUsingNetworkFeatureUnsupported(int networkType, String feature) {
+        try {
+            mCm.startUsingNetworkFeature(networkType, feature);
+            fail("stopUsingNetworkFeature is no longer supported in the current API version");
+        } catch (UnsupportedOperationException expected) {}
+    }
+
+    private void assertRequestRouteToHostUnsupported(int networkType, int hostAddress) {
+        try {
+            mCm.requestRouteToHost(networkType, hostAddress);
+            fail("requestRouteToHost is no longer supported in the current API version");
+        } catch (UnsupportedOperationException expected) {}
+    }
+
     public void testStartUsingNetworkFeature() {
 
         final String invalidateFeature = "invalidateFeature";
@@ -185,27 +211,9 @@
         final int wifiOnlyStartFailureCode = PhoneConstants.APN_REQUEST_FAILED;
         final int wifiOnlyStopFailureCode = -1;
 
-        NetworkInfo ni = mCm.getNetworkInfo(TYPE_MOBILE);
-        if (ni != null) {
-            assertEquals(PhoneConstants.APN_REQUEST_FAILED,
-                    mCm.startUsingNetworkFeature(TYPE_MOBILE, invalidateFeature));
-            assertEquals(failureCode, mCm.stopUsingNetworkFeature(TYPE_MOBILE,
-                    invalidateFeature));
-        } else {
-            assertEquals(wifiOnlyStartFailureCode, mCm.startUsingNetworkFeature(TYPE_MOBILE,
-                    invalidateFeature));
-            assertEquals(wifiOnlyStopFailureCode, mCm.stopUsingNetworkFeature(TYPE_MOBILE,
-                    invalidateFeature));
-        }
-
-        ni = mCm.getNetworkInfo(TYPE_WIFI);
-        if (ni != null) {
-            // Should return failure because MMS is not supported on WIFI.
-            assertEquals(PhoneConstants.APN_REQUEST_FAILED, mCm.startUsingNetworkFeature(TYPE_WIFI,
-                    mmsFeature));
-            assertEquals(failureCode, mCm.stopUsingNetworkFeature(TYPE_WIFI,
-                    mmsFeature));
-        }
+        assertStartUsingNetworkFeatureUnsupported(TYPE_MOBILE, invalidateFeature);
+        assertStopUsingNetworkFeatureUnsupported(TYPE_MOBILE, invalidateFeature);
+        assertStartUsingNetworkFeatureUnsupported(TYPE_WIFI, mmsFeature);
     }
 
     private boolean isSupported(int networkType) {
@@ -216,11 +224,6 @@
                (networkType == ConnectivityManager.TYPE_VPN);
     }
 
-    // true if only the system can turn it on
-    private boolean isNetworkProtected(int networkType) {
-        return mProtectedNetworks.contains(networkType);
-    }
-
     public void testIsNetworkSupported() {
         for (int type = -1; type <= ConnectivityManager.MAX_NETWORK_TYPE; type++) {
             boolean supported = mCm.isNetworkSupported(type);
@@ -234,120 +237,234 @@
 
     public void testRequestRouteToHost() {
         for (int type = -1 ; type <= ConnectivityManager.MAX_NETWORK_TYPE; type++) {
-            NetworkInfo ni = mCm.getNetworkInfo(type);
-            boolean expectToWork = isSupported(type) && !isNetworkProtected(type) &&
-                    ni != null && ni.isConnected();
-
-            try {
-                assertTrue("Network type " + type,
-                        mCm.requestRouteToHost(type, HOST_ADDRESS) == expectToWork);
-            } catch (Exception e) {
-                Log.d(TAG, "got exception in requestRouteToHost for type " + type);
-                assertFalse("Exception received for type " + type, expectToWork);
-            }
-
-            //TODO verify route table
+            assertRequestRouteToHostUnsupported(type, HOST_ADDRESS);
         }
-
-        assertFalse(mCm.requestRouteToHost(-1, HOST_ADDRESS));
     }
 
     public void testTest() {
         mCm.getBackgroundDataSetting();
     }
 
-    /** Test that hipri can be brought up when Wifi is enabled. */
-    public void testStartUsingNetworkFeature_enableHipri() throws Exception {
-        if (!mPackageManager.hasSystemFeature(PackageManager.FEATURE_TELEPHONY)
-                || !mPackageManager.hasSystemFeature(PackageManager.FEATURE_WIFI)) {
-            // This test requires a mobile data connection and WiFi.
+    /**
+     * Exercises both registerNetworkCallback and unregisterNetworkCallback. This checks to
+     * see if we get a callback for the TRANSPORT_WIFI transport type being available.
+     *
+     * <p>In order to test that a NetworkCallback occurs, we need some change in the network
+     * state (either a transport or capability is now available). The most straightforward is
+     * WiFi. We could add a version that uses the telephony data connection but it's not clear
+     * that it would increase test coverage by much (how many devices have 3G radio but not Wifi?).
+     */
+    public void testRegisterNetworkCallback() {
+        if (!mPackageManager.hasSystemFeature(PackageManager.FEATURE_WIFI)) {
+            Log.i(TAG, "testRegisterNetworkCallback cannot execute unless device supports WiFi");
             return;
         }
 
-        boolean isWifiEnabled = mWifiManager.isWifiEnabled();
-        boolean isWifiConnected = false;
+        // We will register for a WIFI network being available or lost.
+        NetworkRequest request = new NetworkRequest.Builder()
+                .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
+                .build();
+        TestNetworkCallback callback = new TestNetworkCallback();
+        mCm.registerNetworkCallback(request, callback);
 
-        NetworkInfo nwInfo = mCm.getNetworkInfo(ConnectivityManager.TYPE_WIFI);
-        if (nwInfo != null) {
-            isWifiConnected = nwInfo.isConnected();
-        }
+        boolean previousWifiEnabledState = mWifiManager.isWifiEnabled();
+
         try {
-            // Make sure WiFi is connected to an access point.
-            if (!isWifiConnected) {
+            // Make sure WiFi is connected to an access point to start with.
+            if (!previousWifiEnabledState) {
                 connectToWifi();
             }
 
-            // Register a receiver that will capture the connectivity change for hipri.
-            ConnectivityActionReceiver receiver =
-                    new ConnectivityActionReceiver(ConnectivityManager.TYPE_MOBILE_HIPRI);
-            IntentFilter filter = new IntentFilter();
-            filter.addAction(ConnectivityManager.CONNECTIVITY_ACTION);
-            mContext.registerReceiver(receiver, filter);
-
-            // Try to start using the hipri feature...
-            int result = mCm.startUsingNetworkFeature(ConnectivityManager.TYPE_MOBILE,
-                    FEATURE_ENABLE_HIPRI);
-            assertTrue("Couldn't start using the HIPRI feature.", result != -1);
-
-            // Check that the ConnectivityManager reported that it connected using hipri...
-            assertTrue("Couldn't connect using hipri...", receiver.waitForConnection());
-
-            assertTrue("Couldn't requestRouteToHost using HIPRI.",
-                    mCm.requestRouteToHost(ConnectivityManager.TYPE_MOBILE_HIPRI, HOST_ADDRESS));
-            // TODO check dns selection
-            // TODO check routes
+            // Now we should expect to get a network callback about availability of the wifi
+            // network even if it was already connected as a state-based action when the callback
+            // is registered.
+            assertTrue("Did not receive NetworkCallback.onAvailable for TRANSPORT_WIFI",
+                    callback.waitForAvailable());
         } catch (InterruptedException e) {
-            fail("Broadcast receiver waiting for ConnectivityManager interrupted.");
+            fail("Broadcast receiver or NetworkCallback wait was interrupted.");
         } finally {
-            mCm.stopUsingNetworkFeature(ConnectivityManager.TYPE_MOBILE,
-                    FEATURE_ENABLE_HIPRI);
-            // TODO wait for HIPRI to go
-            // TODO check dns selection
-            // TODO check routes
-            if (!isWifiEnabled) {
-                mWifiManager.setWifiEnabled(false);
+            mCm.unregisterNetworkCallback(callback);
+
+            // Return WiFI to its original enabled/disabled state.
+            if (!previousWifiEnabledState) {
+                disconnectFromWifi();
             }
         }
     }
 
-    private void connectToWifi() throws InterruptedException {
-        ConnectivityActionReceiver receiver =
-                new ConnectivityActionReceiver(ConnectivityManager.TYPE_WIFI);
+    /**
+     * Tests both registerNetworkCallback and unregisterNetworkCallback similarly to
+     * {@link #testRegisterNetworkCallback} except that a {@code PendingIntent} is used instead
+     * of a {@code NetworkCallback}.
+     */
+    public void testRegisterNetworkCallback_withPendingIntent() {
+        if (!mPackageManager.hasSystemFeature(PackageManager.FEATURE_WIFI)) {
+            Log.i(TAG, "testRegisterNetworkCallback cannot execute unless device supports WiFi");
+            return;
+        }
+
+        // Create a ConnectivityActionReceiver that has an IntentFilter for our locally defined
+        // action, NETWORK_CALLBACK_ACTION.
+        IntentFilter filter = new IntentFilter();
+        filter.addAction(NETWORK_CALLBACK_ACTION);
+
+        ConnectivityActionReceiver receiver = new ConnectivityActionReceiver(
+                ConnectivityManager.TYPE_WIFI, NetworkInfo.State.CONNECTED);
+        mContext.registerReceiver(receiver, filter);
+
+        // Create a broadcast PendingIntent for NETWORK_CALLBACK_ACTION.
+        Intent intent = new Intent(NETWORK_CALLBACK_ACTION);
+        PendingIntent pendingIntent = PendingIntent.getBroadcast(
+                mContext, 0, intent, PendingIntent.FLAG_CANCEL_CURRENT);
+
+        // We will register for a WIFI network being available or lost.
+        NetworkRequest request = new NetworkRequest.Builder()
+                .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
+                .build();
+        mCm.registerNetworkCallback(request, pendingIntent);
+
+        boolean previousWifiEnabledState = mWifiManager.isWifiEnabled();
+
+        try {
+            // Make sure WiFi is connected to an access point to start with.
+            if (!previousWifiEnabledState) {
+                connectToWifi();
+            }
+
+            // Now we expect to get the Intent delivered notifying of the availability of the wifi
+            // network even if it was already connected as a state-based action when the callback
+            // is registered.
+            assertTrue("Did not receive expected Intent " + intent + " for TRANSPORT_WIFI",
+                    receiver.waitForState());
+        } catch (InterruptedException e) {
+            fail("Broadcast receiver or NetworkCallback wait was interrupted.");
+        } finally {
+            mCm.unregisterNetworkCallback(pendingIntent);
+            pendingIntent.cancel();
+            mContext.unregisterReceiver(receiver);
+
+            // Return WiFI to its original enabled/disabled state.
+            if (!previousWifiEnabledState) {
+                disconnectFromWifi();
+            }
+        }
+    }
+
+    /** Enable WiFi and wait for it to become connected to a network. */
+    private void connectToWifi() {
+        ConnectivityActionReceiver receiver = new ConnectivityActionReceiver(
+                ConnectivityManager.TYPE_WIFI, NetworkInfo.State.CONNECTED);
         IntentFilter filter = new IntentFilter();
         filter.addAction(ConnectivityManager.CONNECTIVITY_ACTION);
         mContext.registerReceiver(receiver, filter);
 
-        assertTrue(mWifiManager.setWifiEnabled(true));
-        assertTrue("Wifi must be configured to connect to an access point for this test.",
-                receiver.waitForConnection());
+        boolean connected = false;
+        try {
+            assertTrue(mWifiManager.setWifiEnabled(true));
+            connected = receiver.waitForState();
+        } catch (InterruptedException ex) {
+            fail("connectToWifi was interrupted");
+        } finally {
+            mContext.unregisterReceiver(receiver);
+        }
 
-        mContext.unregisterReceiver(receiver);
+        assertTrue("Wifi must be configured to connect to an access point for this test.",
+                connected);
     }
 
-    /** Receiver that captures the last connectivity change's network type and state. */
+    /** Disable WiFi and wait for it to become disconnected from the network. */
+    private void disconnectFromWifi() {
+        ConnectivityActionReceiver receiver = new ConnectivityActionReceiver(
+                ConnectivityManager.TYPE_WIFI, NetworkInfo.State.DISCONNECTED);
+        IntentFilter filter = new IntentFilter();
+        filter.addAction(ConnectivityManager.CONNECTIVITY_ACTION);
+        mContext.registerReceiver(receiver, filter);
+
+        boolean disconnected = false;
+        try {
+            assertTrue(mWifiManager.setWifiEnabled(false));
+            disconnected = receiver.waitForState();
+        } catch (InterruptedException ex) {
+            fail("disconnectFromWifi was interrupted");
+        } finally {
+            mContext.unregisterReceiver(receiver);
+        }
+
+        assertTrue("Wifi failed to reach DISCONNECTED state.", disconnected);
+    }
+
+    /**
+     * Receiver that captures the last connectivity change's network type and state. Recognizes
+     * both {@code CONNECTIVITY_ACTION} and {@code NETWORK_CALLBACK_ACTION} intents.
+     */
     private class ConnectivityActionReceiver extends BroadcastReceiver {
 
         private final CountDownLatch mReceiveLatch = new CountDownLatch(1);
 
         private final int mNetworkType;
+        private final NetworkInfo.State mNetState;
 
-        ConnectivityActionReceiver(int networkType) {
+        ConnectivityActionReceiver(int networkType, NetworkInfo.State netState) {
             mNetworkType = networkType;
+            mNetState = netState;
         }
 
         public void onReceive(Context context, Intent intent) {
-            NetworkInfo networkInfo = intent.getExtras()
-                    .getParcelable(ConnectivityManager.EXTRA_NETWORK_INFO);
+            String action = intent.getAction();
+            NetworkInfo networkInfo = null;
+
+            // When receiving ConnectivityManager.CONNECTIVITY_ACTION, the NetworkInfo parcelable
+            // is stored in EXTRA_NETWORK_INFO. With a NETWORK_CALLBACK_ACTION, the Network is
+            // sent in EXTRA_NETWORK and we need to ask the ConnectivityManager for the NetworkInfo.
+            if (ConnectivityManager.CONNECTIVITY_ACTION.equals(action)) {
+                networkInfo = intent.getExtras()
+                        .getParcelable(ConnectivityManager.EXTRA_NETWORK_INFO);
+                assertNotNull("ConnectivityActionReceiver expected EXTRA_NETWORK_INFO", networkInfo);
+            } else if (NETWORK_CALLBACK_ACTION.equals(action)) {
+                Network network = intent.getExtras()
+                        .getParcelable(ConnectivityManager.EXTRA_NETWORK);
+                assertNotNull("ConnectivityActionReceiver expected EXTRA_NETWORK", network);
+                networkInfo = mCm.getNetworkInfo(network);
+                if (networkInfo == null) {
+                    // When disconnecting, it seems like we get an intent sent with an invalid
+                    // Network; that is, by the time we call ConnectivityManager.getNetworkInfo(),
+                    // it is invalid. Ignore these.
+                    Log.i(TAG, "ConnectivityActionReceiver NETWORK_CALLBACK_ACTION ignoring "
+                            + "invalid network");
+                    return;
+                }
+            } else {
+                fail("ConnectivityActionReceiver received unxpected intent action: " + action);
+            }
+
+            assertNotNull("ConnectivityActionReceiver didn't find NetworkInfo", networkInfo);
             int networkType = networkInfo.getType();
             State networkState = networkInfo.getState();
             Log.i(TAG, "Network type: " + networkType + " state: " + networkState);
-            if (networkType == mNetworkType && networkInfo.getState() == State.CONNECTED) {
+            if (networkType == mNetworkType && networkInfo.getState() == mNetState) {
                 mReceiveLatch.countDown();
             }
         }
 
-        public boolean waitForConnection() throws InterruptedException {
+        public boolean waitForState() throws InterruptedException {
             return mReceiveLatch.await(30, TimeUnit.SECONDS);
         }
     }
+
+    /**
+     * Callback used in testRegisterNetworkCallback that allows caller to block on
+     * {@code onAvailable}.
+     */
+    private static class TestNetworkCallback extends ConnectivityManager.NetworkCallback {
+        private final CountDownLatch mAvailableLatch = new CountDownLatch(1);
+
+        public boolean waitForAvailable() throws InterruptedException {
+            return mAvailableLatch.await(30, TimeUnit.SECONDS);
+        }
+
+        @Override
+        public void onAvailable(Network network) {
+            mAvailableLatch.countDown();
+        }
+    }
 }