Add the ability to affect ordering in the tethering unit tests. am: 3e09f80e8a am: bdb012f733

Original change: https://android-review.googlesource.com/c/platform/packages/modules/Connectivity/+/1635021

Change-Id: Iee27da6e7f6972564ee062e789905197e8994aa9
diff --git a/Tethering/tests/unit/src/com/android/networkstack/tethering/TestConnectivityManager.java b/Tethering/tests/unit/src/com/android/networkstack/tethering/TestConnectivityManager.java
index 3636b03..7ac80b0 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/TestConnectivityManager.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/TestConnectivityManager.java
@@ -30,12 +30,10 @@
 import android.net.NetworkRequest;
 import android.os.Handler;
 import android.os.UserHandle;
+import android.util.ArrayMap;
 
-import java.util.HashMap;
-import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Set;
 
 /**
  * Simulates upstream switching and sending NetworkCallbacks and CONNECTIVITY_ACTION broadcasts.
@@ -60,12 +58,12 @@
  *   that state changes), this may become less important or unnecessary.
  */
 public class TestConnectivityManager extends ConnectivityManager {
-    public Map<NetworkCallback, Handler> allCallbacks = new HashMap<>();
-    public Set<NetworkCallback> trackingDefault = new HashSet<>();
+    public Map<NetworkCallback, NetworkRequestInfo> allCallbacks = new ArrayMap<>();
+    public Map<NetworkCallback, NetworkRequestInfo> trackingDefault = new ArrayMap<>();
     public TestNetworkAgent defaultNetwork = null;
-    public Map<NetworkCallback, NetworkRequest> listening = new HashMap<>();
-    public Map<NetworkCallback, NetworkRequest> requested = new HashMap<>();
-    public Map<NetworkCallback, Integer> legacyTypeMap = new HashMap<>();
+    public Map<NetworkCallback, NetworkRequestInfo> listening = new ArrayMap<>();
+    public Map<NetworkCallback, NetworkRequestInfo> requested = new ArrayMap<>();
+    public Map<NetworkCallback, Integer> legacyTypeMap = new ArrayMap<>();
 
     private final NetworkRequest mDefaultRequest;
     private final Context mContext;
@@ -86,6 +84,15 @@
         mDefaultRequest = defaultRequest;
     }
 
