Merge changes I24580ea4,Id48b8eab

* changes:
  Allow the system to register 250 NetworkCallbacks.
  Add more test coverage for limiting the number of callbacks.
diff --git a/framework/api/current.txt b/framework/api/current.txt
index ad44b27..0a9560a 100644
--- a/framework/api/current.txt
+++ b/framework/api/current.txt
@@ -68,6 +68,7 @@
     method public boolean bindProcessToNetwork(@Nullable android.net.Network);
     method @NonNull public android.net.SocketKeepalive createSocketKeepalive(@NonNull android.net.Network, @NonNull android.net.IpSecManager.UdpEncapsulationSocket, @NonNull java.net.InetAddress, @NonNull java.net.InetAddress, @NonNull java.util.concurrent.Executor, @NonNull android.net.SocketKeepalive.Callback);
     method @Nullable @RequiresPermission(android.Manifest.permission.ACCESS_NETWORK_STATE) public android.net.Network getActiveNetwork();
+    method @Nullable @RequiresPermission(android.Manifest.permission.NETWORK_STACK) public android.net.Network getActiveNetworkForUid(int);
     method @Deprecated @Nullable @RequiresPermission(android.Manifest.permission.ACCESS_NETWORK_STATE) public android.net.NetworkInfo getActiveNetworkInfo();
     method @Deprecated @NonNull @RequiresPermission(android.Manifest.permission.ACCESS_NETWORK_STATE) public android.net.NetworkInfo[] getAllNetworkInfo();
     method @NonNull @RequiresPermission(android.Manifest.permission.ACCESS_NETWORK_STATE) public android.net.Network[] getAllNetworks();
@@ -387,7 +388,9 @@
   public class NetworkRequest implements android.os.Parcelable {
     method public boolean canBeSatisfiedBy(@Nullable android.net.NetworkCapabilities);
     method public int describeContents();
+    method @NonNull public int[] getCapabilities();
     method @Nullable public android.net.NetworkSpecifier getNetworkSpecifier();
+    method @NonNull public int[] getTransportTypes();
     method public boolean hasCapability(int);
     method public boolean hasTransport(int);
     method public void writeToParcel(android.os.Parcel, int);
diff --git a/framework/api/module-lib-current.txt b/framework/api/module-lib-current.txt
index 35e45ec..cd96a1b 100644
--- a/framework/api/module-lib-current.txt
+++ b/framework/api/module-lib-current.txt
@@ -120,6 +120,7 @@
   }
 
   public class NetworkRequest implements android.os.Parcelable {
+    method @NonNull public int[] getUnwantedCapabilities();
     method public boolean hasUnwantedCapability(int);
   }
 
diff --git a/framework/api/system-current.txt b/framework/api/system-current.txt
index 1ee79a4..95ad694 100644
--- a/framework/api/system-current.txt
+++ b/framework/api/system-current.txt
@@ -219,7 +219,7 @@
     method public void onAutomaticReconnectDisabled();
     method public void onBandwidthUpdateRequested();
     method public void onNetworkCreated();
-    method public void onNetworkDisconnected();
+    method public void onNetworkDestroyed();
     method public void onNetworkUnwanted();
     method public void onQosCallbackRegistered(int, @NonNull android.net.QosFilter);
     method public void onQosCallbackUnregistered(int);
@@ -238,6 +238,7 @@
     method public final void sendQosSessionLost(int, int, int);
     method public final void sendSocketKeepaliveEvent(int, int);
     method @Deprecated public void setLegacySubtype(int, @NonNull String);
+    method public void setTeardownDelayMs(@IntRange(from=0, to=0x1388) int);
     method public final void setUnderlyingNetworks(@Nullable java.util.List<android.net.Network>);
     method public void unregister();
     field public static final int VALIDATION_STATUS_NOT_VALID = 2; // 0x2
