diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerBaseTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerBaseTest.java
index 020bc0078d807..d85b0594344f5 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerBaseTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerBaseTest.java
@@ -4,9 +4,17 @@ package com.android.systemui.statusbar.policy;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import java.io.PrintWriter;
+import java.io.StringWriter;
+import java.util.ArrayList;
+import java.util.List;
+
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
+
 import android.content.Intent;
 import android.net.ConnectivityManager;
-import android.net.NetworkInfo;
+import android.net.NetworkCapabilities;
 import android.net.wifi.WifiManager;
 import android.telephony.PhoneStateListener;
 import android.telephony.ServiceState;
@@ -22,14 +30,6 @@ import com.android.systemui.statusbar.policy.NetworkController.NetworkSignalChan
 import com.android.systemui.statusbar.policy.NetworkControllerImpl.Config;
 import com.android.systemui.statusbar.policy.NetworkControllerImpl.SignalCluster;
 
-import org.mockito.ArgumentCaptor;
-import org.mockito.Mockito;
-
-import java.io.PrintWriter;
-import java.io.StringWriter;
-import java.util.ArrayList;
-import java.util.List;
-
 public class NetworkControllerBaseTest extends AndroidTestCase {
     private static final String TAG = "NetworkControllerBaseTest";
     protected static final int DEFAULT_LEVEL = 2;
@@ -52,6 +52,8 @@ public class NetworkControllerBaseTest extends AndroidTestCase {
     protected TelephonyManager mMockTm;
     protected Config mConfig;
 
+    private NetworkCapabilities mNetCapabilities;
+
     @Override
     protected void setUp() throws Exception {
         super.setUp();
@@ -63,7 +65,10 @@ public class NetworkControllerBaseTest extends AndroidTestCase {
         mMockTm = mock(TelephonyManager.class);
         mMockSm = mock(SubscriptionManager.class);
         mMockCm = mock(ConnectivityManager.class);
+        mNetCapabilities = new NetworkCapabilities();
         when(mMockCm.isNetworkSupported(ConnectivityManager.TYPE_MOBILE)).thenReturn(true);
+        when(mMockCm.getDefaultNetworkCapabilitiesForUser(0)).thenReturn(
+                new NetworkCapabilities[] { mNetCapabilities });
 
         mSignalStrength = mock(SignalStrength.class);
         mServiceState = mock(ServiceState.class);
@@ -115,13 +120,18 @@ public class NetworkControllerBaseTest extends AndroidTestCase {
 
     public void setConnectivity(int inetCondition, int networkType, boolean isConnected) {
         Intent i = new Intent(ConnectivityManager.INET_CONDITION_ACTION);
-        NetworkInfo networkInfo = mock(NetworkInfo.class);
-        when(networkInfo.isConnected()).thenReturn(isConnected);
-        when(networkInfo.getType()).thenReturn(networkType);
-        when(networkInfo.getTypeName()).thenReturn("");
-        when(mMockCm.getActiveNetworkInfo()).thenReturn(networkInfo);
+        // TODO: Separate out into several NetworkCapabilities.
+        if (isConnected) {
+            mNetCapabilities.addTransportType(networkType);
+        } else {
+            mNetCapabilities.removeTransportType(networkType);
+        }
+        if (inetCondition != 0) {
+            mNetCapabilities.addCapability(NetworkCapabilities.NET_CAPABILITY_VALIDATED);
+        } else {
+            mNetCapabilities.removeCapability(NetworkCapabilities.NET_CAPABILITY_VALIDATED);
+        }
 
-        i.putExtra(ConnectivityManager.EXTRA_INET_CONDITION, inetCondition);
         mNetworkController.onReceive(mContext, i);
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerSignalTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerSignalTest.java
index 33eb4d67ba27b..27a4052f3941b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerSignalTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerSignalTest.java
@@ -60,7 +60,7 @@ public class NetworkControllerSignalTest extends NetworkControllerBaseTest {
                     DEFAULT_ICON);
 
             // Verify low inet number indexing.
-            setConnectivity(10, ConnectivityManager.TYPE_MOBILE, true);
+            setConnectivity(0, ConnectivityManager.TYPE_MOBILE, true);
             verifyLastMobileDataIndicators(true,
                     TelephonyIcons.TELEPHONY_SIGNAL_STRENGTH[0][testStrength], 0);
         }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerWifiTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerWifiTest.java
index 7f0a8f4dd1ae7..2e0e9a3c14f97 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerWifiTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerWifiTest.java
@@ -27,7 +27,7 @@ public class NetworkControllerWifiTest extends NetworkControllerBaseTest {
 
             setConnectivity(100, ConnectivityManager.TYPE_WIFI, true);
             verifyLastWifiIcon(true, WifiIcons.WIFI_SIGNAL_STRENGTH[1][testLevel]);
-            setConnectivity(10, ConnectivityManager.TYPE_WIFI, true);
+            setConnectivity(0, ConnectivityManager.TYPE_WIFI, true);
             verifyLastWifiIcon(true, WifiIcons.WIFI_SIGNAL_STRENGTH[0][testLevel]);
         }
     }
@@ -48,7 +48,7 @@ public class NetworkControllerWifiTest extends NetworkControllerBaseTest {
             setConnectivity(100, ConnectivityManager.TYPE_WIFI, true);
             verifyLastQsWifiIcon(true, true, WifiIcons.QS_WIFI_SIGNAL_STRENGTH[1][testLevel],
                     testSsid);
-            setConnectivity(10, ConnectivityManager.TYPE_WIFI, true);
+            setConnectivity(0, ConnectivityManager.TYPE_WIFI, true);
             verifyLastQsWifiIcon(true, true, WifiIcons.QS_WIFI_SIGNAL_STRENGTH[0][testLevel],
                     testSsid);
         }
@@ -75,22 +75,6 @@ public class NetworkControllerWifiTest extends NetworkControllerBaseTest {
         verifyLastQsDataDirection(true, true);
     }
 
-    public void testNoDataIconDuringWifi() {
-        // Setup normal connection
-        String testSsid = "Test SSID";
-        int testLevel = 2;
-        setWifiEnabled(true);
-        setWifiState(true, testSsid);
-        setWifiLevel(testLevel);
-        setConnectivity(100, ConnectivityManager.TYPE_WIFI, true);
-        verifyLastWifiIcon(true, WifiIcons.WIFI_SIGNAL_STRENGTH[1][testLevel]);
-
-        setupDefaultSignal();
-        // Still be on wifi though.
-        setConnectivity(100, ConnectivityManager.TYPE_WIFI, true);
-        verifyLastMobileDataIndicators(true, DEFAULT_SIGNAL_STRENGTH, 0 /* No icon */);
-    }
-
     public void testRoamingIconDuringWifi() {
         // Setup normal connection
         String testSsid = "Test SSID";