+    class NetworkRequestInfo {
+        public final NetworkRequest request;
+        public final Handler handler;
+        NetworkRequestInfo(NetworkRequest r, Handler h) {
+            request = r;
+            handler = h;
+        }
+    }
+
     boolean hasNoCallbacks() {
         return allCallbacks.isEmpty()
                 && trackingDefault.isEmpty()
@@ -106,8 +113,8 @@
         final NetworkCapabilities empty = new NetworkCapabilities();
         empty.clearAll();
 
-        for (NetworkRequest req : listening.values()) {
-            if (req.networkCapabilities.equalRequestableCapabilities(empty)) {
+        for (NetworkRequestInfo nri : listening.values()) {
+            if (nri.request.networkCapabilities.equalRequestableCapabilities(empty)) {
                 return true;
             }
         }
@@ -118,40 +125,52 @@
         return ++mNetworkId;
     }
 
-    void makeDefaultNetwork(TestNetworkAgent agent) {
-        if (Objects.equals(defaultNetwork, agent)) return;
-
-        final TestNetworkAgent formerDefault = defaultNetwork;
-        defaultNetwork = agent;
-
+    private void sendDefaultNetworkBroadcasts(TestNetworkAgent formerDefault,
+            TestNetworkAgent defaultNetwork) {
         if (formerDefault != null) {
             sendConnectivityAction(formerDefault.legacyType, false /* connected */);
         }
         if (defaultNetwork != null) {
             sendConnectivityAction(defaultNetwork.legacyType, true /* connected */);
         }
+    }
 
-        for (NetworkCallback cb : trackingDefault) {
+    private void sendDefaultNetworkCallbacks(TestNetworkAgent formerDefault,
+            TestNetworkAgent defaultNetwork) {
+        for (NetworkCallback cb : trackingDefault.keySet()) {
+            final NetworkRequestInfo nri = trackingDefault.get(cb);
             if (defaultNetwork != null) {
-                cb.onAvailable(defaultNetwork.networkId);
-                cb.onCapabilitiesChanged(
-                        defaultNetwork.networkId, defaultNetwork.networkCapabilities);
-                cb.onLinkPropertiesChanged(
-                        defaultNetwork.networkId, defaultNetwork.linkProperties);
+                nri.handler.post(() -> cb.onAvailable(defaultNetwork.networkId));
+                nri.handler.post(() -> cb.onCapabilitiesChanged(
+                        defaultNetwork.networkId, defaultNetwork.networkCapabilities));
+                nri.handler.post(() -> cb.onLinkPropertiesChanged(
+                        defaultNetwork.networkId, defaultNetwork.linkProperties));
+            } else if (formerDefault != null) {
+                nri.handler.post(() -> cb.onLost(formerDefault.networkId));
             }
         }
     }
 
+    void makeDefaultNetwork(TestNetworkAgent agent) {
+        if (Objects.equals(defaultNetwork, agent)) return;
+
+        final TestNetworkAgent formerDefault = defaultNetwork;
+        defaultNetwork = agent;
+
+        sendDefaultNetworkCallbacks(formerDefault, defaultNetwork);
+        sendDefaultNetworkBroadcasts(formerDefault, defaultNetwork);
+    }
+
     @Override
     public void requestNetwork(NetworkRequest req, NetworkCallback cb, Handler h) {
         assertFalse(allCallbacks.containsKey(cb));
-        allCallbacks.put(cb, h);
+        allCallbacks.put(cb, new NetworkRequestInfo(req, h));
         if (mDefaultRequest.equals(req)) {
-            assertFalse(trackingDefault.contains(cb));
-            trackingDefault.add(cb);
+            assertFalse(trackingDefault.containsKey(cb));
+            trackingDefault.put(cb, new NetworkRequestInfo(req, h));
         } else {
             assertFalse(requested.containsKey(cb));
-            requested.put(cb, req);
+            requested.put(cb, new NetworkRequestInfo(req, h));
         }
     }
 
@@ -164,9 +183,9 @@
     public void requestNetwork(NetworkRequest req,
             int timeoutMs, int legacyType, Handler h, NetworkCallback cb) {
         assertFalse(allCallbacks.containsKey(cb));
-        allCallbacks.put(cb, h);
+        allCallbacks.put(cb, new NetworkRequestInfo(req, h));
         assertFalse(requested.containsKey(cb));
-        requested.put(cb, req);
+        requested.put(cb, new NetworkRequestInfo(req, h));
         assertFalse(legacyTypeMap.containsKey(cb));
         if (legacyType != ConnectivityManager.TYPE_NONE) {
             legacyTypeMap.put(cb, legacyType);
@@ -176,9 +195,9 @@
     @Override
     public void registerNetworkCallback(NetworkRequest req, NetworkCallback cb, Handler h) {
         assertFalse(allCallbacks.containsKey(cb));
-        allCallbacks.put(cb, h);
+        allCallbacks.put(cb, new NetworkRequestInfo(req, h));
         assertFalse(listening.containsKey(cb));
-        listening.put(cb, req);
+        listening.put(cb, new NetworkRequestInfo(req, h));
     }
 
     @Override
@@ -198,7 +217,7 @@
 
     @Override
     public void unregisterNetworkCallback(NetworkCallback cb) {
-        if (trackingDefault.contains(cb)) {
+        if (trackingDefault.containsKey(cb)) {
             trackingDefault.remove(cb);
         } else if (listening.containsKey(cb)) {
             listening.remove(cb);
@@ -211,7 +230,7 @@
         allCallbacks.remove(cb);
 
         assertFalse(allCallbacks.containsKey(cb));
-        assertFalse(trackingDefault.contains(cb));
+        assertFalse(trackingDefault.containsKey(cb));
         assertFalse(listening.containsKey(cb));
         assertFalse(requested.containsKey(cb));
     }
@@ -275,23 +294,25 @@
         }
 
         public void fakeConnect() {
-            for (NetworkRequest request : cm.requested.values()) {
-                if (matchesLegacyType(request.legacyType)) {
+            for (NetworkRequestInfo nri : cm.requested.values()) {
+                if (matchesLegacyType(nri.request.legacyType)) {
                     cm.sendConnectivityAction(legacyType, true /* connected */);
                     // In practice, a given network can match only one legacy type.
                     break;
                 }
             }
             for (NetworkCallback cb : cm.listening.keySet()) {
-                cb.onAvailable(networkId);
-                cb.onCapabilitiesChanged(networkId, copy(networkCapabilities));
-                cb.onLinkPropertiesChanged(networkId, copy(linkProperties));
+                final NetworkRequestInfo nri = cm.listening.get(cb);
+                nri.handler.post(() -> cb.onAvailable(networkId));
+                nri.handler.post(() -> cb.onCapabilitiesChanged(
+                        networkId, copy(networkCapabilities)));
+                nri.handler.post(() -> cb.onLinkPropertiesChanged(networkId, copy(linkProperties)));
             }
         }
 
         public void fakeDisconnect() {
-            for (NetworkRequest request : cm.requested.values()) {
-                if (matchesLegacyType(request.legacyType)) {
+            for (NetworkRequestInfo nri : cm.requested.values()) {
+                if (matchesLegacyType(nri.request.legacyType)) {
                     cm.sendConnectivityAction(legacyType, false /* connected */);
                     break;
                 }
diff --git a/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java b/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java
index d5c7add..a7d61bd 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java
@@ -618,6 +618,7 @@
         when(mOffloadHardwareInterface.getForwardedStats(any())).thenReturn(mForwardedStats);
 
         mServiceContext = new TestContext(mContext);
+        mServiceContext.setUseRegisteredHandlers(true);
         mContentResolver = new MockContentResolver(mServiceContext);
         mContentResolver.addProvider(Settings.AUTHORITY, new FakeSettingsProvider());
         setTetheringSupported(true /* supported */);
@@ -716,6 +717,7 @@
         final Intent intent = new Intent(WifiManager.WIFI_AP_STATE_CHANGED_ACTION);
         intent.putExtra(EXTRA_WIFI_AP_STATE, state);
         mServiceContext.sendStickyBroadcastAsUser(intent, UserHandle.ALL);
+        mLooper.dispatchAll();
     }
 
     private void sendWifiApStateChanged(int state, String ifname, int ipmode) {
@@ -724,6 +726,7 @@
         intent.putExtra(EXTRA_WIFI_AP_INTERFACE_NAME, ifname);
         intent.putExtra(EXTRA_WIFI_AP_MODE, ipmode);
         mServiceContext.sendStickyBroadcastAsUser(intent, UserHandle.ALL);
+        mLooper.dispatchAll();
     }
 
     private static final String[] P2P_RECEIVER_PERMISSIONS_FOR_BROADCAST = {
@@ -750,6 +753,7 @@
 
         mServiceContext.sendBroadcastAsUserMultiplePermissions(intent, UserHandle.ALL,
                 P2P_RECEIVER_PERMISSIONS_FOR_BROADCAST);
+        mLooper.dispatchAll();
     }
 
     private void sendUsbBroadcast(boolean connected, boolean configured, boolean function,
@@ -763,11 +767,13 @@
             intent.putExtra(USB_FUNCTION_NCM, function);
         }
         mServiceContext.sendStickyBroadcastAsUser(intent, UserHandle.ALL);
+        mLooper.dispatchAll();
     }
 
     private void sendConfigurationChanged() {
         final Intent intent = new Intent(Intent.ACTION_CONFIGURATION_CHANGED);
         mServiceContext.sendStickyBroadcastAsUser(intent, UserHandle.ALL);
+        mLooper.dispatchAll();
     }
 
     private void verifyDefaultNetworkRequestFiled() {
@@ -809,7 +815,6 @@
             mTethering.interfaceStatusChanged(TEST_WLAN_IFNAME, true);
         }
         sendWifiApStateChanged(WIFI_AP_STATE_ENABLED);
-        mLooper.dispatchAll();
 
         // If, and only if, Tethering received an interface status changed then
         // it creates a IpServer and sends out a broadcast indicating that the
@@ -857,7 +862,6 @@
 
         // Pretend we then receive USB configured broadcast.
         sendUsbBroadcast(true, true, true, TETHERING_USB);
-        mLooper.dispatchAll();
         // Now we should see the start of tethering mechanics (in this case:
         // tetherMatchingInterfaces() which starts by fetching all interfaces).
         verify(mNetd, times(1)).interfaceGetList();
@@ -886,7 +890,6 @@
             mTethering.interfaceStatusChanged(TEST_WLAN_IFNAME, true);
         }
         sendWifiApStateChanged(WIFI_AP_STATE_ENABLED, TEST_WLAN_IFNAME, IFACE_IP_MODE_LOCAL_ONLY);
-        mLooper.dispatchAll();
 
         verifyInterfaceServingModeStarted(TEST_WLAN_IFNAME);
         verifyTetheringBroadcast(TEST_WLAN_IFNAME, EXTRA_AVAILABLE_TETHER);
@@ -948,7 +951,6 @@
         initTetheringUpstream(upstreamState);
         prepareUsbTethering();
         sendUsbBroadcast(true, true, true, TETHERING_USB);
-        mLooper.dispatchAll();
     }
 
     private void assertSetIfaceToDadProxy(final int numOfCalls, final String ifaceName) {
@@ -1099,7 +1101,6 @@
         // Start USB tethering with no current upstream.
         prepareUsbTethering();
         sendUsbBroadcast(true, true, true, TETHERING_USB);
-        mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).startObserveAllNetworks();
         inOrder.verify(mUpstreamNetworkMonitor).setTryCell(true);
 
@@ -1157,7 +1158,6 @@
     private void runNcmTethering() {
         prepareNcmTethering();
         sendUsbBroadcast(true, true, true, TETHERING_NCM);
-        mLooper.dispatchAll();
     }
 
     @Test
@@ -1205,7 +1205,6 @@
         // tethering mode is to be started.
         mTethering.interfaceStatusChanged(TEST_WLAN_IFNAME, true);
         sendWifiApStateChanged(WIFI_AP_STATE_ENABLED);
-        mLooper.dispatchAll();
 
         // There is 1 IpServer state change event: STATE_AVAILABLE
         verify(mNotificationUpdater, times(1)).onDownstreamChanged(DOWNSTREAM_NONE);
@@ -1233,7 +1232,6 @@
         // tethering mode is to be started.
         mTethering.interfaceStatusChanged(TEST_WLAN_IFNAME, true);
         sendWifiApStateChanged(WIFI_AP_STATE_ENABLED, TEST_WLAN_IFNAME, IFACE_IP_MODE_TETHERED);
-        mLooper.dispatchAll();
 
         verifyInterfaceServingModeStarted(TEST_WLAN_IFNAME);
         verifyTetheringBroadcast(TEST_WLAN_IFNAME, EXTRA_AVAILABLE_TETHER);
@@ -1310,7 +1308,6 @@
         // tethering mode is to be started.
         mTethering.interfaceStatusChanged(TEST_WLAN_IFNAME, true);
         sendWifiApStateChanged(WIFI_AP_STATE_ENABLED, TEST_WLAN_IFNAME, IFACE_IP_MODE_TETHERED);
-        mLooper.dispatchAll();
 
         // We verify get/set called three times here: twice for setup and once during
         // teardown because all events happen over the course of the single
@@ -1634,7 +1631,6 @@
 
         mTethering.startTethering(createTetheringRequestParcel(TETHERING_WIFI), null);
         sendWifiApStateChanged(WIFI_AP_STATE_ENABLED, TEST_WLAN_IFNAME, IFACE_IP_MODE_TETHERED);
-        mLooper.dispatchAll();
         tetherState = callback.pollTetherStatesChanged();
         assertArrayEquals(tetherState.tetheredList, new String[] {TEST_WLAN_IFNAME});
         callback.expectUpstreamChanged(upstreamState.network);
@@ -1656,7 +1652,6 @@
         mLooper.dispatchAll();
         mTethering.stopTethering(TETHERING_WIFI);
         sendWifiApStateChanged(WifiManager.WIFI_AP_STATE_DISABLED);
-        mLooper.dispatchAll();
         tetherState = callback2.pollTetherStatesChanged();
         assertArrayEquals(tetherState.availableList, new String[] {TEST_WLAN_IFNAME});
         mLooper.dispatchAll();
@@ -1749,7 +1744,6 @@
             mTethering.interfaceStatusChanged(TEST_P2P_IFNAME, true);
         }
         sendWifiP2pConnectionChanged(true, true, TEST_P2P_IFNAME);
-        mLooper.dispatchAll();
 
         verifyInterfaceServingModeStarted(TEST_P2P_IFNAME);
         verifyTetheringBroadcast(TEST_P2P_IFNAME, EXTRA_AVAILABLE_TETHER);
@@ -1767,7 +1761,6 @@
         // is being removed.
         sendWifiP2pConnectionChanged(false, true, TEST_P2P_IFNAME);
         mTethering.interfaceRemoved(TEST_P2P_IFNAME);
-        mLooper.dispatchAll();
 
         verify(mNetd, times(1)).tetherApplyDnsInterfaces();
         verify(mNetd, times(1)).tetherInterfaceRemove(TEST_P2P_IFNAME);
@@ -1790,7 +1783,6 @@
             mTethering.interfaceStatusChanged(TEST_P2P_IFNAME, true);
         }
         sendWifiP2pConnectionChanged(true, false, TEST_P2P_IFNAME);
-        mLooper.dispatchAll();
 
         verify(mNetd, never()).interfaceSetCfg(any(InterfaceConfigurationParcel.class));
         verify(mNetd, never()).tetherInterfaceAdd(TEST_P2P_IFNAME);
@@ -1802,7 +1794,6 @@
         // is being removed.
         sendWifiP2pConnectionChanged(false, false, TEST_P2P_IFNAME);
         mTethering.interfaceRemoved(TEST_P2P_IFNAME);
-        mLooper.dispatchAll();
 
         verify(mNetd, never()).tetherApplyDnsInterfaces();
         verify(mNetd, never()).tetherInterfaceRemove(TEST_P2P_IFNAME);
@@ -1838,7 +1829,6 @@
             mTethering.interfaceStatusChanged(TEST_P2P_IFNAME, true);
         }
         sendWifiP2pConnectionChanged(true, true, TEST_P2P_IFNAME);
-        mLooper.dispatchAll();
 
         verify(mNetd, never()).interfaceSetCfg(any(InterfaceConfigurationParcel.class));
         verify(mNetd, never()).tetherInterfaceAdd(TEST_P2P_IFNAME);
@@ -1968,7 +1958,6 @@
         // Expect that when USB comes up, the DHCP server is configured with the requested address.
         mTethering.interfaceStatusChanged(TEST_USB_IFNAME, true);
         sendUsbBroadcast(true, true, true, TETHERING_USB);
-        mLooper.dispatchAll();
         verify(mDhcpServer, timeout(DHCPSERVER_START_TIMEOUT_MS).times(1)).startWithCallbacks(
                 any(), any());
         verify(mNetd).interfaceSetCfg(argThat(cfg -> serverAddr.equals(cfg.ipv4Addr)));
@@ -1988,7 +1977,6 @@
         verify(mUsbManager, times(1)).setCurrentFunctions(UsbManager.FUNCTION_RNDIS);
         mTethering.interfaceStatusChanged(TEST_USB_IFNAME, true);
         sendUsbBroadcast(true, true, true, TETHERING_USB);
-        mLooper.dispatchAll();
         verify(mNetd).interfaceSetCfg(argThat(cfg -> serverAddr.equals(cfg.ipv4Addr)));
         verify(mIpServerDependencies, times(1)).makeDhcpServer(any(), dhcpParamsCaptor.capture(),
                 any());
@@ -2212,7 +2200,6 @@
 
         mTethering.interfaceStatusChanged(TEST_USB_IFNAME, true);
         sendUsbBroadcast(true, true, true, TETHERING_USB);
-        mLooper.dispatchAll();
         assertContains(Arrays.asList(mTethering.getTetherableIfaces()), TEST_USB_IFNAME);
         assertContains(Arrays.asList(mTethering.getTetherableIfaces()), TEST_ETH_IFNAME);
         assertEquals(TETHER_ERROR_IFACE_CFG_ERROR, mTethering.getLastTetherError(TEST_USB_IFNAME));
@@ -2251,7 +2238,6 @@
         // Run local only tethering.
         mTethering.interfaceStatusChanged(TEST_P2P_IFNAME, true);
         sendWifiP2pConnectionChanged(true, true, TEST_P2P_IFNAME);
-        mLooper.dispatchAll();
         verify(mDhcpServer, timeout(DHCPSERVER_START_TIMEOUT_MS)).startWithCallbacks(
                 any(), dhcpEventCbsCaptor.capture());
         eventCallbacks = dhcpEventCbsCaptor.getValue();
@@ -2268,7 +2254,6 @@
         // Run wifi tethering.
         mTethering.interfaceStatusChanged(TEST_WLAN_IFNAME, true);
         sendWifiApStateChanged(WIFI_AP_STATE_ENABLED, TEST_WLAN_IFNAME, IFACE_IP_MODE_TETHERED);
-        mLooper.dispatchAll();
         verify(mDhcpServer, timeout(DHCPSERVER_START_TIMEOUT_MS)).startWithCallbacks(
                 any(), dhcpEventCbsCaptor.capture());
         eventCallbacks = dhcpEventCbsCaptor.getValue();
diff --git a/Tethering/tests/unit/src/com/android/networkstack/tethering/UpstreamNetworkMonitorTest.java b/Tethering/tests/unit/src/com/android/networkstack/tethering/UpstreamNetworkMonitorTest.java
index 43d32b5..40e4f83 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/UpstreamNetworkMonitorTest.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/UpstreamNetworkMonitorTest.java
@@ -50,13 +50,16 @@
 import android.net.NetworkRequest;
 import android.net.util.SharedLog;
 import android.os.Handler;
+import android.os.Looper;
 import android.os.Message;
+import android.os.test.TestLooper;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.util.State;
 import com.android.internal.util.StateMachine;
+import com.android.networkstack.tethering.TestConnectivityManager.NetworkRequestInfo;
 import com.android.networkstack.tethering.TestConnectivityManager.TestNetworkAgent;
 
 import org.junit.After;
@@ -100,6 +103,8 @@
     private TestConnectivityManager mCM;
     private UpstreamNetworkMonitor mUNM;
 
+    private final TestLooper mLooper = new TestLooper();
+
     @Before public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
         reset(mContext);
@@ -109,7 +114,7 @@
         when(mEntitleMgr.isCellularUpstreamPermitted()).thenReturn(true);
 
         mCM = spy(new TestConnectivityManager(mContext, mCS, sDefaultRequest));
-        mSM = new TestStateMachine();
+        mSM = new TestStateMachine(mLooper.getLooper());
         mUNM = new UpstreamNetworkMonitor(mCM, mSM, mLog, EVENT_UNM_UPDATE);
     }
 
@@ -295,14 +300,17 @@
 
         final TestNetworkAgent wifiAgent = new TestNetworkAgent(mCM, WIFI_CAPABILITIES);
         wifiAgent.fakeConnect();
+        mLooper.dispatchAll();
         // WiFi is up, we should prefer it.
         assertSatisfiesLegacyType(TYPE_WIFI, mUNM.selectPreferredUpstreamType(preferredTypes));
         wifiAgent.fakeDisconnect();
+        mLooper.dispatchAll();
         // There are no networks, so there is nothing to select.
         assertSatisfiesLegacyType(TYPE_NONE, mUNM.selectPreferredUpstreamType(preferredTypes));
 
         final TestNetworkAgent cellAgent = new TestNetworkAgent(mCM, CELL_CAPABILITIES);
         cellAgent.fakeConnect();
+        mLooper.dispatchAll();
         assertSatisfiesLegacyType(TYPE_NONE, mUNM.selectPreferredUpstreamType(preferredTypes));
 
         preferredTypes.add(TYPE_MOBILE_DUN);
@@ -319,7 +327,7 @@
                 mUNM.selectPreferredUpstreamType(preferredTypes));
         // Check to see we filed an explicit request.
         assertEquals(1, mCM.requested.size());
-        NetworkRequest netReq = (NetworkRequest) mCM.requested.values().toArray()[0];
+        NetworkRequest netReq = ((NetworkRequestInfo) mCM.requested.values().toArray()[0]).request;
         assertTrue(netReq.networkCapabilities.hasTransport(TRANSPORT_CELLULAR));
         assertFalse(netReq.networkCapabilities.hasCapability(NET_CAPABILITY_DUN));
         // mobile is not permitted, we should not use HIPRI.
@@ -332,6 +340,7 @@
                 mUNM.selectPreferredUpstreamType(preferredTypes));
 
         wifiAgent.fakeConnect();
+        mLooper.dispatchAll();
         // WiFi is up, and we should prefer it over cell.
         assertSatisfiesLegacyType(TYPE_WIFI, mUNM.selectPreferredUpstreamType(preferredTypes));
         assertEquals(0, mCM.requested.size());
@@ -344,17 +353,19 @@
 
         final TestNetworkAgent dunAgent = new TestNetworkAgent(mCM, DUN_CAPABILITIES);
         dunAgent.fakeConnect();
+        mLooper.dispatchAll();
 
         // WiFi is still preferred.
         assertSatisfiesLegacyType(TYPE_WIFI, mUNM.selectPreferredUpstreamType(preferredTypes));
 
         // WiFi goes down, cell and DUN are still up but only DUN is preferred.
         wifiAgent.fakeDisconnect();
+        mLooper.dispatchAll();
         assertSatisfiesLegacyType(TYPE_MOBILE_DUN,
                 mUNM.selectPreferredUpstreamType(preferredTypes));
         // Check to see we filed an explicit request.
         assertEquals(1, mCM.requested.size());
-        netReq = (NetworkRequest) mCM.requested.values().toArray()[0];
+        netReq = ((NetworkRequestInfo) mCM.requested.values().toArray()[0]).request;
         assertTrue(netReq.networkCapabilities.hasTransport(TRANSPORT_CELLULAR));
         assertTrue(netReq.networkCapabilities.hasCapability(NET_CAPABILITY_DUN));
         // mobile is not permitted, we should not use DUN.
@@ -378,6 +389,7 @@
         final TestNetworkAgent cellAgent = new TestNetworkAgent(mCM, CELL_CAPABILITIES);
         cellAgent.fakeConnect();
         mCM.makeDefaultNetwork(cellAgent);
+        mLooper.dispatchAll();
         assertEquals(cellAgent.networkId, mUNM.getCurrentPreferredUpstream().network);
         assertEquals(0, mCM.requested.size());
 
@@ -390,11 +402,13 @@
         // [2] WiFi connects but not validated/promoted to default -> mobile selected.
         final TestNetworkAgent wifiAgent = new TestNetworkAgent(mCM, WIFI_CAPABILITIES);
         wifiAgent.fakeConnect();
+        mLooper.dispatchAll();
         assertEquals(cellAgent.networkId, mUNM.getCurrentPreferredUpstream().network);
         assertEquals(0, mCM.requested.size());
 
         // [3] WiFi validates and is promoted to the default network -> WiFi selected.
         mCM.makeDefaultNetwork(wifiAgent);
+        mLooper.dispatchAll();
         assertEquals(wifiAgent.networkId, mUNM.getCurrentPreferredUpstream().network);
         assertEquals(0, mCM.requested.size());
 
@@ -406,6 +420,7 @@
 
         // [5] WiFi no longer validated, mobile becomes default, DUN required -> null selected.
         mCM.makeDefaultNetwork(cellAgent);
+        mLooper.dispatchAll();
         assertEquals(null, mUNM.getCurrentPreferredUpstream());
         assertEquals(1, mCM.requested.size());
         assertTrue(isDunRequested());
@@ -415,6 +430,7 @@
         dunAgent.networkCapabilities.addCapability(NET_CAPABILITY_DUN);
         dunAgent.networkCapabilities.removeCapability(NET_CAPABILITY_INTERNET);
         dunAgent.fakeConnect();
+        mLooper.dispatchAll();
         assertEquals(dunAgent.networkId, mUNM.getCurrentPreferredUpstream().network);
         assertEquals(1, mCM.requested.size());
 
@@ -461,6 +477,7 @@
         }
         wifiAgent.fakeConnect();
         wifiAgent.sendLinkProperties();
+        mLooper.dispatchAll();
 
         local = mUNM.getLocalPrefixes();
         assertPrefixSet(local, INCLUDES, alreadySeen);
@@ -485,6 +502,7 @@
         }
         cellAgent.fakeConnect();
         cellAgent.sendLinkProperties();
+        mLooper.dispatchAll();
 
         local = mUNM.getLocalPrefixes();
         assertPrefixSet(local, INCLUDES, alreadySeen);
@@ -506,6 +524,7 @@
         }
         dunAgent.fakeConnect();
         dunAgent.sendLinkProperties();