diff --git a/framework/src/android/net/ConnectivityManager.java b/framework/src/android/net/ConnectivityManager.java
index d196c1a..c6f4e0b 100644
--- a/framework/src/android/net/ConnectivityManager.java
+++ b/framework/src/android/net/ConnectivityManager.java
@@ -921,6 +921,7 @@
             BLOCKED_REASON_DOZE,
             BLOCKED_REASON_APP_STANDBY,
             BLOCKED_REASON_RESTRICTED_MODE,
+            BLOCKED_REASON_LOCKDOWN_VPN,
             BLOCKED_METERED_REASON_DATA_SAVER,
             BLOCKED_METERED_REASON_USER_RESTRICTED,
             BLOCKED_METERED_REASON_ADMIN_DISABLED,
@@ -1189,8 +1190,7 @@
      *
      * @return a {@link Network} object for the current default network for the
      *         given UID or {@code null} if no default network is currently active
-     *
-     * @hide
+     * TODO: b/183465229 Cleanup getActiveNetworkForUid once b/165835257 is fixed
      */
     @RequiresPermission(android.Manifest.permission.NETWORK_STACK)
     @Nullable
@@ -3659,7 +3659,8 @@
         public void onBlockedStatusChanged(@NonNull Network network, boolean blocked) {}
 
         /**
-         * Called when access to the specified network is blocked or unblocked.
+         * Called when access to the specified network is blocked or unblocked, or the reason for
+         * access being blocked changes.
          *
          * If a NetworkCallback object implements this method,
          * {@link #onBlockedStatusChanged(Network, boolean)} will not be called.
diff --git a/framework/src/android/net/INetworkAgent.aidl b/framework/src/android/net/INetworkAgent.aidl
index f9d3994..d941d4b 100644
--- a/framework/src/android/net/INetworkAgent.aidl
+++ b/framework/src/android/net/INetworkAgent.aidl
@@ -47,5 +47,5 @@
     void onQosFilterCallbackRegistered(int qosCallbackId, in QosFilterParcelable filterParcel);
     void onQosCallbackUnregistered(int qosCallbackId);
     void onNetworkCreated();
-    void onNetworkDisconnected();
+    void onNetworkDestroyed();
 }
diff --git a/framework/src/android/net/INetworkAgentRegistry.aidl b/framework/src/android/net/INetworkAgentRegistry.aidl
index cbd6193..26cb1ed 100644
--- a/framework/src/android/net/INetworkAgentRegistry.aidl
+++ b/framework/src/android/net/INetworkAgentRegistry.aidl
@@ -41,4 +41,5 @@
     void sendNrQosSessionAvailable(int callbackId, in QosSession session, in NrQosSessionAttributes attributes);
     void sendQosSessionLost(int qosCallbackId, in QosSession session);
     void sendQosCallbackError(int qosCallbackId, int exceptionType);
+    void sendTeardownDelayMs(int teardownDelayMs);
 }
diff --git a/framework/src/android/net/NetworkAgent.java b/framework/src/android/net/NetworkAgent.java
index 6b55bb7..c57da53 100644
--- a/framework/src/android/net/NetworkAgent.java
+++ b/framework/src/android/net/NetworkAgent.java
@@ -185,6 +185,20 @@
     public static final int EVENT_UNDERLYING_NETWORKS_CHANGED = BASE + 5;
 
     /**
+     * Sent by the NetworkAgent to ConnectivityService to pass the current value of the teardown
+     * delay.
+     * arg1 = teardown delay in milliseconds
+     * @hide
+     */
+    public static final int EVENT_TEARDOWN_DELAY_CHANGED = BASE + 6;
+
+    /**
+     * The maximum value for the teardown delay, in milliseconds.
+     * @hide
+     */
+    public static final int MAX_TEARDOWN_DELAY_MS = 5000;
+
+    /**
      * Sent by ConnectivityService to the NetworkAgent to inform the agent of the
      * networks status - whether we could use the network or could not, due to
      * either a bad network configuration (no internet link) or captive portal.
@@ -197,7 +211,6 @@
      */
     public static final int CMD_REPORT_NETWORK_STATUS = BASE + 7;
 
