diff --git a/Tethering/bpf_progs/Android.bp b/Tethering/bpf_progs/Android.bp
index 2b10f89..289d75d 100644
--- a/Tethering/bpf_progs/Android.bp
+++ b/Tethering/bpf_progs/Android.bp
@@ -51,8 +51,6 @@
     include_dirs: [
         // TODO: get rid of system/netd.
         "system/netd/bpf_progs",             // for bpf_net_helpers.h
-        "system/netd/libnetdbpf/include",    // for bpf_shared.h
-        "system/netd/libnetdutils/include",  // for UidConstants.h
     ],
 }
 
@@ -66,7 +64,5 @@
     include_dirs: [
         // TODO: get rid of system/netd.
         "system/netd/bpf_progs",             // for bpf_net_helpers.h
-        "system/netd/libnetdbpf/include",    // for bpf_shared.h
-        "system/netd/libnetdutils/include",  // for UidConstants.h
     ],
 }
diff --git a/Tethering/src/android/net/ip/RouterAdvertisementDaemon.java b/Tethering/src/android/net/ip/RouterAdvertisementDaemon.java
index 7c0b7cc..afccd0a 100644
--- a/Tethering/src/android/net/ip/RouterAdvertisementDaemon.java
+++ b/Tethering/src/android/net/ip/RouterAdvertisementDaemon.java
@@ -37,7 +37,7 @@
 import android.util.Log;
 
 import com.android.internal.annotations.GuardedBy;
-import com.android.internal.util.TrafficStatsConstants;
+import com.android.net.module.util.NetworkStackConstants;
 
 import java.io.FileDescriptor;
 import java.io.IOException;
@@ -589,7 +589,7 @@
         final int send_timout_ms = 300;
 
         final int oldTag = TrafficStats.getAndSetThreadStatsTag(
-                TrafficStatsConstants.TAG_SYSTEM_NEIGHBOR);
+                NetworkStackConstants.TAG_SYSTEM_NEIGHBOR);
         try {
             mSocket = Os.socket(AF_INET6, SOCK_RAW, IPPROTO_ICMPV6);
             // Setting SNDTIMEO is purely for defensive purposes.
diff --git a/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java b/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java
index 74eb87b..7fb73b4 100644
--- a/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java
+++ b/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java
@@ -814,7 +814,7 @@
                 final int upstreamIfindex = rule.upstreamIfindex;
                 pw.println(String.format("%d(%s) %d(%s) %s %s %s", upstreamIfindex,
                         mInterfaceNames.get(upstreamIfindex), rule.downstreamIfindex,
-                        downstreamIface, rule.address, rule.srcMac, rule.dstMac));
+                        downstreamIface, rule.address.getHostAddress(), rule.srcMac, rule.dstMac));
             }
             pw.decreaseIndent();
         }
@@ -851,9 +851,10 @@
         } catch (UnknownHostException impossible) {
             throw new AssertionError("4-byte array not valid IPv4 address!");
         }
-        return String.format("%d(%s) %d(%s) %s:%d -> %s:%d -> %s:%d",
-                key.iif, getIfName(key.iif), value.oif, getIfName(value.oif),
-                private4, key.srcPort, public4, value.srcPort, dst4, key.dstPort);
+        return String.format("[%s] %d(%s) %s:%d -> %d(%s) %s:%d -> %s:%d",
+                key.dstMac, key.iif, getIfName(key.iif), private4, key.srcPort,
+                value.oif, getIfName(value.oif),
+                public4, value.srcPort, dst4, key.dstPort);
     }
 
     private void dumpIpv4ForwardingRules(IndentingPrintWriter pw) {
@@ -866,7 +867,7 @@
                 pw.println("No IPv4 rules");
                 return;
             }
-            pw.println("[IPv4]: iif(iface) oif(iface) src nat dst");
+            pw.println("IPv4: [inDstMac] iif(iface) src -> nat -> dst");
             pw.increaseIndent();
             map.forEach((k, v) -> pw.println(ipv4RuleToString(k, v)));
         } catch (ErrnoException e) {
diff --git a/Tethering/src/com/android/networkstack/tethering/Tethering.java b/Tethering/src/com/android/networkstack/tethering/Tethering.java
index ac5857d..f795747 100644
--- a/Tethering/src/com/android/networkstack/tethering/Tethering.java
+++ b/Tethering/src/com/android/networkstack/tethering/Tethering.java
@@ -442,7 +442,8 @@
     // NOTE: This is always invoked on the mLooper thread.
     private void updateConfiguration() {
         mConfig = mDeps.generateTetheringConfiguration(mContext, mLog, mActiveDataSubId);
-        mUpstreamNetworkMonitor.updateMobileRequiresDun(mConfig.isDunRequired);
+        mUpstreamNetworkMonitor.setUpstreamConfig(mConfig.chooseUpstreamAutomatically,
+                mConfig.isDunRequired);
         reportConfigurationChanged(mConfig.toStableParcelable());
     }
 
@@ -1559,7 +1560,7 @@
                             config.preferredUpstreamIfaceTypes);
             if (ns == null) {
                 if (tryCell) {
-                    mUpstreamNetworkMonitor.registerMobileNetworkRequest();
+                    mUpstreamNetworkMonitor.setTryCell(true);
                     // We think mobile should be coming up; don't set a retry.
                 } else {
                     sendMessageDelayed(CMD_RETRY_UPSTREAM, UPSTREAM_SETTLE_TIME_MS);
@@ -1718,6 +1719,12 @@
                     break;
             }
 
+            if (mConfig.chooseUpstreamAutomatically
+                    && arg1 == UpstreamNetworkMonitor.EVENT_DEFAULT_SWITCHED) {
+                chooseUpstreamType(true);
+                return;
+            }
+
             if (ns == null || !pertainsToCurrentUpstream(ns)) {
                 // TODO: In future, this is where upstream evaluation and selection
                 // could be handled for notifications which include sufficient data.
@@ -1852,7 +1859,7 @@
                         // longer desired, release any mobile requests.
                         final boolean previousUpstreamWanted = updateUpstreamWanted();
                         if (previousUpstreamWanted && !mUpstreamWanted) {
-                            mUpstreamNetworkMonitor.releaseMobileNetworkRequest();
+                            mUpstreamNetworkMonitor.setTryCell(false);
                         }
                         break;
                     }
diff --git a/Tethering/src/com/android/networkstack/tethering/UpstreamNetworkMonitor.java b/Tethering/src/com/android/networkstack/tethering/UpstreamNetworkMonitor.java
index b17065c..e39145b 100644
--- a/Tethering/src/com/android/networkstack/tethering/UpstreamNetworkMonitor.java
+++ b/Tethering/src/com/android/networkstack/tethering/UpstreamNetworkMonitor.java
@@ -42,11 +42,14 @@
 import android.util.Log;
 import android.util.SparseIntArray;
 
+import androidx.annotation.NonNull;
+
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.StateMachine;
 
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Objects;
 import java.util.Set;
 
 
@@ -60,7 +63,7 @@
  * Calling #startObserveAllNetworks() to observe all networks. Listening all
  * networks is necessary while the expression of preferred upstreams remains
  * a list of legacy connectivity types.  In future, this can be revisited.
- * Calling #registerMobileNetworkRequest() to bring up mobile DUN/HIPRI network.
+ * Calling #setTryCell() to request bringing up mobile DUN or HIPRI.
  *
  * The methods and data members of this class are only to be accessed and
  * modified from the tethering main state machine thread. Any other
@@ -82,6 +85,7 @@
     public static final int EVENT_ON_CAPABILITIES   = 1;
     public static final int EVENT_ON_LINKPROPERTIES = 2;
     public static final int EVENT_ON_LOST           = 3;
+    public static final int EVENT_DEFAULT_SWITCHED  = 4;
     public static final int NOTIFY_LOCAL_PREFIXES   = 10;
     // This value is used by deprecated preferredUpstreamIfaceTypes selection which is default
     // disabled.
@@ -114,7 +118,14 @@
     private NetworkCallback mListenAllCallback;
     private NetworkCallback mDefaultNetworkCallback;
     private NetworkCallback mMobileNetworkCallback;
+
+    /** Whether Tethering has requested a cellular upstream. */
+    private boolean mTryCell;
+    /** Whether the carrier requires DUN. */
     private boolean mDunRequired;
+    /** Whether automatic upstream selection is enabled. */
+    private boolean mAutoUpstream;
+
     // Whether the current default upstream is mobile or not.
     private boolean mIsDefaultCellularUpstream;
     // The current system default network (not really used yet).
@@ -190,23 +201,49 @@
         mNetworkMap.clear();
     }
 