+        mLooper.dispatchAll();
 
         local = mUNM.getLocalPrefixes();
         assertPrefixSet(local, INCLUDES, alreadySeen);
@@ -517,6 +536,7 @@
         // [4] Pretend Wi-Fi disconnected.  It's addresses/prefixes should no
         // longer be included (should be properly removed).
         wifiAgent.fakeDisconnect();
+        mLooper.dispatchAll();
         local = mUNM.getLocalPrefixes();
         assertPrefixSet(local, EXCLUDES, wifiLinkPrefixes);
         assertPrefixSet(local, INCLUDES, cellLinkPrefixes);
@@ -524,6 +544,7 @@
 
         // [5] Pretend mobile disconnected.
         cellAgent.fakeDisconnect();
+        mLooper.dispatchAll();
         local = mUNM.getLocalPrefixes();
         assertPrefixSet(local, EXCLUDES, wifiLinkPrefixes);
         assertPrefixSet(local, EXCLUDES, cellLinkPrefixes);
@@ -531,6 +552,7 @@
 
         // [6] Pretend DUN disconnected.
         dunAgent.fakeDisconnect();
+        mLooper.dispatchAll();
         local = mUNM.getLocalPrefixes();
         assertTrue(local.isEmpty());
     }
@@ -550,6 +572,7 @@
         // Setup mobile network.
         final TestNetworkAgent cellAgent = new TestNetworkAgent(mCM, CELL_CAPABILITIES);
         cellAgent.fakeConnect();
+        mLooper.dispatchAll();
 
         assertSatisfiesLegacyType(TYPE_MOBILE_HIPRI,
                 mUNM.selectPreferredUpstreamType(preferredTypes));
@@ -575,8 +598,8 @@
     }
 
     private boolean isDunRequested() {
-        for (NetworkRequest req : mCM.requested.values()) {
-            if (req.networkCapabilities.hasCapability(NET_CAPABILITY_DUN)) {
+        for (NetworkRequestInfo nri : mCM.requested.values()) {
+            if (nri.request.networkCapabilities.hasCapability(NET_CAPABILITY_DUN)) {
                 return true;
             }
         }
@@ -602,8 +625,8 @@
             }
         }
 
-        public TestStateMachine() {
-            super("UpstreamNetworkMonitor.TestStateMachine");
+        public TestStateMachine(Looper looper) {
+            super("UpstreamNetworkMonitor.TestStateMachine", looper);
             addState(mLoggingState);
             setInitialState(mLoggingState);
             super.start();