-
     /**
      * Network validation suceeded.
      * Corresponds to {@link NetworkCapabilities.NET_CAPABILITY_VALIDATED}.
@@ -376,7 +389,7 @@
      *
      * @hide
      */
-    public static final int CMD_NETWORK_DISCONNECTED = BASE + 23;
+    public static final int CMD_NETWORK_DESTROYED = BASE + 23;
 
     private static NetworkInfo getLegacyNetworkInfo(final NetworkAgentConfig config) {
         final NetworkInfo ni = new NetworkInfo(config.legacyType, config.legacySubType,
@@ -581,8 +594,8 @@
                     onNetworkCreated();
                     break;
                 }
-                case CMD_NETWORK_DISCONNECTED: {
-                    onNetworkDisconnected();
+                case CMD_NETWORK_DESTROYED: {
+                    onNetworkDestroyed();
                     break;
                 }
             }
@@ -732,8 +745,8 @@
         }
 
         @Override
-        public void onNetworkDisconnected() {
-            mHandler.sendMessage(mHandler.obtainMessage(CMD_NETWORK_DISCONNECTED));
+        public void onNetworkDestroyed() {
+            mHandler.sendMessage(mHandler.obtainMessage(CMD_NETWORK_DESTROYED));
         }
     }
 
@@ -851,6 +864,29 @@
     }
 
     /**
+     * Sets the value of the teardown delay.
+     *
+     * The teardown delay is the time between when the network disconnects and when the native
+     * network corresponding to this {@code NetworkAgent} is destroyed. By default, the native
+     * network is destroyed immediately. If {@code teardownDelayMs} is non-zero, then when this
+     * network disconnects, the system will instead immediately mark the network as restricted
+     * and unavailable to unprivileged apps, but will defer destroying the native network until the
+     * teardown delay timer expires.
+     *
+     * The interfaces in use by this network will remain in use until the native network is
+     * destroyed and cannot be reused until {@link #onNetworkDestroyed()} is called.
+     *
+     * This method may be called at any time while the network is connected. It has no effect if
+     * the network is already disconnected and the teardown delay timer is running.
+     *
+     * @param teardownDelayMs the teardown delay to set, or 0 to disable teardown delay.
+     */
+    public void setTeardownDelayMs(
+            @IntRange(from = 0, to = MAX_TEARDOWN_DELAY_MS) int teardownDelayMs) {
+        queueOrSendMessage(reg -> reg.sendTeardownDelayMs(teardownDelayMs));
+    }
+
+    /**
      * Change the legacy subtype of this network agent.
      *
      * This is only for backward compatibility and should not be used by non-legacy network agents,
@@ -1053,7 +1089,7 @@
     /**
      * Called when ConnectivityService has successfully destroy this NetworkAgent's native network.
      */
-    public void onNetworkDisconnected() {}
+    public void onNetworkDestroyed() {}
 
     /**
      * Requests that the network hardware send the specified packet at the specified interval.
diff --git a/framework/src/android/net/NetworkRequest.java b/framework/src/android/net/NetworkRequest.java
index bcbc04f7..38691ef 100644
--- a/framework/src/android/net/NetworkRequest.java
+++ b/framework/src/android/net/NetworkRequest.java
@@ -699,4 +699,43 @@
     public int hashCode() {
         return Objects.hash(requestId, legacyType, networkCapabilities, type);
     }
+
+    /**
+     * Gets all the capabilities set on this {@code NetworkRequest} instance.
+     *
+     * @return an array of capability values for this instance.
+     */
+    @NonNull
+    public @NetCapability int[] getCapabilities() {
+        // No need to make a defensive copy here as NC#getCapabilities() already returns
+        // a new array.
+        return networkCapabilities.getCapabilities();
+    }
+
+    /**
+     * Gets all the unwanted capabilities set on this {@code NetworkRequest} instance.
+     *
+     * @return an array of unwanted capability values for this instance.
+     *
+     * @hide
+     */
+    @NonNull
+    @SystemApi(client = SystemApi.Client.MODULE_LIBRARIES)
+    public @NetCapability int[] getUnwantedCapabilities() {
+        // No need to make a defensive copy here as NC#getUnwantedCapabilities() already returns
+        // a new array.
+        return networkCapabilities.getUnwantedCapabilities();
+    }
+
+    /**
+     * Gets all the transports set on this {@code NetworkRequest} instance.
+     *
+     * @return an array of transport type values for this instance.
+     */
+    @NonNull
+    public @Transport int[] getTransportTypes() {
+        // No need to make a defensive copy here as NC#getTransportTypes() already returns
+        // a new array.
+        return networkCapabilities.getTransportTypes();
+    }
 }
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 77a2d79..7ea3738 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -1556,16 +1556,16 @@
         mNetworkInfoBlockingLogs.log(action + " " + uid);
     }
 