-    /** Setup or teardown DUN connection according to |dunRequired|. */
-    public void updateMobileRequiresDun(boolean dunRequired) {
-        final boolean valueChanged = (mDunRequired != dunRequired);
+    private void reevaluateUpstreamRequirements(boolean tryCell, boolean autoUpstream,
+            boolean dunRequired) {
+        final boolean mobileRequestRequired = tryCell && (dunRequired || !autoUpstream);
+        final boolean dunRequiredChanged = (mDunRequired != dunRequired);
+
+        mTryCell = tryCell;
         mDunRequired = dunRequired;
-        if (valueChanged && mobileNetworkRequested()) {
-            releaseMobileNetworkRequest();
+        mAutoUpstream = autoUpstream;
+
+        if (mobileRequestRequired && !mobileNetworkRequested()) {
             registerMobileNetworkRequest();
+        } else if (mobileNetworkRequested() && !mobileRequestRequired) {
+            releaseMobileNetworkRequest();
+        } else if (mobileNetworkRequested() && dunRequiredChanged) {
+            releaseMobileNetworkRequest();
+            if (mobileRequestRequired) {
+                registerMobileNetworkRequest();
+            }
         }
     }
 
+    /**
+     * Informs UpstreamNetworkMonitor that a cellular upstream is desired.
+     *
+     * This may result in filing a NetworkRequest for DUN if it is required, or for MOBILE_HIPRI if
+     * automatic upstream selection is disabled and MOBILE_HIPRI is the preferred upstream.
+     */
+    public void setTryCell(boolean tryCell) {
+        reevaluateUpstreamRequirements(tryCell, mAutoUpstream, mDunRequired);
+    }
+
+    /** Informs UpstreamNetworkMonitor of upstream configuration parameters. */
+    public void setUpstreamConfig(boolean autoUpstream, boolean dunRequired) {
+        reevaluateUpstreamRequirements(mTryCell, autoUpstream, dunRequired);
+    }
+
     /** Whether mobile network is requested. */
     public boolean mobileNetworkRequested() {
         return (mMobileNetworkCallback != null);
     }
 
     /** Request mobile network if mobile upstream is permitted. */
-    public void registerMobileNetworkRequest() {
+    private void registerMobileNetworkRequest() {
         if (!isCellularUpstreamPermitted()) {
             mLog.i("registerMobileNetworkRequest() is not permitted");
             releaseMobileNetworkRequest();
@@ -241,14 +278,16 @@
         // TODO: Change the timeout from 0 (no onUnavailable callback) to some
         // moderate callback timeout. This might be useful for updating some UI.
         // Additionally, we log a message to aid in any subsequent debugging.
-        mLog.i("requesting mobile upstream network: " + mobileUpstreamRequest);
+        mLog.i("requesting mobile upstream network: " + mobileUpstreamRequest
+                + " mTryCell=" + mTryCell + " mAutoUpstream=" + mAutoUpstream
+                + " mDunRequired=" + mDunRequired);
 
         cm().requestNetwork(mobileUpstreamRequest, 0, legacyType, mHandler,
                 mMobileNetworkCallback);
     }
 
     /** Release mobile network request. */
-    public void releaseMobileNetworkRequest() {
+    private void releaseMobileNetworkRequest() {
         if (mMobileNetworkCallback == null) return;
 
         cm().unregisterNetworkCallback(mMobileNetworkCallback);
@@ -363,7 +402,7 @@
         notifyTarget(EVENT_ON_CAPABILITIES, network);
     }
 
-    private void handleLinkProp(Network network, LinkProperties newLp) {
+    private void updateLinkProperties(Network network, LinkProperties newLp) {
         final UpstreamNetworkState prev = mNetworkMap.get(network);
         if (prev == null || newLp.equals(prev.linkProperties)) {
             // Ignore notifications about networks for which we have not yet
@@ -379,8 +418,10 @@
 
         mNetworkMap.put(network, new UpstreamNetworkState(
                 newLp, prev.networkCapabilities, network));
-        // TODO: If sufficient information is available to select a more
-        // preferable upstream, do so now and notify the target.
+    }
+
+    private void handleLinkProp(Network network, LinkProperties newLp) {
+        updateLinkProperties(network, newLp);
         notifyTarget(EVENT_ON_LINKPROPERTIES, network);
     }
 
@@ -410,6 +451,24 @@
         notifyTarget(EVENT_ON_LOST, mNetworkMap.remove(network));
     }
 
+    private void maybeHandleNetworkSwitch(@NonNull Network network) {
+        if (Objects.equals(mDefaultInternetNetwork, network)) return;
+
+        final UpstreamNetworkState ns = mNetworkMap.get(network);
+        if (ns == null) {
+            // Can never happen unless there is a bug in ConnectivityService. Entries are only
+            // removed from mNetworkMap when receiving onLost, and onLost for a given network can
+            // never be followed by any other callback on that network.
+            Log.wtf(TAG, "maybeHandleNetworkSwitch: no UpstreamNetworkState for " + network);
+            return;
+        }
+
+        // Default network changed. Update local data and notify tethering.
+        Log.d(TAG, "New default Internet network: " + network);
+        mDefaultInternetNetwork = network;
+        notifyTarget(EVENT_DEFAULT_SWITCHED, ns);
+    }
+
     private void recomputeLocalPrefixes() {
         final HashSet<IpPrefix> localPrefixes = allLocalPrefixes(mNetworkMap.values());
         if (!mLocalPrefixes.equals(localPrefixes)) {
@@ -447,7 +506,22 @@
         @Override
         public void onCapabilitiesChanged(Network network, NetworkCapabilities newNc) {
             if (mCallbackType == CALLBACK_DEFAULT_INTERNET) {
-                mDefaultInternetNetwork = network;
+                // mDefaultInternetNetwork is not updated here because upstream selection must only
+                // run when the LinkProperties have been updated as well as the capabilities. If
+                // this callback is due to a default network switch, then the system will invoke
+                // onLinkPropertiesChanged right after this method and mDefaultInternetNetwork will
+                // be updated then.
+                //
+                // Technically, not updating here isn't necessary, because the notifications to
+                // Tethering sent by notifyTarget are messages sent to a state machine running on
+                // the same thread as this method, and so cannot arrive until after this method has
+                // returned. However, it is not a good idea to rely on that because fact that
+                // Tethering uses multiple state machines running on the same thread is a major
+                // source of race conditions and something that should be fixed.
+                //
+                // TODO: is it correct that this code always updates EntitlementManager?
+                // This code runs when the default network connects or changes capabilities, but the
+                // default network might not be the tethering upstream.
                 final boolean newIsCellular = isCellular(newNc);
                 if (mIsDefaultCellularUpstream != newIsCellular) {
                     mIsDefaultCellularUpstream = newIsCellular;
@@ -461,7 +535,15 @@
 
         @Override
         public void onLinkPropertiesChanged(Network network, LinkProperties newLp) {
-            if (mCallbackType == CALLBACK_DEFAULT_INTERNET) return;
+            if (mCallbackType == CALLBACK_DEFAULT_INTERNET) {
+                updateLinkProperties(network, newLp);
+                // When the default network callback calls onLinkPropertiesChanged, it means that
+                // all the network information for the default network is known (because
+                // onLinkPropertiesChanged is called after onAvailable and onCapabilitiesChanged).
+                // Inform tethering that the default network might have changed.
+                maybeHandleNetworkSwitch(network);
+                return;
+            }
 
             handleLinkProp(network, newLp);
             // Any non-LISTEN_ALL callback will necessarily concern a network that will
@@ -478,6 +560,8 @@
                 mDefaultInternetNetwork = null;
                 mIsDefaultCellularUpstream = false;
                 mEntitlementMgr.notifyUpstream(false);
+                Log.d(TAG, "Lost default Internet network: " + network);
+                notifyTarget(EVENT_DEFAULT_SWITCHED, null);
                 return;
             }
 
diff --git a/Tethering/tests/privileged/src/android/net/ip/DadProxyTest.java b/Tethering/tests/privileged/src/android/net/ip/DadProxyTest.java
index 42a91aa..a933e1b 100644
--- a/Tethering/tests/privileged/src/android/net/ip/DadProxyTest.java
+++ b/Tethering/tests/privileged/src/android/net/ip/DadProxyTest.java
@@ -21,9 +21,8 @@
 import static com.android.net.module.util.IpUtils.icmpv6Checksum;
 import static com.android.net.module.util.NetworkStackConstants.ETHER_SRC_ADDR_OFFSET;
 
-import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertTrue;
 
 import android.app.Instrumentation;
 import android.content.Context;
@@ -52,13 +51,14 @@
 import org.junit.runner.RunWith;
 import org.mockito.MockitoAnnotations;
 
+import java.io.IOException;
 import java.nio.ByteBuffer;
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
 public class DadProxyTest {
     private static final int DATA_BUFFER_LEN = 4096;
-    private static final int PACKET_TIMEOUT_MS = 5_000;
+    private static final int PACKET_TIMEOUT_MS = 2_000;  // Long enough for DAD to succeed.
 
     // Start the readers manually on a common handler shared with DadProxy, for simplicity
     @Rule
@@ -119,16 +119,18 @@
         }
     }
 
-    private void setupTapInterfaces() {
+    private void setupTapInterfaces() throws Exception {
         // Create upstream test iface.
         mUpstreamReader.start(mHandler);
-        mUpstreamParams = InterfaceParams.getByName(mUpstreamReader.iface.getInterfaceName());
+        final String upstreamIface = mUpstreamReader.iface.getInterfaceName();
+        mUpstreamParams = InterfaceParams.getByName(upstreamIface);
         assertNotNull(mUpstreamParams);
         mUpstreamPacketReader = mUpstreamReader.getReader();
 
         // Create tethered test iface.
         mTetheredReader.start(mHandler);
-        mTetheredParams = InterfaceParams.getByName(mTetheredReader.getIface().getInterfaceName());
+        final String tetheredIface = mTetheredReader.getIface().getInterfaceName();
+        mTetheredParams = InterfaceParams.getByName(tetheredIface);
         assertNotNull(mTetheredParams);
         mTetheredPacketReader = mTetheredReader.getReader();
     }
@@ -224,6 +226,12 @@
         return false;
     }
 
+    private ByteBuffer copy(ByteBuffer buf) {
+        // There does not seem to be a way to copy ByteBuffers. ByteBuffer does not implement
+        // clone() and duplicate() copies the metadata but shares the contents.
+        return ByteBuffer.wrap(buf.array().clone());
+    }
+
     private void updateDstMac(ByteBuffer buf, MacAddress mac) {
         buf.put(mac.toByteArray());
         buf.rewind();
@@ -234,14 +242,50 @@
         buf.rewind();
     }
 
+    private void receivePacketAndMaybeExpectForwarded(boolean expectForwarded,
+            ByteBuffer in, TapPacketReader inReader, ByteBuffer out, TapPacketReader outReader)
+            throws IOException {
+
+        inReader.sendResponse(in);
+        if (waitForPacket(out, outReader)) return;
+
+        // When the test runs, DAD may be in progress, because the interface has just been created.
+        // If so, the DAD proxy will get EADDRNOTAVAIL when trying to send packets. It is not
+        // possible to work around this using IPV6_FREEBIND or IPV6_TRANSPARENT options because the
+        // kernel rawv6 code doesn't consider those options either when binding or when sending, and
+        // doesn't get the source address from the packet even in IPPROTO_RAW/HDRINCL mode (it only
+        // gets it from the socket or from cmsg).
+        //
+        // If DAD was in progress when the above was attempted, try again and expect the packet to
+        // be forwarded. Don't disable DAD in the test because if we did, the test would not notice
+        // if, for example, the DAD proxy code just crashed if it received EADDRNOTAVAIL.
+        final String msg = expectForwarded
+                ? "Did not receive expected packet even after waiting for DAD:"
+                : "Unexpectedly received packet:";
+
+        inReader.sendResponse(in);
+        assertEquals(msg, expectForwarded, waitForPacket(out, outReader));
+    }
+
+    private void receivePacketAndExpectForwarded(ByteBuffer in, TapPacketReader inReader,
+            ByteBuffer out, TapPacketReader outReader) throws IOException {
+        receivePacketAndMaybeExpectForwarded(true, in, inReader, out, outReader);
+    }
+
+    private void receivePacketAndExpectNotForwarded(ByteBuffer in, TapPacketReader inReader,
+            ByteBuffer out, TapPacketReader outReader) throws IOException {
+        receivePacketAndMaybeExpectForwarded(false, in, inReader, out, outReader);
+    }
+
     @Test
     public void testNaForwardingFromUpstreamToTether() throws Exception {
         ByteBuffer na = createDadPacket(NeighborPacketForwarder.ICMPV6_NEIGHBOR_ADVERTISEMENT);
 
-        mUpstreamPacketReader.sendResponse(na);
-        updateDstMac(na, MacAddress.fromString("33:33:00:00:00:01"));
-        updateSrcMac(na, mTetheredParams);
-        assertTrue(waitForPacket(na, mTetheredPacketReader));
+        ByteBuffer out = copy(na);
+        updateDstMac(out, MacAddress.fromString("33:33:00:00:00:01"));
+        updateSrcMac(out, mTetheredParams);
+
+        receivePacketAndExpectForwarded(na, mUpstreamPacketReader, out, mTetheredPacketReader);
     }
 
     @Test
@@ -249,19 +293,21 @@
     public void testNaForwardingFromTetherToUpstream() throws Exception {
         ByteBuffer na = createDadPacket(NeighborPacketForwarder.ICMPV6_NEIGHBOR_ADVERTISEMENT);
 
-        mTetheredPacketReader.sendResponse(na);
-        updateDstMac(na, MacAddress.fromString("33:33:00:00:00:01"));
-        updateSrcMac(na, mTetheredParams);
-        assertFalse(waitForPacket(na, mUpstreamPacketReader));
+        ByteBuffer out = copy(na);
+        updateDstMac(out, MacAddress.fromString("33:33:00:00:00:01"));
+        updateSrcMac(out, mTetheredParams);
+
+        receivePacketAndExpectNotForwarded(na, mTetheredPacketReader, out, mUpstreamPacketReader);
     }
 
     @Test
     public void testNsForwardingFromTetherToUpstream() throws Exception {
         ByteBuffer ns = createDadPacket(NeighborPacketForwarder.ICMPV6_NEIGHBOR_SOLICITATION);
 
-        mTetheredPacketReader.sendResponse(ns);
-        updateSrcMac(ns, mUpstreamParams);
-        assertTrue(waitForPacket(ns, mUpstreamPacketReader));
+        ByteBuffer out = copy(ns);
+        updateSrcMac(out, mUpstreamParams);
+
+        receivePacketAndExpectForwarded(ns, mTetheredPacketReader, out, mUpstreamPacketReader);
     }
 
     @Test
@@ -269,8 +315,9 @@
     public void testNsForwardingFromUpstreamToTether() throws Exception {
         ByteBuffer ns = createDadPacket(NeighborPacketForwarder.ICMPV6_NEIGHBOR_SOLICITATION);
 
-        mUpstreamPacketReader.sendResponse(ns);
+        ByteBuffer out = copy(ns);
         updateSrcMac(ns, mUpstreamParams);
-        assertFalse(waitForPacket(ns, mTetheredPacketReader));
+
+        receivePacketAndExpectNotForwarded(ns, mUpstreamPacketReader, out, mTetheredPacketReader);
     }
 }
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..d045bf1 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,12 @@
 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 androidx.annotation.Nullable;
+
 import java.util.Map;
 import java.util.Objects;
-import java.util.Set;
 
 /**
  * Simulates upstream switching and sending NetworkCallbacks and CONNECTIVITY_ACTION broadcasts.
@@ -60,17 +60,20 @@
  *   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 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 static final boolean BROADCAST_FIRST = false;
+    public static final boolean CALLBACKS_FIRST = true;
+
+    final Map<NetworkCallback, NetworkRequestInfo> mAllCallbacks = new ArrayMap<>();
+    final Map<NetworkCallback, NetworkRequestInfo> mTrackingDefault = new ArrayMap<>();
+    final Map<NetworkCallback, NetworkRequestInfo> mListening = new ArrayMap<>();
+    final Map<NetworkCallback, NetworkRequestInfo> mRequested = new ArrayMap<>();
+    final Map<NetworkCallback, Integer> mLegacyTypeMap = new ArrayMap<>();
 
     private final NetworkRequest mDefaultRequest;
     private final Context mContext;
 
     private int mNetworkId = 100;
+    private TestNetworkAgent mDefaultNetwork = null;
 
     /**
      * Constructs a TestConnectivityManager.
@@ -86,28 +89,37 @@
         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()
-                && listening.isEmpty()
-                && requested.isEmpty()
-                && legacyTypeMap.isEmpty();
+        return mAllCallbacks.isEmpty()
+                && mTrackingDefault.isEmpty()
+                && mListening.isEmpty()
+                && mRequested.isEmpty()
+                && mLegacyTypeMap.isEmpty();
     }
 
     boolean onlyHasDefaultCallbacks() {
-        return (allCallbacks.size() == 1)
-                && (trackingDefault.size() == 1)
-                && listening.isEmpty()
-                && requested.isEmpty()
-                && legacyTypeMap.isEmpty();
+        return (mAllCallbacks.size() == 1)
+                && (mTrackingDefault.size() == 1)
+                && mListening.isEmpty()
+                && mRequested.isEmpty()
+                && mLegacyTypeMap.isEmpty();
     }
 
     boolean isListeningForAll() {
         final NetworkCapabilities empty = new NetworkCapabilities();
         empty.clearAll();
 
-        for (NetworkRequest req : listening.values()) {
-            if (req.networkCapabilities.equalRequestableCapabilities(empty)) {
+        for (NetworkRequestInfo nri : mListening.values()) {
+            if (nri.request.networkCapabilities.equalRequestableCapabilities(empty)) {
                 return true;
             }
         }
@@ -118,40 +130,67 @@
         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 : mTrackingDefault.keySet()) {
+            final NetworkRequestInfo nri = mTrackingDefault.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, boolean order, @Nullable Runnable inBetween) {
+        if (Objects.equals(mDefaultNetwork, agent)) return;
+
+        final TestNetworkAgent formerDefault = mDefaultNetwork;
+        mDefaultNetwork = agent;
+
+        if (order == CALLBACKS_FIRST) {
+            sendDefaultNetworkCallbacks(formerDefault, mDefaultNetwork);
+            if (inBetween != null) inBetween.run();
+            sendDefaultNetworkBroadcasts(formerDefault, mDefaultNetwork);
+        } else {
+            sendDefaultNetworkBroadcasts(formerDefault, mDefaultNetwork);
+            if (inBetween != null) inBetween.run();
+            sendDefaultNetworkCallbacks(formerDefault, mDefaultNetwork);
+        }
+    }
+
+    void makeDefaultNetwork(TestNetworkAgent agent, boolean order) {
+        makeDefaultNetwork(agent, order, null /* inBetween */);
+    }
+
+    void makeDefaultNetwork(TestNetworkAgent agent) {
+        makeDefaultNetwork(agent, BROADCAST_FIRST, null /* inBetween */);
+    }
+
     @Override
     public void requestNetwork(NetworkRequest req, NetworkCallback cb, Handler h) {
-        assertFalse(allCallbacks.containsKey(cb));
-        allCallbacks.put(cb, h);
+        assertFalse(mAllCallbacks.containsKey(cb));
+        mAllCallbacks.put(cb, new NetworkRequestInfo(req, h));
         if (mDefaultRequest.equals(req)) {
-            assertFalse(trackingDefault.contains(cb));
-            trackingDefault.add(cb);
+            assertFalse(mTrackingDefault.containsKey(cb));
+            mTrackingDefault.put(cb, new NetworkRequestInfo(req, h));
         } else {
-            assertFalse(requested.containsKey(cb));
-            requested.put(cb, req);
+            assertFalse(mRequested.containsKey(cb));
+            mRequested.put(cb, new NetworkRequestInfo(req, h));
         }
     }
 
@@ -163,22 +202,22 @@
     @Override
     public void requestNetwork(NetworkRequest req,
             int timeoutMs, int legacyType, Handler h, NetworkCallback cb) {
-        assertFalse(allCallbacks.containsKey(cb));
-        allCallbacks.put(cb, h);
-        assertFalse(requested.containsKey(cb));
-        requested.put(cb, req);
-        assertFalse(legacyTypeMap.containsKey(cb));
+        assertFalse(mAllCallbacks.containsKey(cb));
+        mAllCallbacks.put(cb, new NetworkRequestInfo(req, h));
+        assertFalse(mRequested.containsKey(cb));
+        mRequested.put(cb, new NetworkRequestInfo(req, h));
+        assertFalse(mLegacyTypeMap.containsKey(cb));
         if (legacyType != ConnectivityManager.TYPE_NONE) {
-            legacyTypeMap.put(cb, legacyType);
+            mLegacyTypeMap.put(cb, legacyType);
         }
     }
 
     @Override
     public void registerNetworkCallback(NetworkRequest req, NetworkCallback cb, Handler h) {
-        assertFalse(allCallbacks.containsKey(cb));
-        allCallbacks.put(cb, h);
-        assertFalse(listening.containsKey(cb));
-        listening.put(cb, req);
+        assertFalse(mAllCallbacks.containsKey(cb));
+        mAllCallbacks.put(cb, new NetworkRequestInfo(req, h));
+        assertFalse(mListening.containsKey(cb));
+        mListening.put(cb, new NetworkRequestInfo(req, h));
     }
 
     @Override
@@ -198,22 +237,22 @@
 
     @Override
     public void unregisterNetworkCallback(NetworkCallback cb) {
-        if (trackingDefault.contains(cb)) {
-            trackingDefault.remove(cb);
-        } else if (listening.containsKey(cb)) {
-            listening.remove(cb);
-        } else if (requested.containsKey(cb)) {
-            requested.remove(cb);
-            legacyTypeMap.remove(cb);
+        if (mTrackingDefault.containsKey(cb)) {
+            mTrackingDefault.remove(cb);
+        } else if (mListening.containsKey(cb)) {
+            mListening.remove(cb);
+        } else if (mRequested.containsKey(cb)) {
+            mRequested.remove(cb);
+            mLegacyTypeMap.remove(cb);
         } else {
             fail("Unexpected callback removed");
         }
-        allCallbacks.remove(cb);
+        mAllCallbacks.remove(cb);
 
-        assertFalse(allCallbacks.containsKey(cb));
-        assertFalse(trackingDefault.contains(cb));
-        assertFalse(listening.containsKey(cb));
-        assertFalse(requested.containsKey(cb));
+        assertFalse(mAllCallbacks.containsKey(cb));
+        assertFalse(mTrackingDefault.containsKey(cb));
+        assertFalse(mListening.containsKey(cb));
+        assertFalse(mRequested.containsKey(cb));
     }
 
     private void sendConnectivityAction(int type, boolean connected) {
@@ -275,34 +314,38 @@
         }
 
         public void fakeConnect() {
-            for (NetworkRequest request : cm.requested.values()) {
-                if (matchesLegacyType(request.legacyType)) {
+            for (NetworkRequestInfo nri : cm.mRequested.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));
+            for (NetworkCallback cb : cm.mListening.keySet()) {
+                final NetworkRequestInfo nri = cm.mListening.get(cb);
+                nri.handler.post(() -> cb.onAvailable(networkId));
+                nri.handler.post(() -> cb.onCapabilitiesChanged(
+                        networkId, copy(networkCapabilities)));
+                nri.handler.post(() -> cb.onLinkPropertiesChanged(networkId, copy(linkProperties)));
             }
+            // mTrackingDefault will be updated if/when the caller calls makeDefaultNetwork
         }
 
         public void fakeDisconnect() {
-            for (NetworkRequest request : cm.requested.values()) {
-                if (matchesLegacyType(request.legacyType)) {
+            for (NetworkRequestInfo nri : cm.mRequested.values()) {
+                if (matchesLegacyType(nri.request.legacyType)) {
                     cm.sendConnectivityAction(legacyType, false /* connected */);
                     break;
                 }
             }
-            for (NetworkCallback cb : cm.listening.keySet()) {
+            for (NetworkCallback cb : cm.mListening.keySet()) {
                 cb.onLost(networkId);
             }
+            // mTrackingDefault will be updated if/when the caller calls makeDefaultNetwork
         }
 
         public void sendLinkProperties() {
-            for (NetworkCallback cb : cm.listening.keySet()) {
+            for (NetworkCallback cb : cm.mListening.keySet()) {
                 cb.onLinkPropertiesChanged(networkId, copy(linkProperties));
             }
         }
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 0611086..d18d990 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java
@@ -59,6 +59,8 @@
 
 import static com.android.net.module.util.Inet4AddressUtils.inet4AddressToIntHTH;
 import static com.android.net.module.util.Inet4AddressUtils.intToInet4AddressHTH;
+import static com.android.networkstack.tethering.TestConnectivityManager.BROADCAST_FIRST;
+import static com.android.networkstack.tethering.TestConnectivityManager.CALLBACKS_FIRST;
 import static com.android.networkstack.tethering.Tethering.UserRestrictionActionListener;
 import static com.android.networkstack.tethering.TetheringNotificationUpdater.DOWNSTREAM_NONE;
 import static com.android.networkstack.tethering.UpstreamNetworkMonitor.EVENT_ON_CAPABILITIES;
@@ -618,6 +620,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 +719,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 +728,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 +755,7 @@
 
         mServiceContext.sendBroadcastAsUserMultiplePermissions(intent, UserHandle.ALL,
                 P2P_RECEIVER_PERMISSIONS_FOR_BROADCAST);
+        mLooper.dispatchAll();
     }
 
     private void sendUsbBroadcast(boolean connected, boolean configured, boolean function,
@@ -763,11 +769,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 +817,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 +864,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 +892,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 +953,6 @@
         initTetheringUpstream(upstreamState);
         prepareUsbTethering();
         sendUsbBroadcast(true, true, true, TETHERING_USB);
-        mLooper.dispatchAll();
     }
 
     private void assertSetIfaceToDadProxy(final int numOfCalls, final String ifaceName) {
@@ -1099,29 +1103,44 @@
         // Start USB tethering with no current upstream.
         prepareUsbTethering();
         sendUsbBroadcast(true, true, true, TETHERING_USB);
-        mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).startObserveAllNetworks();
-        inOrder.verify(mUpstreamNetworkMonitor).registerMobileNetworkRequest();
+        inOrder.verify(mUpstreamNetworkMonitor).setTryCell(true);
 
         // Pretend cellular connected and expect the upstream to be set.
         TestNetworkAgent mobile = new TestNetworkAgent(mCm, buildMobileDualStackUpstreamState());
         mobile.fakeConnect();
-        mCm.makeDefaultNetwork(mobile);
+        mCm.makeDefaultNetwork(mobile, BROADCAST_FIRST);
         mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(mobile.networkId);
 
         // Switch upstreams a few times.
-        // TODO: there may be a race where if the effects of the CONNECTIVITY_ACTION happen before
-        // UpstreamNetworkMonitor gets onCapabilitiesChanged on CALLBACK_DEFAULT_INTERNET, the
-        // upstream does not change. Extend TestConnectivityManager to simulate this condition and
-        // write a test for this.
         TestNetworkAgent wifi = new TestNetworkAgent(mCm, buildWifiUpstreamState());
         wifi.fakeConnect();
-        mCm.makeDefaultNetwork(wifi);
+        mCm.makeDefaultNetwork(wifi, BROADCAST_FIRST);
         mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(wifi.networkId);
 
-        mCm.makeDefaultNetwork(mobile);
+        // This code has historically been racy, so test different orderings of CONNECTIVITY_ACTION
+        // broadcasts and callbacks, and add mLooper.dispatchAll() calls between the two.
+        final Runnable doDispatchAll = () -> mLooper.dispatchAll();
+
+        mCm.makeDefaultNetwork(mobile, BROADCAST_FIRST, doDispatchAll);
+        mLooper.dispatchAll();
+        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(mobile.networkId);
+
+        mCm.makeDefaultNetwork(wifi, BROADCAST_FIRST, doDispatchAll);
+        mLooper.dispatchAll();
+        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(wifi.networkId);
+
+        mCm.makeDefaultNetwork(mobile, CALLBACKS_FIRST);
+        mLooper.dispatchAll();
+        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(mobile.networkId);
+
+        mCm.makeDefaultNetwork(wifi, CALLBACKS_FIRST);
+        mLooper.dispatchAll();
+        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(wifi.networkId);
+
+        mCm.makeDefaultNetwork(mobile, CALLBACKS_FIRST, doDispatchAll);
         mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(mobile.networkId);
 
@@ -1133,14 +1152,15 @@
         // Lose and regain upstream.
         assertTrue(mUpstreamNetworkMonitor.getCurrentPreferredUpstream().linkProperties
                 .hasIPv4Address());
+        mCm.makeDefaultNetwork(null, BROADCAST_FIRST, doDispatchAll);
+        mLooper.dispatchAll();
         mobile.fakeDisconnect();
-        mCm.makeDefaultNetwork(null);
         mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(null);
 
         mobile = new TestNetworkAgent(mCm, buildMobile464xlatUpstreamState());
         mobile.fakeConnect();
-        mCm.makeDefaultNetwork(mobile);
+        mCm.makeDefaultNetwork(mobile, BROADCAST_FIRST, doDispatchAll);
         mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(mobile.networkId);
 
@@ -1148,16 +1168,26 @@
         // mobile upstream, even though the netId is (unrealistically) the same.
         assertFalse(mUpstreamNetworkMonitor.getCurrentPreferredUpstream().linkProperties
                 .hasIPv4Address());
+
+        // Lose and regain upstream again.
+        mCm.makeDefaultNetwork(null, CALLBACKS_FIRST, doDispatchAll);
         mobile.fakeDisconnect();
-        mCm.makeDefaultNetwork(null);
         mLooper.dispatchAll();
         inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(null);
+
+        mobile = new TestNetworkAgent(mCm, buildMobileDualStackUpstreamState());
+        mobile.fakeConnect();
+        mCm.makeDefaultNetwork(mobile, CALLBACKS_FIRST, doDispatchAll);
+        mLooper.dispatchAll();
+        inOrder.verify(mUpstreamNetworkMonitor).setCurrentUpstream(mobile.networkId);
+
+        assertTrue(mUpstreamNetworkMonitor.getCurrentPreferredUpstream().linkProperties
+                .hasIPv4Address());
     }
 
     private void runNcmTethering() {
         prepareNcmTethering();
         sendUsbBroadcast(true, true, true, TETHERING_NCM);
-        mLooper.dispatchAll();
     }
 
     @Test
@@ -1205,7 +1235,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 +1262,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);
@@ -1251,7 +1279,7 @@
         verify(mUpstreamNetworkMonitor, times(1)).startObserveAllNetworks();
         // In tethering mode, in the default configuration, an explicit request
         // for a mobile network is also made.
-        verify(mUpstreamNetworkMonitor, times(1)).registerMobileNetworkRequest();
+        verify(mUpstreamNetworkMonitor, times(1)).setTryCell(true);
         // There are 2 IpServer state change events: STATE_AVAILABLE -> STATE_TETHERED
         verify(mNotificationUpdater, times(1)).onDownstreamChanged(DOWNSTREAM_NONE);
         verify(mNotificationUpdater, times(1)).onDownstreamChanged(eq(1 << TETHERING_WIFI));
@@ -1310,7 +1338,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 +1661,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 +1682,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 +1774,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 +1791,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 +1813,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 +1824,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 +1859,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 +1988,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 +2007,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 +2230,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 +2268,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 +2284,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 7d735fc..bc21692 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/UpstreamNetworkMonitorTest.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/UpstreamNetworkMonitorTest.java
@@ -41,7 +41,6 @@
 import static org.mockito.Mockito.when;
 
 import android.content.Context;
-import android.net.ConnectivityManager;
 import android.net.ConnectivityManager.NetworkCallback;
 import android.net.IConnectivityManager;
 import android.net.IpPrefix;
@@ -51,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;
@@ -101,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);
@@ -110,9 +114,8 @@
         when(mEntitleMgr.isCellularUpstreamPermitted()).thenReturn(true);
 
         mCM = spy(new TestConnectivityManager(mContext, mCS, sDefaultRequest));
-        mSM = new TestStateMachine();
-        mUNM = new UpstreamNetworkMonitor(
-                (ConnectivityManager) mCM, mSM, mLog, EVENT_UNM_UPDATE);
+        mSM = new TestStateMachine(mLooper.getLooper());
+        mUNM = new UpstreamNetworkMonitor(mCM, mSM, mLog, EVENT_UNM_UPDATE);
     }
 
     @After public void tearDown() throws Exception {
@@ -134,9 +137,9 @@
         assertTrue(mCM.hasNoCallbacks());
         assertFalse(mUNM.mobileNetworkRequested());
 
-        mUNM.updateMobileRequiresDun(true);
+        mUNM.setUpstreamConfig(false /* autoUpstream */, true /* dunRequired */);
         assertTrue(mCM.hasNoCallbacks());
-        mUNM.updateMobileRequiresDun(false);
+        mUNM.setUpstreamConfig(false /* autoUpstream */, false /* dunRequired */);
         assertTrue(mCM.hasNoCallbacks());
     }
 
@@ -146,7 +149,7 @@
         mUNM.startTrackDefaultNetwork(sDefaultRequest, mEntitleMgr);
 
         mUNM.startObserveAllNetworks();
-        assertEquals(1, mCM.trackingDefault.size());
+        assertEquals(1, mCM.mTrackingDefault.size());
 
         mUNM.stop();
         assertTrue(mCM.onlyHasDefaultCallbacks());
@@ -154,11 +157,11 @@
 
     @Test
     public void testListensForAllNetworks() throws Exception {
-        assertTrue(mCM.listening.isEmpty());
+        assertTrue(mCM.mListening.isEmpty());
 
         mUNM.startTrackDefaultNetwork(sDefaultRequest, mEntitleMgr);
         mUNM.startObserveAllNetworks();
-        assertFalse(mCM.listening.isEmpty());
+        assertFalse(mCM.mListening.isEmpty());
         assertTrue(mCM.isListeningForAll());
 
         mUNM.stop();
@@ -181,17 +184,17 @@
     @Test
     public void testRequestsMobileNetwork() throws Exception {
         assertFalse(mUNM.mobileNetworkRequested());
-        assertEquals(0, mCM.requested.size());
+        assertEquals(0, mCM.mRequested.size());
 
         mUNM.startObserveAllNetworks();
         assertFalse(mUNM.mobileNetworkRequested());
-        assertEquals(0, mCM.requested.size());
+        assertEquals(0, mCM.mRequested.size());
 
-        mUNM.updateMobileRequiresDun(false);
+        mUNM.setUpstreamConfig(false /* autoUpstream */, false /* dunRequired */);
         assertFalse(mUNM.mobileNetworkRequested());
-        assertEquals(0, mCM.requested.size());
+        assertEquals(0, mCM.mRequested.size());
 
-        mUNM.registerMobileNetworkRequest();
+        mUNM.setTryCell(true);
         assertTrue(mUNM.mobileNetworkRequested());
         assertUpstreamTypeRequested(TYPE_MOBILE_HIPRI);
         assertFalse(isDunRequested());
@@ -204,16 +207,16 @@
     @Test
     public void testDuplicateMobileRequestsIgnored() throws Exception {
         assertFalse(mUNM.mobileNetworkRequested());
-        assertEquals(0, mCM.requested.size());
+        assertEquals(0, mCM.mRequested.size());
 
         mUNM.startObserveAllNetworks();
         verify(mCM, times(1)).registerNetworkCallback(
                 any(NetworkRequest.class), any(NetworkCallback.class), any(Handler.class));
         assertFalse(mUNM.mobileNetworkRequested());
-        assertEquals(0, mCM.requested.size());
+        assertEquals(0, mCM.mRequested.size());
 
-        mUNM.updateMobileRequiresDun(true);
-        mUNM.registerMobileNetworkRequest();
+        mUNM.setUpstreamConfig(false /* autoUpstream */, true /* dunRequired */);
+        mUNM.setTryCell(true);
         verify(mCM, times(1)).requestNetwork(
                 any(NetworkRequest.class), anyInt(), anyInt(), any(Handler.class),
                 any(NetworkCallback.class));
@@ -223,9 +226,9 @@
         assertTrue(isDunRequested());
 
         // Try a few things that must not result in any state change.
-        mUNM.registerMobileNetworkRequest();
-        mUNM.updateMobileRequiresDun(true);
-        mUNM.registerMobileNetworkRequest();
+        mUNM.setTryCell(true);
+        mUNM.setUpstreamConfig(false /* autoUpstream */, true /* dunRequired */);
+        mUNM.setTryCell(true);
 
         assertTrue(mUNM.mobileNetworkRequested());
         assertUpstreamTypeRequested(TYPE_MOBILE_DUN);
@@ -240,17 +243,17 @@
     @Test
     public void testRequestsDunNetwork() throws Exception {
         assertFalse(mUNM.mobileNetworkRequested());
-        assertEquals(0, mCM.requested.size());
+        assertEquals(0, mCM.mRequested.size());
 
         mUNM.startObserveAllNetworks();
         assertFalse(mUNM.mobileNetworkRequested());
-        assertEquals(0, mCM.requested.size());
+        assertEquals(0, mCM.mRequested.size());
 
-        mUNM.updateMobileRequiresDun(true);
+        mUNM.setUpstreamConfig(false /* autoUpstream */, true /* dunRequired */);
         assertFalse(mUNM.mobileNetworkRequested());
-        assertEquals(0, mCM.requested.size());
+        assertEquals(0, mCM.mRequested.size());
 
-        mUNM.registerMobileNetworkRequest();
+        mUNM.setTryCell(true);
         assertTrue(mUNM.mobileNetworkRequested());
         assertUpstreamTypeRequested(TYPE_MOBILE_DUN);
         assertTrue(isDunRequested());
@@ -265,18 +268,18 @@
         mUNM.startObserveAllNetworks();
 
         // Test going from no-DUN to DUN correctly re-registers callbacks.
-        mUNM.updateMobileRequiresDun(false);
-        mUNM.registerMobileNetworkRequest();
+        mUNM.setUpstreamConfig(false /* autoUpstream */, false /* dunRequired */);
+        mUNM.setTryCell(true);
         assertTrue(mUNM.mobileNetworkRequested());
         assertUpstreamTypeRequested(TYPE_MOBILE_HIPRI);
         assertFalse(isDunRequested());
-        mUNM.updateMobileRequiresDun(true);
+        mUNM.setUpstreamConfig(false /* autoUpstream */, true /* dunRequired */);
         assertTrue(mUNM.mobileNetworkRequested());
         assertUpstreamTypeRequested(TYPE_MOBILE_DUN);
         assertTrue(isDunRequested());
 
         // Test going from DUN to no-DUN correctly re-registers callbacks.
-        mUNM.updateMobileRequiresDun(false);
+        mUNM.setUpstreamConfig(false /* autoUpstream */, false /* dunRequired */);
         assertTrue(mUNM.mobileNetworkRequested());
         assertUpstreamTypeRequested(TYPE_MOBILE_HIPRI);
         assertFalse(isDunRequested());
@@ -297,72 +300,78 @@
 
         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);
         // This is coupled with preferred types in TetheringConfiguration.
-        mUNM.updateMobileRequiresDun(true);
+        mUNM.setUpstreamConfig(false /* autoUpstream */, true /* dunRequired */);
         // DUN is available, but only use regular cell: no upstream selected.
         assertSatisfiesLegacyType(TYPE_NONE, mUNM.selectPreferredUpstreamType(preferredTypes));
         preferredTypes.remove(TYPE_MOBILE_DUN);
         // No WiFi, but our preferred flavour of cell is up.
         preferredTypes.add(TYPE_MOBILE_HIPRI);
         // This is coupled with preferred types in TetheringConfiguration.
-        mUNM.updateMobileRequiresDun(false);
+        mUNM.setUpstreamConfig(false /* autoUpstream */, false /* dunRequired */);
         assertSatisfiesLegacyType(TYPE_MOBILE_HIPRI,
                 mUNM.selectPreferredUpstreamType(preferredTypes));
         // Check to see we filed an explicit request.
-        assertEquals(1, mCM.requested.size());
-        NetworkRequest netReq = (NetworkRequest) mCM.requested.values().toArray()[0];
+        assertEquals(1, mCM.mRequested.size());
+        NetworkRequest netReq = ((NetworkRequestInfo) mCM.mRequested.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.
         when(mEntitleMgr.isCellularUpstreamPermitted()).thenReturn(false);
         assertSatisfiesLegacyType(TYPE_NONE, mUNM.selectPreferredUpstreamType(preferredTypes));
-        assertEquals(0, mCM.requested.size());
+        assertEquals(0, mCM.mRequested.size());
         // mobile change back to permitted, HIRPI should come back
         when(mEntitleMgr.isCellularUpstreamPermitted()).thenReturn(true);
         assertSatisfiesLegacyType(TYPE_MOBILE_HIPRI,
                 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());
+        assertEquals(0, mCM.mRequested.size());
 
         preferredTypes.remove(TYPE_MOBILE_HIPRI);
         preferredTypes.add(TYPE_MOBILE_DUN);
         // This is coupled with preferred types in TetheringConfiguration.
-        mUNM.updateMobileRequiresDun(true);
+        mUNM.setUpstreamConfig(false /* autoUpstream */, true /* dunRequired */);
         assertSatisfiesLegacyType(TYPE_WIFI, mUNM.selectPreferredUpstreamType(preferredTypes));
 
         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];
+        assertEquals(1, mCM.mRequested.size());
+        netReq = ((NetworkRequestInfo) mCM.mRequested.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.
         when(mEntitleMgr.isCellularUpstreamPermitted()).thenReturn(false);
         assertSatisfiesLegacyType(TYPE_NONE, mUNM.selectPreferredUpstreamType(preferredTypes));
-        assertEquals(0, mCM.requested.size());
+        assertEquals(0, mCM.mRequested.size());
         // mobile change back to permitted, DUN should come back
         when(mEntitleMgr.isCellularUpstreamPermitted()).thenReturn(true);
         assertSatisfiesLegacyType(TYPE_MOBILE_DUN,
@@ -373,49 +382,72 @@
     public void testGetCurrentPreferredUpstream() throws Exception {
         mUNM.startTrackDefaultNetwork(sDefaultRequest, mEntitleMgr);
         mUNM.startObserveAllNetworks();
-        mUNM.updateMobileRequiresDun(false);
+        mUNM.setUpstreamConfig(true /* autoUpstream */, false /* dunRequired */);
+        mUNM.setTryCell(true);
 
         // [0] Mobile connects, DUN not required -> mobile selected.
         final TestNetworkAgent cellAgent = new TestNetworkAgent(mCM, CELL_CAPABILITIES);
         cellAgent.fakeConnect();
         mCM.makeDefaultNetwork(cellAgent);
+        mLooper.dispatchAll();
         assertEquals(cellAgent.networkId, mUNM.getCurrentPreferredUpstream().network);
+        assertEquals(0, mCM.mRequested.size());
 
         // [1] Mobile connects but not permitted -> null selected
         when(mEntitleMgr.isCellularUpstreamPermitted()).thenReturn(false);
         assertEquals(null, mUNM.getCurrentPreferredUpstream());
         when(mEntitleMgr.isCellularUpstreamPermitted()).thenReturn(true);
+        assertEquals(0, mCM.mRequested.size());
 
         // [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.mRequested.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.mRequested.size());
 
         // [4] DUN required, no other changes -> WiFi still selected
-        mUNM.updateMobileRequiresDun(true);
+        mUNM.setUpstreamConfig(false /* autoUpstream */, true /* dunRequired */);
         assertEquals(wifiAgent.networkId, mUNM.getCurrentPreferredUpstream().network);
+        assertEquals(1, mCM.mRequested.size());
+        assertTrue(isDunRequested());
 
         // [5] WiFi no longer validated, mobile becomes default, DUN required -> null selected.
         mCM.makeDefaultNetwork(cellAgent);
+        mLooper.dispatchAll();
         assertEquals(null, mUNM.getCurrentPreferredUpstream());
-        // TODO: make sure that a DUN request has been filed. This is currently
-        // triggered by code over in Tethering, but once that has been moved
-        // into UNM we should test for this here.
+        assertEquals(1, mCM.mRequested.size());
+        assertTrue(isDunRequested());
 
         // [6] DUN network arrives -> DUN selected
         final TestNetworkAgent dunAgent = new TestNetworkAgent(mCM, CELL_CAPABILITIES);
         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.mRequested.size());
 
         // [7] Mobile is not permitted -> null selected
         when(mEntitleMgr.isCellularUpstreamPermitted()).thenReturn(false);
         assertEquals(null, mUNM.getCurrentPreferredUpstream());
+        assertEquals(1, mCM.mRequested.size());
+
+        // [7] Mobile is permitted again -> DUN selected
+        when(mEntitleMgr.isCellularUpstreamPermitted()).thenReturn(true);
+        assertEquals(dunAgent.networkId, mUNM.getCurrentPreferredUpstream().network);
+        assertEquals(1, mCM.mRequested.size());
+
+        // [8] DUN no longer required -> request is withdrawn
+        mUNM.setUpstreamConfig(true /* autoUpstream */, false /* dunRequired */);
+        assertEquals(0, mCM.mRequested.size());
+        assertFalse(isDunRequested());
     }
 
     @Test
@@ -445,6 +477,7 @@
         }
         wifiAgent.fakeConnect();
         wifiAgent.sendLinkProperties();
+        mLooper.dispatchAll();
 
         local = mUNM.getLocalPrefixes();
         assertPrefixSet(local, INCLUDES, alreadySeen);
@@ -469,6 +502,7 @@
         }
         cellAgent.fakeConnect();
         cellAgent.sendLinkProperties();
+        mLooper.dispatchAll();
 
         local = mUNM.getLocalPrefixes();
         assertPrefixSet(local, INCLUDES, alreadySeen);
@@ -490,6 +524,7 @@
         }
         dunAgent.fakeConnect();
         dunAgent.sendLinkProperties();
+        mLooper.dispatchAll();
 
         local = mUNM.getLocalPrefixes();
         assertPrefixSet(local, INCLUDES, alreadySeen);
@@ -501,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);
@@ -508,6 +544,7 @@
 
         // [5] Pretend mobile disconnected.
         cellAgent.fakeDisconnect();
+        mLooper.dispatchAll();
         local = mUNM.getLocalPrefixes();
         assertPrefixSet(local, EXCLUDES, wifiLinkPrefixes);
         assertPrefixSet(local, EXCLUDES, cellLinkPrefixes);
@@ -515,6 +552,7 @@
 
         // [6] Pretend DUN disconnected.
         dunAgent.fakeDisconnect();
+        mLooper.dispatchAll();
         local = mUNM.getLocalPrefixes();
         assertTrue(local.isEmpty());
     }
@@ -534,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));
@@ -552,15 +591,15 @@
     }
 
     private void assertUpstreamTypeRequested(int upstreamType) throws Exception {
-        assertEquals(1, mCM.requested.size());
-        assertEquals(1, mCM.legacyTypeMap.size());
+        assertEquals(1, mCM.mRequested.size());
+        assertEquals(1, mCM.mLegacyTypeMap.size());
         assertEquals(Integer.valueOf(upstreamType),
-                mCM.legacyTypeMap.values().iterator().next());
+                mCM.mLegacyTypeMap.values().iterator().next());
     }
 
     private boolean isDunRequested() {
-        for (NetworkRequest req : mCM.requested.values()) {
-            if (req.networkCapabilities.hasCapability(NET_CAPABILITY_DUN)) {
+        for (NetworkRequestInfo nri : mCM.mRequested.values()) {
+            if (nri.request.networkCapabilities.hasCapability(NET_CAPABILITY_DUN)) {
                 return true;
             }
         }
@@ -586,8 +625,8 @@
             }
         }
 
-        public TestStateMachine() {
-            super("UpstreamNetworkMonitor.TestStateMachine");
+        public TestStateMachine(Looper looper) {
+            super("UpstreamNetworkMonitor.TestStateMachine", looper);
             addState(mLoggingState);
             setInitialState(mLoggingState);
             super.start();
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java
index f423503..48923b9 100644
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java
@@ -25,7 +25,8 @@
 import static com.android.cts.net.hostside.NetworkPolicyTestUtils.getConnectivityManager;
 import static com.android.cts.net.hostside.NetworkPolicyTestUtils.getContext;
 import static com.android.cts.net.hostside.NetworkPolicyTestUtils.getInstrumentation;
-import static com.android.cts.net.hostside.NetworkPolicyTestUtils.getWifiManager;
+import static com.android.cts.net.hostside.NetworkPolicyTestUtils.isAppStandbySupported;
+import static com.android.cts.net.hostside.NetworkPolicyTestUtils.isBatterySaverSupported;
 import static com.android.cts.net.hostside.NetworkPolicyTestUtils.isDozeModeSupported;
 import static com.android.cts.net.hostside.NetworkPolicyTestUtils.restrictBackgroundValueToString;
 
@@ -46,12 +47,10 @@
 import android.net.ConnectivityManager;
 import android.net.NetworkInfo.DetailedState;
 import android.net.NetworkInfo.State;
-import android.net.wifi.WifiManager;
 import android.os.BatteryManager;
 import android.os.Binder;
 import android.os.Bundle;
 import android.os.SystemClock;
-import android.provider.Settings;
 import android.service.notification.NotificationListenerService;
 import android.util.Log;
 
@@ -628,6 +627,9 @@
     }
 
     protected void setBatterySaverMode(boolean enabled) throws Exception {
+        if (!isBatterySaverSupported()) {
+            return;
+        }
         Log.i(TAG, "Setting Battery Saver Mode to " + enabled);
         if (enabled) {
             turnBatteryOn();
@@ -639,8 +641,9 @@
     }
 
     protected void setDozeMode(boolean enabled) throws Exception {
-        // Check doze mode is supported.
-        assertTrue("Device does not support Doze Mode", isDozeModeSupported());
+        if (!isDozeModeSupported()) {
+            return;
+        }
 
         Log.i(TAG, "Setting Doze Mode to " + enabled);
         if (enabled) {
@@ -660,12 +663,18 @@
     }
 
     protected void setAppIdle(boolean enabled) throws Exception {
+        if (!isAppStandbySupported()) {
+            return;
+        }
         Log.i(TAG, "Setting app idle to " + enabled);
         executeSilentShellCommand("am set-inactive " + TEST_APP2_PKG + " " + enabled );
         assertAppIdle(enabled);
     }
 
     protected void setAppIdleNoAssert(boolean enabled) throws Exception {
+        if (!isAppStandbySupported()) {
+            return;
+        }
         Log.i(TAG, "Setting app idle to " + enabled);
         executeSilentShellCommand("am set-inactive " + TEST_APP2_PKG + " " + enabled );
     }
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/DumpOnFailureRule.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/DumpOnFailureRule.java
index 5ecb399..66cb935 100644
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/DumpOnFailureRule.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/DumpOnFailureRule.java
@@ -24,6 +24,8 @@
 import android.os.ParcelFileDescriptor;
 import android.util.Log;
 
+import androidx.test.platform.app.InstrumentationRegistry;
+
 import com.android.compatibility.common.util.OnFailureRule;
 
 import org.junit.AssumptionViolatedException;
@@ -37,23 +39,20 @@
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 
-import androidx.test.platform.app.InstrumentationRegistry;
-
 public class DumpOnFailureRule extends OnFailureRule {
     private File mDumpDir = new File(Environment.getExternalStorageDirectory(),
             "CtsHostsideNetworkTests");
 
     @Override
     public void onTestFailure(Statement base, Description description, Throwable throwable) {
-        final String testName = description.getClassName() + "_" + description.getMethodName();
-
         if (throwable instanceof AssumptionViolatedException) {
+            final String testName = description.getClassName() + "_" + description.getMethodName();
             Log.d(TAG, "Skipping test " + testName + ": " + throwable);
             return;
         }
 
         prepareDumpRootDir();
-        final File dumpFile = new File(mDumpDir, "dump-" + testName);
+        final File dumpFile = new File(mDumpDir, "dump-" + getShortenedTestName(description));
         Log.i(TAG, "Dumping debug info for " + description + ": " + dumpFile.getPath());
         try (FileOutputStream out = new FileOutputStream(dumpFile)) {
             for (String cmd : new String[] {
@@ -71,6 +70,17 @@
         }
     }
 
+    private String getShortenedTestName(Description description) {
+        final String qualifiedClassName = description.getClassName();
+        final String className = qualifiedClassName.substring(
+                qualifiedClassName.lastIndexOf(".") + 1);
+        final String shortenedClassName = className.chars()
+                .filter(Character::isUpperCase)
+                .collect(StringBuilder::new, StringBuilder::appendCodePoint, StringBuilder::append)
+                .toString();
+        return shortenedClassName + "_" + description.getMethodName();
+    }
+
     void dumpCommandOutput(FileOutputStream out, String cmd) {
         final ParcelFileDescriptor pfd = InstrumentationRegistry.getInstrumentation()
                 .getUiAutomation().executeShellCommand(cmd);
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/NetworkPolicyTestUtils.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/NetworkPolicyTestUtils.java
index b61535b..0a13408 100644
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/NetworkPolicyTestUtils.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/NetworkPolicyTestUtils.java
@@ -22,12 +22,13 @@
 import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_METERED;
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
+import static android.net.wifi.WifiConfiguration.METERED_OVERRIDE_METERED;
+import static android.net.wifi.WifiConfiguration.METERED_OVERRIDE_NONE;
 
 import static com.android.compatibility.common.util.SystemUtil.runShellCommand;
 import static com.android.cts.net.hostside.AbstractRestrictBackgroundNetworkTestCase.TAG;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
@@ -35,19 +36,21 @@
 
 import android.app.ActivityManager;
 import android.app.Instrumentation;
+import android.app.UiAutomation;
 import android.content.Context;
 import android.location.LocationManager;
 import android.net.ConnectivityManager;
 import android.net.ConnectivityManager.NetworkCallback;
 import android.net.Network;
 import android.net.NetworkCapabilities;
+import android.net.wifi.WifiConfiguration;
 import android.net.wifi.WifiManager;
+import android.net.wifi.WifiManager.ActionListener;
 import android.os.PersistableBundle;
 import android.os.Process;
 import android.telephony.CarrierConfigManager;
 import android.telephony.SubscriptionManager;
 import android.telephony.data.ApnSetting;
-import android.text.TextUtils;
 import android.util.Log;
 
 import androidx.test.platform.app.InstrumentationRegistry;
@@ -57,7 +60,12 @@
 import com.android.compatibility.common.util.ShellIdentityUtils;
 import com.android.compatibility.common.util.ThrowingRunnable;
 
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 
 public class NetworkPolicyTestUtils {
@@ -98,11 +106,11 @@
         if (mDataSaverSupported == null) {
             assertMyRestrictBackgroundStatus(RESTRICT_BACKGROUND_STATUS_DISABLED);
             try {
-                setRestrictBackground(true);
+                setRestrictBackgroundInternal(true);
                 mDataSaverSupported = !isMyRestrictBackgroundStatus(
                         RESTRICT_BACKGROUND_STATUS_DISABLED);
             } finally {
-                setRestrictBackground(false);
+                setRestrictBackgroundInternal(false);
             }
         }
         return mDataSaverSupported;
@@ -182,19 +190,14 @@
     }
 
     private static String getWifiSsid() {
-        final boolean isLocationEnabled = isLocationEnabled();
+        final UiAutomation uiAutomation = getInstrumentation().getUiAutomation();
         try {
-            if (!isLocationEnabled) {
-                setLocationEnabled(true);
-            }
-            final String ssid = unquoteSSID(getWifiManager().getConnectionInfo().getSSID());
+            uiAutomation.adoptShellPermissionIdentity();
+            final String ssid = getWifiManager().getConnectionInfo().getSSID();
             assertNotEquals(WifiManager.UNKNOWN_SSID, ssid);
             return ssid;
         } finally {
-            // Reset the location enabled state
-            if (!isLocationEnabled) {
-                setLocationEnabled(false);
-            }
+            uiAutomation.dropShellPermissionIdentity();
         }
     }
 
@@ -205,18 +208,71 @@
     }
 
     private static void setWifiMeteredStatus(String ssid, boolean metered) throws Exception {
-        assertFalse("SSID should not be empty", TextUtils.isEmpty(ssid));
-        final String cmd = "cmd netpolicy set metered-network " + ssid + " " + metered;
-        executeShellCommand(cmd);
-        assertWifiMeteredStatus(ssid, metered);
-        assertActiveNetworkMetered(metered);
+        final UiAutomation uiAutomation = getInstrumentation().getUiAutomation();
+        try {
+            uiAutomation.adoptShellPermissionIdentity();
+            final WifiConfiguration currentConfig = getWifiConfiguration(ssid);
+            currentConfig.meteredOverride = metered
+                    ? METERED_OVERRIDE_METERED : METERED_OVERRIDE_NONE;
+            BlockingQueue<Integer> blockingQueue = new LinkedBlockingQueue<>();
+            getWifiManager().save(currentConfig, createActionListener(
+                    blockingQueue, Integer.MAX_VALUE));
+            Integer resultCode = blockingQueue.poll(TIMEOUT_CHANGE_METEREDNESS_MS,
+                    TimeUnit.MILLISECONDS);
+            if (resultCode == null) {
+                fail("Timed out waiting for meteredness to change; ssid=" + ssid
+                        + ", metered=" + metered);
+            } else if (resultCode != Integer.MAX_VALUE) {
+                fail("Error overriding the meteredness; ssid=" + ssid
+                        + ", metered=" + metered + ", error=" + resultCode);
+            }
+            final boolean success = assertActiveNetworkMetered(metered, false /* throwOnFailure */);
+            if (!success) {
+                Log.i(TAG, "Retry connecting to wifi; ssid=" + ssid);
+                blockingQueue = new LinkedBlockingQueue<>();
+                getWifiManager().connect(currentConfig, createActionListener(
+                        blockingQueue, Integer.MAX_VALUE));
+                resultCode = blockingQueue.poll(TIMEOUT_CHANGE_METEREDNESS_MS,
+                        TimeUnit.MILLISECONDS);
+                if (resultCode == null) {
+                    fail("Timed out waiting for wifi to connect; ssid=" + ssid);
+                } else if (resultCode != Integer.MAX_VALUE) {
+                    fail("Error connecting to wifi; ssid=" + ssid
+                            + ", error=" + resultCode);
+                }
+                assertActiveNetworkMetered(metered, true /* throwOnFailure */);
+            }
+        } finally {
+            uiAutomation.dropShellPermissionIdentity();
+        }
     }
 
-    private static void assertWifiMeteredStatus(String ssid, boolean expectedMeteredStatus) {
-        final String result = executeShellCommand("cmd netpolicy list wifi-networks");
-        final String expectedLine = ssid + ";" + expectedMeteredStatus;
-        assertTrue("Expected line: " + expectedLine + "; Actual result: " + result,
-                result.contains(expectedLine));
+    private static WifiConfiguration getWifiConfiguration(String ssid) {
+        final List<String> ssids = new ArrayList<>();
+        for (WifiConfiguration config : getWifiManager().getConfiguredNetworks()) {
+            if (config.SSID.equals(ssid)) {
+                return config;
+            }
+            ssids.add(config.SSID);
+        }
+        fail("Couldn't find the wifi config; ssid=" + ssid
+                + ", all=" + Arrays.toString(ssids.toArray()));
+        return null;
+    }
+
+    private static ActionListener createActionListener(BlockingQueue<Integer> blockingQueue,
+            int successCode) {
+        return new ActionListener() {
+            @Override
+            public void onSuccess() {
+                blockingQueue.offer(successCode);
+            }
+
+            @Override
+            public void onFailure(int reason) {
+                blockingQueue.offer(reason);
+            }
+        };
     }
 
     private static void setCellularMeteredStatus(int subId, boolean metered) throws Exception {
@@ -225,11 +281,11 @@
                 new String[] {ApnSetting.TYPE_MMS_STRING});
         ShellIdentityUtils.invokeMethodWithShellPermissionsNoReturn(getCarrierConfigManager(),
                 (cm) -> cm.overrideConfig(subId, metered ? null : bundle));
-        assertActiveNetworkMetered(metered);
+        assertActiveNetworkMetered(metered, true /* throwOnFailure */);
     }
 
-    // Copied from cts/tests/tests/net/src/android/net/cts/ConnectivityManagerTest.java
-    private static void assertActiveNetworkMetered(boolean expectedMeteredStatus) throws Exception {
+    private static boolean assertActiveNetworkMetered(boolean expectedMeteredStatus,
+            boolean throwOnFailure) throws Exception {
         final CountDownLatch latch = new CountDownLatch(1);
         final NetworkCallback networkCallback = new NetworkCallback() {
             @Override
@@ -246,16 +302,29 @@
         getConnectivityManager().registerDefaultNetworkCallback(networkCallback);
         try {
             if (!latch.await(TIMEOUT_CHANGE_METEREDNESS_MS, TimeUnit.MILLISECONDS)) {
-                fail("Timed out waiting for active network metered status to change to "
-                        + expectedMeteredStatus + "; network = "
-                        + getConnectivityManager().getActiveNetwork());
+                final String errorMsg = "Timed out waiting for active network metered status "
+                        + "to change to " + expectedMeteredStatus + "; network = "
+                        + getConnectivityManager().getActiveNetwork();
+                if (throwOnFailure) {
+                    fail(errorMsg);
+                }
+                Log.w(TAG, errorMsg);
+                return false;
             }
+            return true;
         } finally {
             getConnectivityManager().unregisterNetworkCallback(networkCallback);
         }
     }
 
     public static void setRestrictBackground(boolean enabled) {
+        if (!isDataSaverSupported()) {
+            return;
+        }
+        setRestrictBackgroundInternal(enabled);
+    }
+
+    private static void setRestrictBackgroundInternal(boolean enabled) {
         executeShellCommand("cmd netpolicy set restrict-background " + enabled);
         final String output = executeShellCommand("cmd netpolicy get restrict-background");
         final String expectedSuffix = enabled ? "enabled" : "disabled";
diff --git a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
index bd1139d..acf7d26 100644
--- a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
+++ b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
@@ -18,6 +18,8 @@
 import android.app.Instrumentation
 import android.content.Context
 import android.net.ConnectivityManager
+import android.net.INetworkAgent
+import android.net.INetworkAgentRegistry
 import android.net.InetAddresses
 import android.net.IpPrefix
 import android.net.KeepalivePacketData
@@ -66,8 +68,6 @@
 import android.os.Message
 import android.util.DebugUtils.valueToString
 import androidx.test.InstrumentationRegistry
-import com.android.connectivity.aidl.INetworkAgent
-import com.android.connectivity.aidl.INetworkAgentRegistry
 import com.android.modules.utils.build.SdkLevel
 import com.android.net.module.util.ArrayTrackRecord
 import com.android.testutils.CompatUtil
@@ -547,6 +547,7 @@
     @Test
     @IgnoreUpTo(Build.VERSION_CODES.R)
     fun testSetUnderlyingNetworksAndVpnSpecifier() {
+        val mySessionId = "MySession12345"
         val request = NetworkRequest.Builder()
                 .addTransportType(TRANSPORT_TEST)
                 .addTransportType(TRANSPORT_VPN)
@@ -560,7 +561,7 @@
             addTransportType(TRANSPORT_TEST)
             addTransportType(TRANSPORT_VPN)
             removeCapability(NET_CAPABILITY_NOT_VPN)
-            setTransportInfo(VpnTransportInfo(VpnManager.TYPE_VPN_SERVICE))
+            setTransportInfo(VpnTransportInfo(VpnManager.TYPE_VPN_SERVICE, mySessionId))
             if (SdkLevel.isAtLeastS()) {
                 addCapability(NET_CAPABILITY_NOT_VCN_MANAGED)
             }
@@ -580,6 +581,8 @@
         assertNotNull(vpnNc)
         assertEquals(VpnManager.TYPE_VPN_SERVICE,
                 (vpnNc.transportInfo as VpnTransportInfo).type)
+        // TODO: b/183938194 please fix the issue and enable following check.
+        // assertEquals(mySessionId, (vpnNc.transportInfo as VpnTransportInfo).sessionId)
 
         val testAndVpn = intArrayOf(TRANSPORT_TEST, TRANSPORT_VPN)
         assertTrue(hasAllTransports(vpnNc, testAndVpn))
diff --git a/tests/cts/net/src/android/net/cts/NetworkRequestTest.java b/tests/cts/net/src/android/net/cts/NetworkRequestTest.java
index 30c4e72..9906c30 100644
--- a/tests/cts/net/src/android/net/cts/NetworkRequestTest.java
+++ b/tests/cts/net/src/android/net/cts/NetworkRequestTest.java
@@ -39,17 +39,20 @@
 import android.net.NetworkCapabilities;
 import android.net.NetworkRequest;
 import android.net.NetworkSpecifier;
-import android.net.UidRange;
 import android.net.wifi.WifiNetworkSpecifier;
 import android.os.Build;
 import android.os.PatternMatcher;
 import android.os.Process;
 import android.util.ArraySet;
+import android.util.Range;
 
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.modules.utils.build.SdkLevel;
 import com.android.networkstack.apishim.ConstantsShim;
+import com.android.networkstack.apishim.NetworkRequestShimImpl;
+import com.android.networkstack.apishim.common.NetworkRequestShim;
+import com.android.networkstack.apishim.common.UnsupportedApiLevelException;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
 
@@ -57,6 +60,8 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.util.Set;
+
 @RunWith(AndroidJUnit4.class)
 public class NetworkRequestTest {
     @Rule
@@ -225,6 +230,14 @@
         assertTrue(requestCellularInternet.canBeSatisfiedBy(capCellularVpnMmsInternet));
     }
 
+    private void setUids(NetworkRequest.Builder builder, Set<Range<Integer>> ranges)
+            throws UnsupportedApiLevelException {
+        if (SdkLevel.isAtLeastS()) {
+            final NetworkRequestShim networkRequestShim = NetworkRequestShimImpl.newInstance();
+            networkRequestShim.setUids(builder, ranges);
+        }
+    }
+
     @Test
     @IgnoreUpTo(Build.VERSION_CODES.Q)
     public void testInvariantInCanBeSatisfiedBy() {
@@ -232,15 +245,26 @@
         // NetworkCapabilities.satisfiedByNetworkCapabilities().
         final LocalNetworkSpecifier specifier1 = new LocalNetworkSpecifier(1234 /* id */);
         final int uid = Process.myUid();
-        final ArraySet<UidRange> ranges = new ArraySet<>();
-        ranges.add(new UidRange(uid, uid));
-        final NetworkRequest requestCombination = new NetworkRequest.Builder()
+        final NetworkRequest.Builder nrBuilder = new NetworkRequest.Builder()
                 .addTransportType(TRANSPORT_CELLULAR)
                 .addCapability(NET_CAPABILITY_INTERNET)
                 .setLinkUpstreamBandwidthKbps(1000)
                 .setNetworkSpecifier(specifier1)
-                .setSignalStrength(-123)
-                .setUids(ranges).build();
+                .setSignalStrength(-123);
+
+        // The uid ranges should be set into the request, but setUids() takes a set of UidRange
+        // that is hidden and inaccessible from shims. Before, S setUids will be a no-op. But
+        // because NetworkRequest.Builder sets the UID of the request to the current UID, the
+        // request contains the current UID both on S and before S.
+        final Set<Range<Integer>> ranges = new ArraySet<>();
+        ranges.add(new Range<Integer>(uid, uid));
+        try {
+            setUids(nrBuilder, ranges);
+        } catch (UnsupportedApiLevelException e) {
+            // Not supported before API31.
+        }
+        final NetworkRequest requestCombination = nrBuilder.build();
+
         final NetworkCapabilities capCell = new NetworkCapabilities.Builder()
                 .addTransportType(TRANSPORT_CELLULAR).build();
         assertCorrectlySatisfies(false, requestCombination, capCell);
diff --git a/tests/cts/net/src/android/net/cts/VpnServiceTest.java b/tests/cts/net/src/android/net/cts/VpnServiceTest.java
index 15af23c..5c7b5ca 100644
--- a/tests/cts/net/src/android/net/cts/VpnServiceTest.java
+++ b/tests/cts/net/src/android/net/cts/VpnServiceTest.java
@@ -47,6 +47,7 @@
         assertEquals(1, count);
     }
 
+    @AppModeFull(reason = "establish() requires prepare(), which requires PackageManager access")
     public void testEstablish() throws Exception {
         ParcelFileDescriptor descriptor = null;
         try {
@@ -62,7 +63,7 @@
         }
     }
 
-    @AppModeFull(reason = "Socket cannot bind in instant app mode")
+    @AppModeFull(reason = "Protecting sockets requires prepare(), which requires PackageManager")
     public void testProtect_DatagramSocket() throws Exception {
         DatagramSocket socket = new DatagramSocket();
         try {
@@ -77,6 +78,7 @@
         }
     }
 
+    @AppModeFull(reason = "Protecting sockets requires prepare(), which requires PackageManager")
     public void testProtect_Socket() throws Exception {
         Socket socket = new Socket();
         try {
@@ -91,7 +93,7 @@
         }
     }
 
-    @AppModeFull(reason = "Socket cannot bind in instant app mode")
+    @AppModeFull(reason = "Protecting sockets requires prepare(), which requires PackageManager")
     public void testProtect_int() throws Exception {
         DatagramSocket socket = new DatagramSocket();
         ParcelFileDescriptor descriptor = ParcelFileDescriptor.fromDatagramSocket(socket);