-    private void maybeLogBlockedStatusChanged(NetworkRequestInfo nri, Network net,
-            boolean blocked) {
+    private void maybeLogBlockedStatusChanged(NetworkRequestInfo nri, Network net, int blocked) {
         if (nri == null || net == null || !LOGD_BLOCKED_NETWORKINFO) {
             return;
         }
-        final String action = blocked ? "BLOCKED" : "UNBLOCKED";
+        final String action = (blocked != 0) ? "BLOCKED" : "UNBLOCKED";
         final int requestId = nri.getActiveRequest() != null
                 ? nri.getActiveRequest().requestId : nri.mRequests.get(0).requestId;
         mNetworkInfoBlockingLogs.log(String.format(
-                "%s %d(%d) on netId %d", action, nri.mAsUid, requestId, net.getNetId()));
+                "%s %d(%d) on netId %d: %s", action, nri.mAsUid, requestId, net.getNetId(),
+                blockedReasonsToString(blocked)));
     }
 
     /**
@@ -3127,6 +3127,13 @@
                     }
                     break;
                 }
+                case NetworkAgent.EVENT_TEARDOWN_DELAY_CHANGED: {
+                    if (msg.arg1 >= 0 && msg.arg1 <= NetworkAgent.MAX_TEARDOWN_DELAY_MS) {
+                        nai.teardownDelayMs = msg.arg1;
+                    } else {
+                        logwtf(nai.toShortString() + " set invalid teardown delay " + msg.arg1);
+                    }
+                }
             }
         }
 
@@ -3697,6 +3704,23 @@
         mLegacyTypeTracker.remove(nai, wasDefault);
         rematchAllNetworksAndRequests();
         mLingerMonitor.noteDisconnect(nai);
+
+        // Immediate teardown.
+        if (nai.teardownDelayMs == 0) {
+            destroyNetwork(nai);
+            return;
+        }
+
+        // Delayed teardown.
+        try {
+            mNetd.networkSetPermissionForNetwork(nai.network.netId, INetd.PERMISSION_SYSTEM);
+        } catch (RemoteException e) {
+            Log.d(TAG, "Error marking network restricted during teardown: " + e);
+        }
+        mHandler.postDelayed(() -> destroyNetwork(nai), nai.teardownDelayMs);
+    }
+
+    private void destroyNetwork(NetworkAgentInfo nai) {
         if (nai.created) {
             // Tell netd to clean up the configuration for this network
             // (routing rules, DNS, etc).
@@ -3709,7 +3733,7 @@
             mDnsManager.removeNetwork(nai.network);
         }
         mNetIdManager.releaseNetId(nai.network.getNetId());
-        nai.onNetworkDisconnected();
+        nai.onNetworkDestroyed();
     }
 
     private boolean createNativeNetwork(@NonNull NetworkAgentInfo networkAgent) {
@@ -7367,7 +7391,7 @@
                 break;
             }
             case ConnectivityManager.CALLBACK_BLK_CHANGED: {
-                maybeLogBlockedStatusChanged(nri, networkAgent.network, arg1 != 0);
+                maybeLogBlockedStatusChanged(nri, networkAgent.network, arg1);
                 msg.arg1 = arg1;
                 break;
             }
diff --git a/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java b/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java
index 97df5bf..ee32fbf 100644
--- a/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java
+++ b/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java
@@ -201,6 +201,9 @@
     // Set to true when partial connectivity was detected.
     public boolean partialConnectivity;
 
+    // Delay between when the network is disconnected and when the native network is destroyed.
+    public int teardownDelayMs;
+
     // Captive portal info of the network from RFC8908, if any.
     // Obtained by ConnectivityService and merged into NetworkAgent-provided information.
     public CaptivePortalData capportApiData;
@@ -589,13 +592,13 @@
     }
 
     /**
-     * Notify the NetworkAgent that the network is disconnected and destroyed.
+     * Notify the NetworkAgent that the native network has been destroyed.
      */
-    public void onNetworkDisconnected() {
+    public void onNetworkDestroyed() {
         try {
-            networkAgent.onNetworkDisconnected();
+            networkAgent.onNetworkDestroyed();
         } catch (RemoteException e) {
-            Log.e(TAG, "Error sending network disconnected event", e);
+            Log.e(TAG, "Error sending network destroyed event", e);
         }
     }
 
@@ -675,6 +678,12 @@
                 @QosCallbackException.ExceptionType final int exceptionType) {
             mQosCallbackTracker.sendEventQosCallbackError(qosCallbackId, exceptionType);
         }
+
+        @Override
+        public void sendTeardownDelayMs(int teardownDelayMs) {
+            mHandler.obtainMessage(NetworkAgent.EVENT_TEARDOWN_DELAY_CHANGED,
+                    teardownDelayMs, 0, new Pair<>(NetworkAgentInfo.this, null)).sendToTarget();
+        }
     }
 
     /**
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index e38773f..14c72f1 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -32,6 +32,7 @@
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 import static android.net.ConnectivityManager.ACTION_CAPTIVE_PORTAL_SIGN_IN;
 import static android.net.ConnectivityManager.BLOCKED_METERED_REASON_DATA_SAVER;
+import static android.net.ConnectivityManager.BLOCKED_METERED_REASON_MASK;
 import static android.net.ConnectivityManager.BLOCKED_METERED_REASON_USER_RESTRICTED;
 import static android.net.ConnectivityManager.BLOCKED_REASON_BATTERY_SAVER;
 import static android.net.ConnectivityManager.BLOCKED_REASON_NONE;
@@ -716,6 +717,9 @@
         private int mProbesSucceeded;
         private String mNmValidationRedirectUrl = null;
         private boolean mNmProvNotificationRequested = false;
+        private Runnable mCreatedCallback;
+        private Runnable mUnwantedCallback;
+        private Runnable mDisconnectedCallback;
 
         private final ConditionVariable mNetworkStatusReceived = new ConditionVariable();
         // Contains the redirectUrl from networkStatus(). Before reading, wait for
@@ -770,6 +774,24 @@
                     mRedirectUrl = redirectUrl;
                     mNetworkStatusReceived.open();
                 }
+
+                @Override
+                public void onNetworkCreated() {
+                    super.onNetworkCreated();
+                    if (mCreatedCallback != null) mCreatedCallback.run();
+                }
+
+                @Override
+                public void onNetworkUnwanted() {
+                    super.onNetworkUnwanted();
+                    if (mUnwantedCallback != null) mUnwantedCallback.run();
+                }
+
+                @Override
+                public void onNetworkDestroyed() {
+                    super.onNetworkDestroyed();
+                    if (mDisconnectedCallback != null) mDisconnectedCallback.run();
+                }
             };
 
             assertEquals(na.getNetwork().netId, nmNetworkCaptor.getValue().netId);
@@ -971,6 +993,18 @@
             p.timestampMillis = DATA_STALL_TIMESTAMP;
             mNmCallbacks.notifyDataStallSuspected(p);
         }
+
+        public void setCreatedCallback(Runnable r) {
+            mCreatedCallback = r;
+        }
+
+        public void setUnwantedCallback(Runnable r) {
+            mUnwantedCallback = r;
+        }
+
+        public void setDisconnectedCallback(Runnable r) {
+            mDisconnectedCallback = r;
+        }
     }
 
     /**
@@ -1374,10 +1408,21 @@
     }
 
     private void mockUidNetworkingBlocked() {
-        doAnswer(i -> NetworkPolicyManager.isUidBlocked(mBlockedReasons, i.getArgument(1))
+        doAnswer(i -> isUidBlocked(mBlockedReasons, i.getArgument(1))
         ).when(mNetworkPolicyManager).isUidNetworkingBlocked(anyInt(), anyBoolean());
     }
 
+    private boolean isUidBlocked(int blockedReasons, boolean meteredNetwork) {
+        final int blockedOnAllNetworksReason = (blockedReasons & ~BLOCKED_METERED_REASON_MASK);
+        if (blockedOnAllNetworksReason != BLOCKED_REASON_NONE) {
+            return true;
+        }
+        if (meteredNetwork) {
+            return blockedReasons != BLOCKED_REASON_NONE;
+        }
+        return false;
+    }
+
     private void setBlockedReasonChanged(int blockedReasons) {
         mBlockedReasons = blockedReasons;
         mPolicyCallback.onUidBlockedReasonChanged(Process.myUid(), blockedReasons);
@@ -2799,6 +2844,94 @@
     }
 
     @Test
+    public void testNetworkAgentCallbacks() throws Exception {
+        // Keeps track of the order of events that happen in this test.
+        final LinkedBlockingQueue<String> eventOrder = new LinkedBlockingQueue<>();
+
+        final NetworkRequest request = new NetworkRequest.Builder()
+                .addTransportType(TRANSPORT_WIFI).build();
+        final TestNetworkCallback callback = new TestNetworkCallback();
+        final AtomicReference<Network> wifiNetwork = new AtomicReference<>();
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
+
+        // Expectations for state when various callbacks fire. These expectations run on the handler
+        // thread and not on the test thread because they need to prevent the handler thread from
+        // advancing while they examine state.
+
+        // 1. When onCreated fires, netd has been told to create the network.
+        mWiFiNetworkAgent.setCreatedCallback(() -> {
+            eventOrder.offer("onNetworkCreated");
+            wifiNetwork.set(mWiFiNetworkAgent.getNetwork());
+            assertNotNull(wifiNetwork.get());
+            try {
+                verify(mMockNetd).networkCreatePhysical(wifiNetwork.get().getNetId(),
+                        INetd.PERMISSION_NONE);
+            } catch (RemoteException impossible) {
+                fail();
+            }
+        });
+
+        // 2. onNetworkUnwanted isn't precisely ordered with respect to any particular events. Just
+        //    check that it is fired at some point after disconnect.
+        mWiFiNetworkAgent.setUnwantedCallback(() -> eventOrder.offer("onNetworkUnwanted"));
+
+        // 3. While the teardown timer is running, connectivity APIs report the network is gone, but
+        //    netd has not yet been told to destroy it.
+        final Runnable duringTeardown = () -> {
+            eventOrder.offer("timePasses");
+            assertNull(mCm.getLinkProperties(wifiNetwork.get()));
+            try {
+                verify(mMockNetd, never()).networkDestroy(wifiNetwork.get().getNetId());
+            } catch (RemoteException impossible) {
+                fail();
+            }
+        };
+
+        // 4. After onNetworkDisconnected is called, connectivity APIs report the network is gone,
+        // and netd has been told to destroy it.
+        mWiFiNetworkAgent.setDisconnectedCallback(() -> {
+            eventOrder.offer("onNetworkDisconnected");
+            assertNull(mCm.getLinkProperties(wifiNetwork.get()));
+            try {
+                verify(mMockNetd).networkDestroy(wifiNetwork.get().getNetId());
+            } catch (RemoteException impossible) {
+                fail();
+            }
+        });
+
+        // Connect a network, and file a request for it after it has come up, to ensure the nascent
+        // timer is cleared and the test does not have to wait for it. Filing the request after the
+        // network has come up is necessary because ConnectivityService does not appear to clear the
+        // nascent timer if the first request satisfied by the network was filed before the network
+        // connected.
+        // TODO: fix this bug, file the request before connecting, and remove the waitForIdle.
+        mWiFiNetworkAgent.connectWithoutInternet();
+        waitForIdle();
+        mCm.requestNetwork(request, callback);
+        callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
+
+        // Set teardown delay and make sure CS has processed it.
+        mWiFiNetworkAgent.getNetworkAgent().setTeardownDelayMs(300);
+        waitForIdle();
+
+        // Post the duringTeardown lambda to the handler so it fires while teardown is in progress.
+        // The delay must be long enough it will run after the unregisterNetworkCallback has torn
+        // down the network and started the teardown timer, and short enough that the lambda is
+        // scheduled to run before the teardown timer.
+        final Handler h = new Handler(mCsHandlerThread.getLooper());
+        h.postDelayed(duringTeardown, 150);
+
+        // Disconnect the network and check that events happened in the right order.
+        mCm.unregisterNetworkCallback(callback);
+        assertEquals("onNetworkCreated", eventOrder.poll(TIMEOUT_MS, TimeUnit.MILLISECONDS));
+        assertEquals("onNetworkUnwanted", eventOrder.poll(TIMEOUT_MS, TimeUnit.MILLISECONDS));
+        assertEquals("timePasses", eventOrder.poll(TIMEOUT_MS, TimeUnit.MILLISECONDS));
+        assertEquals("onNetworkDisconnected", eventOrder.poll(TIMEOUT_MS, TimeUnit.MILLISECONDS));
+
+        mCm.unregisterNetworkCallback(callback);
+    }
+
+    @Test
     public void testExplicitlySelected() throws Exception {
         NetworkRequest request = new NetworkRequest.Builder()
                 .clearCapabilities().addCapability(NET_CAPABILITY_INTERNET)
diff --git a/tests/net/java/com/android/server/connectivity/IpConnectivityMetricsTest.java b/tests/net/java/com/android/server/connectivity/IpConnectivityMetricsTest.java
index 8c5d1d6..8b072c4 100644
--- a/tests/net/java/com/android/server/connectivity/IpConnectivityMetricsTest.java
+++ b/tests/net/java/com/android/server/connectivity/IpConnectivityMetricsTest.java
@@ -22,7 +22,9 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.content.Context;
@@ -56,6 +58,7 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -80,6 +83,12 @@
 
     IpConnectivityMetrics mService;
     NetdEventListenerService mNetdListener;
+    private static final NetworkCapabilities CAPABILITIES_WIFI = new NetworkCapabilities.Builder()
+            .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
+            .build();
+    private static final NetworkCapabilities CAPABILITIES_CELL = new NetworkCapabilities.Builder()
+            .addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR)
+            .build();
 
     @Before
     public void setUp() {
@@ -263,14 +272,6 @@
         // TODO: instead of comparing textpb to textpb, parse textpb and compare proto to proto.
         IpConnectivityLog logger = new IpConnectivityLog(mService.impl);
 
-        NetworkCapabilities ncWifi = new NetworkCapabilities();
-        NetworkCapabilities ncCell = new NetworkCapabilities();
-        ncWifi.addTransportType(NetworkCapabilities.TRANSPORT_WIFI);
-        ncCell.addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR);
-
-        when(mCm.getNetworkCapabilities(new Network(100))).thenReturn(ncWifi);
-        when(mCm.getNetworkCapabilities(new Network(101))).thenReturn(ncCell);
-
         ApfStats apfStats = new ApfStats.Builder()
                 .setDurationMs(45000)
                 .setReceivedRas(10)
@@ -584,11 +585,21 @@
         return buffer.toString();
     }
 
-    void connectEvent(int netid, int error, int latencyMs, String ipAddr) throws Exception {
-        mNetdListener.onConnectEvent(netid, error, latencyMs, ipAddr, 80, 1);
+    private void setCapabilities(int netId) {
+        final ArgumentCaptor<ConnectivityManager.NetworkCallback> networkCallback =
+                ArgumentCaptor.forClass(ConnectivityManager.NetworkCallback.class);
+        verify(mCm).registerNetworkCallback(any(), networkCallback.capture());
+        networkCallback.getValue().onCapabilitiesChanged(new Network(netId),
+                netId == 100 ? CAPABILITIES_WIFI : CAPABILITIES_CELL);
+    }
+
+    void connectEvent(int netId, int error, int latencyMs, String ipAddr) throws Exception {
+        setCapabilities(netId);
+        mNetdListener.onConnectEvent(netId, error, latencyMs, ipAddr, 80, 1);
     }
 
     void dnsEvent(int netId, int type, int result, int latency) throws Exception {
+        setCapabilities(netId);
         mNetdListener.onDnsEvent(netId, type, result, latency, "", null, 0, 0);
     }
 
diff --git a/tests/net/java/com/android/server/connectivity/NetdEventListenerServiceTest.java b/tests/net/java/com/android/server/connectivity/NetdEventListenerServiceTest.java
index 8ccea1a..50aaaee 100644
--- a/tests/net/java/com/android/server/connectivity/NetdEventListenerServiceTest.java
+++ b/tests/net/java/com/android/server/connectivity/NetdEventListenerServiceTest.java
@@ -23,8 +23,9 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.verify;
 
 import android.content.Context;
 import android.net.ConnectivityManager;
@@ -42,6 +43,7 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 
 import java.io.FileOutputStream;
 import java.io.PrintWriter;
@@ -61,18 +63,16 @@
 
     NetdEventListenerService mService;
     ConnectivityManager mCm;
+    private static final NetworkCapabilities CAPABILITIES_WIFI = new NetworkCapabilities.Builder()
+            .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
+            .build();
+    private static final NetworkCapabilities CAPABILITIES_CELL = new NetworkCapabilities.Builder()
+            .addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR)
+            .build();
 
     @Before
     public void setUp() {
-        NetworkCapabilities ncWifi = new NetworkCapabilities();
-        NetworkCapabilities ncCell = new NetworkCapabilities();
-        ncWifi.addTransportType(NetworkCapabilities.TRANSPORT_WIFI);
-        ncCell.addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR);
-
         mCm = mock(ConnectivityManager.class);
-        when(mCm.getNetworkCapabilities(new Network(100))).thenReturn(ncWifi);
-        when(mCm.getNetworkCapabilities(new Network(101))).thenReturn(ncCell);
-
         mService = new NetdEventListenerService(mCm);
     }
 
@@ -470,7 +470,16 @@
         assertEquals(want, got);
     }
 
+    private void setCapabilities(int netId) {
+        final ArgumentCaptor<ConnectivityManager.NetworkCallback> networkCallback =
+                ArgumentCaptor.forClass(ConnectivityManager.NetworkCallback.class);
+        verify(mCm).registerNetworkCallback(any(), networkCallback.capture());
+        networkCallback.getValue().onCapabilitiesChanged(new Network(netId),
+                netId == 100 ? CAPABILITIES_WIFI : CAPABILITIES_CELL);
+    }
+
     Thread connectEventAction(int netId, int error, int latencyMs, String ipAddr) {
+        setCapabilities(netId);
         return new Thread(() -> {
             try {
                 mService.onConnectEvent(netId, error, latencyMs, ipAddr, 80, 1);
@@ -481,6 +490,7 @@
     }
 
     void dnsEvent(int netId, int type, int result, int latency) throws Exception {
+        setCapabilities(netId);
         mService.onDnsEvent(netId, type, result, latency, "", null, 0, 0);
     }