diff --git a/Tethering/common/TetheringLib/src/android/net/TetheringConfigurationParcel.aidl b/Tethering/common/TetheringLib/src/android/net/TetheringConfigurationParcel.aidl
index 89f3813..168d7f9 100644
--- a/Tethering/common/TetheringLib/src/android/net/TetheringConfigurationParcel.aidl
+++ b/Tethering/common/TetheringLib/src/android/net/TetheringConfigurationParcel.aidl
@@ -21,17 +21,10 @@
  * @hide
  */
 parcelable TetheringConfigurationParcel {
-    int subId;
     String[] tetherableUsbRegexs;
     String[] tetherableWifiRegexs;
     String[] tetherableBluetoothRegexs;
-    boolean isDunRequired;
-    boolean chooseUpstreamAutomatically;
-    int[] preferredUpstreamIfaceTypes;
     String[] legacyDhcpRanges;
-    String[] defaultIPv4DNS;
-    boolean enableLegacyDhcpServer;
     String[] provisioningApp;
     String provisioningAppNoUi;
-    int provisioningCheckPeriod;
 }
diff --git a/Tethering/proguard.flags b/Tethering/proguard.flags
index 2905e28..109bbda 100644
--- a/Tethering/proguard.flags
+++ b/Tethering/proguard.flags
@@ -1,7 +1,10 @@
 # Keep class's integer static field for MessageUtils to parsing their name.
--keep class com.android.networkstack.tethering.Tethering$TetherMainSM {
-    static final int CMD_*;
-    static final int EVENT_*;
+-keepclassmembers class com.android.server.**,android.net.**,com.android.networkstack.** {
+    static final % POLICY_*;
+    static final % NOTIFY_TYPE_*;
+    static final % TRANSPORT_*;
+    static final % CMD_*;
+    static final % EVENT_*;
 }
 
 -keep class com.android.networkstack.tethering.util.BpfMap {
@@ -21,12 +24,8 @@
     *;
 }
 
--keepclassmembers class android.net.ip.IpServer {
-    static final int CMD_*;
-}
-
 # The lite proto runtime uses reflection to access fields based on the names in
 # the schema, keep all the fields.
 -keepclassmembers class * extends com.android.networkstack.tethering.protobuf.MessageLite {
     <fields>;
-}
\ No newline at end of file
+}
diff --git a/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java b/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java
index 74ba209..ebc9d26 100644
--- a/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java
+++ b/Tethering/src/com/android/networkstack/tethering/BpfCoordinator.java
@@ -1473,6 +1473,15 @@
             // to Objects.hash() to avoid autoboxing overhead.
             return Objects.hash(upstreamIfindex, downstreamIfindex, address, srcMac, dstMac);
         }
+
+        @Override
+        public String toString() {
+            return "upstreamIfindex: " + upstreamIfindex
+                    + ", downstreamIfindex: " + downstreamIfindex
+                    + ", address: " + address.getHostAddress()
+                    + ", srcMac: " + srcMac
+                    + ", dstMac: " + dstMac;
+        }
     }
 
     /** Tethering client information class. */
diff --git a/Tethering/src/com/android/networkstack/tethering/TetheringConfiguration.java b/Tethering/src/com/android/networkstack/tethering/TetheringConfiguration.java
index 696a970..8ef5d71 100644
--- a/Tethering/src/com/android/networkstack/tethering/TetheringConfiguration.java
+++ b/Tethering/src/com/android/networkstack/tethering/TetheringConfiguration.java
@@ -652,21 +652,13 @@
      */
     public TetheringConfigurationParcel toStableParcelable() {
         final TetheringConfigurationParcel parcel = new TetheringConfigurationParcel();
-        parcel.subId = activeDataSubId;
         parcel.tetherableUsbRegexs = tetherableUsbRegexs;
         parcel.tetherableWifiRegexs = tetherableWifiRegexs;
         parcel.tetherableBluetoothRegexs = tetherableBluetoothRegexs;
-        parcel.isDunRequired = isDunRequired;
-        parcel.chooseUpstreamAutomatically = chooseUpstreamAutomatically;
-
-        parcel.preferredUpstreamIfaceTypes = toIntArray(preferredUpstreamIfaceTypes);
-
         parcel.legacyDhcpRanges = legacyDhcpRanges;
-        parcel.defaultIPv4DNS = defaultIPv4DNS;
-        parcel.enableLegacyDhcpServer = mEnableLegacyDhcpServer;
         parcel.provisioningApp = provisioningApp;
         parcel.provisioningAppNoUi = provisioningAppNoUi;
-        parcel.provisioningCheckPeriod = provisioningCheckPeriod;
+
         return parcel;
     }
 }
diff --git a/Tethering/src/com/android/networkstack/tethering/metrics/TetheringMetrics.java b/Tethering/src/com/android/networkstack/tethering/metrics/TetheringMetrics.java
index d8e631e..2c6054d 100644
--- a/Tethering/src/com/android/networkstack/tethering/metrics/TetheringMetrics.java
+++ b/Tethering/src/com/android/networkstack/tethering/metrics/TetheringMetrics.java
@@ -75,6 +75,8 @@
                     .setUserType(userTypeToEnum(callerPkg))
                     .setUpstreamType(UpstreamType.UT_UNKNOWN)
                     .setErrorCode(ErrorCode.EC_NO_ERROR)
+                    .setUpstreamEvents(UpstreamEvents.newBuilder())
+                    .setDurationMillis(0)
                     .build();
         mBuilderMap.put(downstreamType, statsBuilder);
     }
@@ -110,7 +112,8 @@
                 reported.getErrorCode().getNumber(),
                 reported.getDownstreamType().getNumber(),
                 reported.getUpstreamType().getNumber(),
-                reported.getUserType().getNumber());
+                reported.getUserType().getNumber(),
+                null, 0);
         if (DBG) {
             Log.d(TAG, "Write errorCode: " + reported.getErrorCode().getNumber()
                     + ", downstreamType: " + reported.getDownstreamType().getNumber()
diff --git a/Tethering/src/com/android/networkstack/tethering/metrics/stats.proto b/Tethering/src/com/android/networkstack/tethering/metrics/stats.proto
index 46a47af..27f2126 100644
--- a/Tethering/src/com/android/networkstack/tethering/metrics/stats.proto
+++ b/Tethering/src/com/android/networkstack/tethering/metrics/stats.proto
@@ -21,12 +21,38 @@
 
 import "frameworks/proto_logging/stats/enums/stats/connectivity/tethering.proto";
 
+// Logs each upstream for a successful switch over
+message UpstreamEvent {
+  // Transport type of upstream network
+  optional .android.stats.connectivity.UpstreamType upstream_type = 1;
+
+  // A time period that an upstream continued
+  optional int64 duration_millis = 2;
+}
+
+message UpstreamEvents {
+  repeated UpstreamEvent upstream_event = 1;
+}
+
 /**
  * Logs Tethering events
  */
 message NetworkTetheringReported {
-   optional .android.stats.connectivity.ErrorCode error_code = 1;
-   optional .android.stats.connectivity.DownstreamType downstream_type = 2;
-   optional .android.stats.connectivity.UpstreamType upstream_type = 3;
-   optional .android.stats.connectivity.UserType user_type = 4;
+    // Tethering error code
+    optional .android.stats.connectivity.ErrorCode error_code = 1;
+
+    // Tethering downstream type
+    optional .android.stats.connectivity.DownstreamType downstream_type = 2;
+
+    // Transport type of upstream network
+    optional .android.stats.connectivity.UpstreamType upstream_type = 3  [deprecated = true];
+
+    // The user type of switching tethering
+    optional .android.stats.connectivity.UserType user_type = 4;
+
+    // Log each transport type of upstream network event
+    optional UpstreamEvents upstream_events = 5;
+
+    // A time period that a downstreams exists
+    optional int64 duration_millis = 6;
 }
diff --git a/Tethering/tests/unit/src/com/android/networkstack/tethering/BpfCoordinatorTest.java b/Tethering/tests/unit/src/com/android/networkstack/tethering/BpfCoordinatorTest.java
index ac92b43..bbca565 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/BpfCoordinatorTest.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/BpfCoordinatorTest.java
@@ -2092,4 +2092,11 @@
         assertNull(mBpfUpstream4Map.getValue(upstream4KeyB));
         assertNull(mBpfDownstream4Map.getValue(downstream4KeyB));
     }
+
+    @Test
+    public void testIpv6ForwardingRuleToString() throws Exception {
+        final Ipv6ForwardingRule rule = buildTestForwardingRule(UPSTREAM_IFINDEX, NEIGH_A, MAC_A);
+        assertEquals("upstreamIfindex: 1001, downstreamIfindex: 1003, address: 2001:db8::1, "
+                + "srcMac: 12:34:56:78:90:ab, dstMac: 00:00:00:00:00:0a", rule.toString());
+    }
 }
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 a468d82..9337b1a 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java
@@ -1855,20 +1855,12 @@
 
         private void assertTetherConfigParcelEqual(@NonNull TetheringConfigurationParcel actual,
                 @NonNull TetheringConfigurationParcel expect) {
-            assertEquals(actual.subId, expect.subId);
             assertArrayEquals(actual.tetherableUsbRegexs, expect.tetherableUsbRegexs);
             assertArrayEquals(actual.tetherableWifiRegexs, expect.tetherableWifiRegexs);
             assertArrayEquals(actual.tetherableBluetoothRegexs, expect.tetherableBluetoothRegexs);
-            assertEquals(actual.isDunRequired, expect.isDunRequired);
-            assertEquals(actual.chooseUpstreamAutomatically, expect.chooseUpstreamAutomatically);
-            assertArrayEquals(actual.preferredUpstreamIfaceTypes,
-                    expect.preferredUpstreamIfaceTypes);
             assertArrayEquals(actual.legacyDhcpRanges, expect.legacyDhcpRanges);
-            assertArrayEquals(actual.defaultIPv4DNS, expect.defaultIPv4DNS);
-            assertEquals(actual.enableLegacyDhcpServer, expect.enableLegacyDhcpServer);
             assertArrayEquals(actual.provisioningApp, expect.provisioningApp);
             assertEquals(actual.provisioningAppNoUi, expect.provisioningAppNoUi);
-            assertEquals(actual.provisioningCheckPeriod, expect.provisioningCheckPeriod);
         }
     }
 
diff --git a/Tethering/tests/unit/src/com/android/networkstack/tethering/metrics/TetheringMetricsTest.java b/Tethering/tests/unit/src/com/android/networkstack/tethering/metrics/TetheringMetricsTest.java
index 6a85718..7fdde97 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/metrics/TetheringMetricsTest.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/metrics/TetheringMetricsTest.java
@@ -88,6 +88,8 @@
                 .setUserType(user)
                 .setUpstreamType(UpstreamType.UT_UNKNOWN)
                 .setErrorCode(error)
+                .setUpstreamEvents(UpstreamEvents.newBuilder())
+                .setDurationMillis(0)
                 .build();
         verify(mTetheringMetrics).write(expectedReport);
     }
diff --git a/bpf_progs/bpf_shared.h b/bpf_progs/bpf_shared.h
index 85b9f86..7b1106a 100644
--- a/bpf_progs/bpf_shared.h
+++ b/bpf_progs/bpf_shared.h
@@ -95,13 +95,13 @@
 
 // 'static' - otherwise these constants end up in .rodata in the resulting .o post compilation
 static const int COOKIE_UID_MAP_SIZE = 10000;
-static const int UID_COUNTERSET_MAP_SIZE = 2000;
+static const int UID_COUNTERSET_MAP_SIZE = 4000;
 static const int APP_STATS_MAP_SIZE = 10000;
 static const int STATS_MAP_SIZE = 5000;
 static const int IFACE_INDEX_NAME_MAP_SIZE = 1000;
 static const int IFACE_STATS_MAP_SIZE = 1000;
 static const int CONFIGURATION_MAP_SIZE = 2;
-static const int UID_OWNER_MAP_SIZE = 2000;
+static const int UID_OWNER_MAP_SIZE = 4000;
 
 #ifdef __cplusplus
 
diff --git a/bpf_progs/offload.c b/bpf_progs/offload.c
index bb9fc34..c7b444d 100644
--- a/bpf_progs/offload.c
+++ b/bpf_progs/offload.c
@@ -91,8 +91,11 @@
 
 // ----- Tethering Error Counters -----
 
-DEFINE_BPF_MAP_GRW(tether_error_map, ARRAY, uint32_t, uint32_t, BPF_TETHER_ERR__MAX,
-                   TETHERING_GID)
+// Note that pre-T devices with Mediatek chipsets may have a kernel bug (bad patch
+// "[ALPS05162612] bpf: fix ubsan error") making it impossible to write to non-zero
+// offset of bpf map ARRAYs.  This file (offload.o) loads on T, but luckily this
+// array is only written by bpf code, and only read by userspace.
+DEFINE_BPF_MAP_RO(tether_error_map, ARRAY, uint32_t, uint32_t, BPF_TETHER_ERR__MAX, TETHERING_GID)
 
 #define COUNT_AND_RETURN(counter, ret) do {                     \
     uint32_t code = BPF_TETHER_ERR_ ## counter;                 \
diff --git a/bpf_progs/test.c b/bpf_progs/test.c
index d42205f..c11c358 100644
--- a/bpf_progs/test.c
+++ b/bpf_progs/test.c
@@ -40,6 +40,10 @@
 #define TETHERING_GID AID_NETWORK_STACK
 #endif
 
+// This is non production code, only used for testing
+// Needed because the bitmap array definition is non-kosher for pre-T OS devices.
+#define THIS_BPF_PROGRAM_IS_FOR_TEST_PURPOSES_ONLY
+
 #include "bpf_helpers.h"
 #include "bpf_net_helpers.h"
 #include "bpf_tethering.h"
diff --git a/framework-t/src/android/net/NetworkIdentity.java b/framework-t/src/android/net/NetworkIdentity.java
index 350ed86..48e5092 100644
--- a/framework-t/src/android/net/NetworkIdentity.java
+++ b/framework-t/src/android/net/NetworkIdentity.java
@@ -34,8 +34,8 @@
 import android.telephony.TelephonyManager;
 import android.util.proto.ProtoOutputStream;
 
+import com.android.net.module.util.BitUtils;
 import com.android.net.module.util.CollectionUtils;
-import com.android.net.module.util.NetworkCapabilitiesUtils;
 import com.android.net.module.util.NetworkIdentityUtils;
 
 import java.lang.annotation.Retention;
@@ -172,7 +172,7 @@
         if (oemManaged == OEM_NONE) {
             return "OEM_NONE";
         }
-        final int[] bitPositions = NetworkCapabilitiesUtils.unpackBits(oemManaged);
+        final int[] bitPositions = BitUtils.unpackBits(oemManaged);
         final ArrayList<String> oemManagedNames = new ArrayList<String>();
         for (int position : bitPositions) {
             oemManagedNames.add(nameOfOemManaged(1 << position));
diff --git a/framework/src/android/net/NetworkCapabilities.java b/framework/src/android/net/NetworkCapabilities.java
index d0cbbe5..e07601f 100644
--- a/framework/src/android/net/NetworkCapabilities.java
+++ b/framework/src/android/net/NetworkCapabilities.java
@@ -17,6 +17,7 @@
 package android.net;
 
 import static com.android.internal.annotations.VisibleForTesting.Visibility.PRIVATE;
+import static com.android.net.module.util.BitUtils.appendStringRepresentationOfBitMaskToStringBuilder;
 
 import android.annotation.IntDef;
 import android.annotation.LongDef;
@@ -36,6 +37,7 @@
 import android.util.Range;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.net.module.util.BitUtils;
 import com.android.net.module.util.CollectionUtils;
 import com.android.net.module.util.NetworkCapabilitiesUtils;
 
@@ -217,7 +219,7 @@
         if (hasCapability(NET_CAPABILITY_ENTERPRISE) && mEnterpriseId == 0) {
             return new int[]{NET_ENTERPRISE_ID_1};
         }
-        return NetworkCapabilitiesUtils.unpackBits(mEnterpriseId);
+        return BitUtils.unpackBits(mEnterpriseId);
     }
 
     /**
@@ -643,7 +645,7 @@
      * Network capabilities that are expected to be mutable, i.e., can change while a particular
      * network is connected.
      */
-    private static final long MUTABLE_CAPABILITIES = NetworkCapabilitiesUtils.packBitList(
+    private static final long MUTABLE_CAPABILITIES = BitUtils.packBitList(
             // TRUSTED can change when user explicitly connects to an untrusted network in Settings.
             // http://b/18206275
             NET_CAPABILITY_TRUSTED,
@@ -678,7 +680,7 @@
     /**
      * Capabilities that are set by default when the object is constructed.
      */
-    private static final long DEFAULT_CAPABILITIES = NetworkCapabilitiesUtils.packBitList(
+    private static final long DEFAULT_CAPABILITIES = BitUtils.packBitList(
             NET_CAPABILITY_NOT_RESTRICTED,
             NET_CAPABILITY_TRUSTED,
             NET_CAPABILITY_NOT_VPN);
@@ -687,7 +689,7 @@
      * Capabilities that are managed by ConnectivityService.
      */
     private static final long CONNECTIVITY_MANAGED_CAPABILITIES =
-            NetworkCapabilitiesUtils.packBitList(
+            BitUtils.packBitList(
                     NET_CAPABILITY_VALIDATED,
                     NET_CAPABILITY_CAPTIVE_PORTAL,
                     NET_CAPABILITY_FOREGROUND,
@@ -700,7 +702,7 @@
      * INTERNET, IMS, SUPL, etc.
      */
     private static final long TEST_NETWORKS_ALLOWED_CAPABILITIES =
-            NetworkCapabilitiesUtils.packBitList(
+            BitUtils.packBitList(
             NET_CAPABILITY_NOT_METERED,
             NET_CAPABILITY_TEMPORARILY_NOT_METERED,
             NET_CAPABILITY_NOT_RESTRICTED,
@@ -800,7 +802,7 @@
      * @return an array of capability values for this instance.
      */
     public @NonNull @NetCapability int[] getCapabilities() {
-        return NetworkCapabilitiesUtils.unpackBits(mNetworkCapabilities);
+        return BitUtils.unpackBits(mNetworkCapabilities);
     }
 
     /**
@@ -810,7 +812,7 @@
      * @hide
      */
     public @NetCapability int[] getForbiddenCapabilities() {
-        return NetworkCapabilitiesUtils.unpackBits(mForbiddenNetworkCapabilities);
+        return BitUtils.unpackBits(mForbiddenNetworkCapabilities);
     }
 
 
@@ -822,8 +824,8 @@
      */
     public void setCapabilities(@NetCapability int[] capabilities,
             @NetCapability int[] forbiddenCapabilities) {
-        mNetworkCapabilities = NetworkCapabilitiesUtils.packBits(capabilities);
-        mForbiddenNetworkCapabilities = NetworkCapabilitiesUtils.packBits(forbiddenCapabilities);
+        mNetworkCapabilities = BitUtils.packBits(capabilities);
+        mForbiddenNetworkCapabilities = BitUtils.packBits(forbiddenCapabilities);
     }
 
     /**
@@ -960,7 +962,7 @@
                 & NON_REQUESTABLE_CAPABILITIES;
 
         if (nonRequestable != 0) {
-            return capabilityNameOf(NetworkCapabilitiesUtils.unpackBits(nonRequestable)[0]);
+            return capabilityNameOf(BitUtils.unpackBits(nonRequestable)[0]);
         }
         if (mLinkUpBandwidthKbps != 0 || mLinkDownBandwidthKbps != 0) return "link bandwidth";
         if (hasSignalStrength()) return "signalStrength";
@@ -1193,7 +1195,7 @@
      * Allowed transports on an unrestricted test network (in addition to TRANSPORT_TEST).
      */
     private static final long UNRESTRICTED_TEST_NETWORKS_ALLOWED_TRANSPORTS =
-            NetworkCapabilitiesUtils.packBitList(
+            BitUtils.packBitList(
                     TRANSPORT_TEST,
                     // Test eth networks are created with EthernetManager#setIncludeTestInterfaces
                     TRANSPORT_ETHERNET,
@@ -1258,7 +1260,7 @@
      */
     @SystemApi
     @NonNull public @Transport int[] getTransportTypes() {
-        return NetworkCapabilitiesUtils.unpackBits(mTransportTypes);
+        return BitUtils.unpackBits(mTransportTypes);
     }
 
     /**
@@ -1268,7 +1270,7 @@
      * @hide
      */
     public void setTransportTypes(@Transport int[] transportTypes) {
-        mTransportTypes = NetworkCapabilitiesUtils.packBits(transportTypes);
+        mTransportTypes = BitUtils.packBits(transportTypes);
     }
 
     /**
@@ -2041,9 +2043,9 @@
         long oldImmutableCapabilities = this.mNetworkCapabilities & mask;
         long newImmutableCapabilities = that.mNetworkCapabilities & mask;
         if (oldImmutableCapabilities != newImmutableCapabilities) {
-            String before = capabilityNamesOf(NetworkCapabilitiesUtils.unpackBits(
+            String before = capabilityNamesOf(BitUtils.unpackBits(
                     oldImmutableCapabilities));
-            String after = capabilityNamesOf(NetworkCapabilitiesUtils.unpackBits(
+            String after = capabilityNamesOf(BitUtils.unpackBits(
                     newImmutableCapabilities));
             joiner.add(String.format("immutable capabilities changed: %s -> %s", before, after));
         }
@@ -2064,6 +2066,36 @@
     }
 
     /**
+     * Returns a short but human-readable string of updates from an older set of capabilities.
+     * @param old the old capabilities to diff from
+     * @return a string fit for logging differences, or null if no differences.
+     *         this never returns the empty string.
+     * @hide
+     */
+    @Nullable
+    public String describeCapsDifferencesFrom(@Nullable final NetworkCapabilities old) {
+        final long oldCaps = null == old ? 0 : old.mNetworkCapabilities;
+        final long changed = oldCaps ^ mNetworkCapabilities;
+        if (0 == changed) return null;
+        // If the control reaches here, there are changes (additions, removals, or both) so
+        // the code below is guaranteed to add something to the string and can't return "".
+        final long removed = oldCaps & changed;
+        final long added = mNetworkCapabilities & changed;
+        final StringBuilder sb = new StringBuilder();
+        if (0 != removed) {
+            sb.append("-");
+            appendStringRepresentationOfBitMaskToStringBuilder(sb, removed,
+                    NetworkCapabilities::capabilityNameOf, "-");
+        }
+        if (0 != added) {
+            sb.append("+");
+            appendStringRepresentationOfBitMaskToStringBuilder(sb, added,
+                    NetworkCapabilities::capabilityNameOf, "+");
+        }
+        return sb.toString();
+    }
+
+    /**
      * Checks that our requestable capabilities are the same as those of the given
      * {@code NetworkCapabilities}.
      *
@@ -2313,32 +2345,6 @@
         return sb.toString();
     }
 
-
-    private interface NameOf {
-        String nameOf(int value);
-    }
-
-    /**
-     * @hide
-     */
-    public static void appendStringRepresentationOfBitMaskToStringBuilder(@NonNull StringBuilder sb,
-            long bitMask, @NonNull NameOf nameFetcher, @NonNull String separator) {
-        int bitPos = 0;
-        boolean firstElementAdded = false;
-        while (bitMask != 0) {
-            if ((bitMask & 1) != 0) {
-                if (firstElementAdded) {
-                    sb.append(separator);
-                } else {
-                    firstElementAdded = true;
-                }
-                sb.append(nameFetcher.nameOf(bitPos));
-            }
-            bitMask >>= 1;
-            ++bitPos;
-        }
-    }
-
     /**
      * @hide
      */
diff --git a/nearby/service/Android.bp b/nearby/service/Android.bp
index d318a80..6065f7f 100644
--- a/nearby/service/Android.bp
+++ b/nearby/service/Android.bp
@@ -85,10 +85,10 @@
     ],
     libs: [
         "androidx.annotation_annotation",
-        "framework-bluetooth.stubs.module_lib", // TODO(b/215722418): Change to framework-bluetooth once fixed
+        "framework-bluetooth",
         "error_prone_annotations",
         "framework-connectivity-t.impl",
-        "framework-statsd.stubs.module_lib",
+        "framework-statsd",
     ],
     static_libs: [
         "androidx.core_core",
diff --git a/service-t/Android.bp b/service-t/Android.bp
index 1b9f2ec..9bf9135 100644
--- a/service-t/Android.bp
+++ b/service-t/Android.bp
@@ -49,7 +49,7 @@
         "framework-annotations-lib",
         "framework-connectivity-pre-jarjar",
         "framework-connectivity-t-pre-jarjar",
-        "framework-tethering.stubs.module_lib",
+        "framework-tethering",
         "service-connectivity-pre-jarjar",
         "service-nearby-pre-jarjar",
         "ServiceConnectivityResources",
diff --git a/service-t/src/com/android/server/net/NetworkStatsService.java b/service-t/src/com/android/server/net/NetworkStatsService.java
index 629bf73..cf53002 100644
--- a/service-t/src/com/android/server/net/NetworkStatsService.java
+++ b/service-t/src/com/android/server/net/NetworkStatsService.java
@@ -430,12 +430,7 @@
     private long mLastStatsSessionPoll;
 
     private final Object mOpenSessionCallsLock = new Object();
-    /**
-     * Map from UID to number of opened sessions. This is used for rate-limt an app to open
-     * session frequently
-     */
-    @GuardedBy("mOpenSessionCallsLock")
-    private final SparseIntArray mOpenSessionCallsPerUid = new SparseIntArray();
+
     /**
      * Map from key {@code OpenSessionKey} to count of opened sessions. This is for recording
      * the caller of open session and it is only for debugging.
@@ -1361,9 +1356,6 @@
                 mOpenSessionCallsPerCaller.put(key, Integer.sum(callsPerCaller, 1));
             }
 
-            int callsPerUid = mOpenSessionCallsPerUid.get(key.uid, 0);
-            mOpenSessionCallsPerUid.put(key.uid, callsPerUid + 1);
-
             if (key.uid == android.os.Process.SYSTEM_UID) {
                 return false;
             }
@@ -2514,8 +2506,7 @@
         for (int uid : uids) {
             deleteKernelTagData(uid);
         }
-        // TODO: Remove the UID's entries from mOpenSessionCallsPerUid and
-        // mOpenSessionCallsPerCaller
+        // TODO: Remove the UID's entries from mOpenSessionCallsPerCaller.
     }
 
     /**
@@ -2959,7 +2950,7 @@
      */
     private NetworkStats getNetworkStatsUidDetail(String[] ifaces)
             throws RemoteException {
-        final NetworkStats uidSnapshot = readNetworkStatsUidDetail(UID_ALL,  ifaces, TAG_ALL);
+        final NetworkStats uidSnapshot = readNetworkStatsUidDetail(UID_ALL, ifaces, TAG_ALL);
 
         // fold tethering stats and operations into uid snapshot
         final NetworkStats tetherSnapshot = getNetworkStatsTethering(STATS_PER_UID);
@@ -2974,6 +2965,7 @@
         uidSnapshot.combineAllValues(providerStats);
 
         uidSnapshot.combineAllValues(mUidOperations);
+        uidSnapshot.filter(UID_ALL, ifaces, TAG_ALL);
 
         return uidSnapshot;
     }
diff --git a/service/Android.bp b/service/Android.bp
index 0a00362..d850015 100644
--- a/service/Android.bp
+++ b/service/Android.bp
@@ -148,12 +148,20 @@
     libs: [
         "framework-annotations-lib",
         "framework-connectivity-pre-jarjar",
+        // The framework-connectivity-t library is only available on T+ platforms
+        // so any calls to it must be protected with a check to ensure that it is
+        // available. The linter will detect any unprotected calls through an API
+        // but not direct calls to the implementation. So, this depends on the
+        // module lib stubs directly to ensure the linter will work correctly
+        // as depending on framework-connectivity-t would cause it to be compiled
+        // against the implementation because the two libraries are in the same
+        // APEX.
         "framework-connectivity-t.stubs.module_lib",
-        "framework-tethering.stubs.module_lib",
-        "framework-wifi.stubs.module_lib",
+        "framework-tethering",
+        "framework-wifi",
         "unsupportedappusage",
         "ServiceConnectivityResources",
-        "framework-statsd.stubs.module_lib",
+        "framework-statsd",
     ],
     static_libs: [
         // Do not add libs here if they are already included
@@ -194,7 +202,7 @@
     libs: [
         "framework-annotations-lib",
         "framework-connectivity-pre-jarjar",
-        "framework-wifi.stubs.module_lib",
+        "framework-wifi",
         "service-connectivity-pre-jarjar",
     ],
     visibility: [
@@ -216,7 +224,9 @@
     apex_available: [
         "com.android.tethering",
     ],
-    lint: { strict_updatability_linting: true },
+    lint: {
+        strict_updatability_linting: true,
+    },
 }
 
 java_defaults {
@@ -246,8 +256,8 @@
         "framework-annotations-lib",
         "framework-connectivity.impl",
         "framework-connectivity-t.impl",
-        "framework-tethering.stubs.module_lib",
-        "framework-wifi.stubs.module_lib",
+        "framework-tethering",
+        "framework-wifi",
         "libprotobuf-java-nano",
     ],
     jarjar_rules: ":connectivity-jarjar-rules",
@@ -257,7 +267,9 @@
     optimize: {
         proguard_flags_files: ["proguard.flags"],
     },
-    lint: { strict_updatability_linting: true },
+    lint: {
+        strict_updatability_linting: true,
+    },
 }
 
 // A special library created strictly for use by the tests as they need the
@@ -337,8 +349,8 @@
 }
 
 genrule {
-  name: "statslog-connectivity-java-gen",
-  tools: ["stats-log-api-gen"],
-  cmd: "$(location stats-log-api-gen) --java $(out) --module connectivity --javaPackage com.android.server --javaClass ConnectivityStatsLog",
-  out: ["com/android/server/ConnectivityStatsLog.java"],
+    name: "statslog-connectivity-java-gen",
+    tools: ["stats-log-api-gen"],
+    cmd: "$(location stats-log-api-gen) --java $(out) --module connectivity --javaPackage com.android.server --javaClass ConnectivityStatsLog",
+    out: ["com/android/server/ConnectivityStatsLog.java"],
 }
diff --git a/service/mdns/com/android/server/connectivity/mdns/MdnsConfigs.java b/service/mdns/com/android/server/connectivity/mdns/MdnsConfigs.java
index 922037b..35a685d 100644
--- a/service/mdns/com/android/server/connectivity/mdns/MdnsConfigs.java
+++ b/service/mdns/com/android/server/connectivity/mdns/MdnsConfigs.java
@@ -89,4 +89,12 @@
     public static boolean preferIpv6() {
         return false;
     }
+
+    public static boolean removeServiceAfterTtlExpires() {
+        return false;
+    }
+
+    public static boolean allowSearchOptionsToRemoveExpiredService() {
+        return false;
+    }
 }
\ No newline at end of file
diff --git a/service/mdns/com/android/server/connectivity/mdns/MdnsSearchOptions.java b/service/mdns/com/android/server/connectivity/mdns/MdnsSearchOptions.java
index 6e90d2c..195bc8e 100644
--- a/service/mdns/com/android/server/connectivity/mdns/MdnsSearchOptions.java
+++ b/service/mdns/com/android/server/connectivity/mdns/MdnsSearchOptions.java
@@ -43,7 +43,7 @@
                 @Override
                 public MdnsSearchOptions createFromParcel(Parcel source) {
                     return new MdnsSearchOptions(source.createStringArrayList(),
-                            source.readBoolean());
+                            source.readBoolean(), source.readBoolean());
                 }
 
                 @Override
@@ -55,14 +55,16 @@
     private final List<String> subtypes;
 
     private final boolean isPassiveMode;
+    private final boolean removeExpiredService;
 
     /** Parcelable constructs for a {@link MdnsServiceInfo}. */
-    MdnsSearchOptions(List<String> subtypes, boolean isPassiveMode) {
+    MdnsSearchOptions(List<String> subtypes, boolean isPassiveMode, boolean removeExpiredService) {
         this.subtypes = new ArrayList<>();
         if (subtypes != null) {
             this.subtypes.addAll(subtypes);
         }
         this.isPassiveMode = isPassiveMode;
+        this.removeExpiredService = removeExpiredService;
     }
 
     /** Returns a {@link Builder} for {@link MdnsSearchOptions}. */
@@ -91,6 +93,11 @@
         return isPassiveMode;
     }
 
+    /** Returns {@code true} if service will be removed after its TTL expires. */
+    public boolean removeExpiredService() {
+        return removeExpiredService;
+    }
+
     @Override
     public int describeContents() {
         return 0;
@@ -100,12 +107,14 @@
     public void writeToParcel(Parcel out, int flags) {
         out.writeStringList(subtypes);
         out.writeBoolean(isPassiveMode);
+        out.writeBoolean(removeExpiredService);
     }
 
     /** A builder to create {@link MdnsSearchOptions}. */
     public static final class Builder {
         private final Set<String> subtypes;
         private boolean isPassiveMode = true;
+        private boolean removeExpiredService;
 
         private Builder() {
             subtypes = new ArraySet<>();
@@ -136,8 +145,7 @@
 
         /**
          * Sets if the passive mode scan should be used. The passive mode scans less frequently in
-         * order
-         * to conserve battery and produce less network traffic.
+         * order to conserve battery and produce less network traffic.
          *
          * @param isPassiveMode If set to {@code true}, passive mode will be used. If set to {@code
          *                      false}, active mode will be used.
@@ -147,9 +155,20 @@
             return this;
         }
 
+        /**
+         * Sets if the service should be removed after TTL.
+         *
+         * @param removeExpiredService If set to {@code true}, the service will be removed after TTL
+         */
+        public Builder setRemoveExpiredService(boolean removeExpiredService) {
+            this.removeExpiredService = removeExpiredService;
+            return this;
+        }
+
         /** Builds a {@link MdnsSearchOptions} with the arguments supplied to this builder. */
         public MdnsSearchOptions build() {
-            return new MdnsSearchOptions(new ArrayList<>(subtypes), isPassiveMode);
+            return new MdnsSearchOptions(
+                    new ArrayList<>(subtypes), isPassiveMode, removeExpiredService);
         }
     }
 }
\ No newline at end of file
diff --git a/service/mdns/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java b/service/mdns/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
index e335de9..4fbc809 100644
--- a/service/mdns/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
+++ b/service/mdns/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
@@ -16,7 +16,11 @@
 
 package com.android.server.connectivity.mdns;
 
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
+
 import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.os.SystemClock;
 import android.text.TextUtils;
 import android.util.ArraySet;
 import android.util.Pair;
@@ -28,12 +32,12 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.TimeUnit;
 
 /**
  * Instance of this class sends and receives mDNS packets of a given service type and invoke
@@ -53,6 +57,12 @@
     private final Object lock = new Object();
     private final Set<MdnsServiceBrowserListener> listeners = new ArraySet<>();
     private final Map<String, MdnsResponse> instanceNameToResponse = new HashMap<>();
+    private final boolean removeServiceAfterTtlExpires =
+            MdnsConfigs.removeServiceAfterTtlExpires();
+    private final boolean allowSearchOptionsToRemoveExpiredService =
+            MdnsConfigs.allowSearchOptionsToRemoveExpiredService();
+
+    @Nullable private MdnsSearchOptions searchOptions;
 
     // The session ID increases when startSendAndReceive() is called where we schedule a
     // QueryTask for
@@ -115,6 +125,7 @@
             @NonNull MdnsServiceBrowserListener listener,
             @NonNull MdnsSearchOptions searchOptions) {
         synchronized (lock) {
+            this.searchOptions = searchOptions;
             if (!listeners.contains(listener)) {
                 listeners.add(listener);
                 for (MdnsResponse existingResponse : instanceNameToResponse.values()) {
@@ -164,10 +175,23 @@
     }
 
     public synchronized void processResponse(@NonNull MdnsResponse response) {
-        if (response.isGoodbye()) {
-            onGoodbyeReceived(response.getServiceInstanceName());
+        if (shouldRemoveServiceAfterTtlExpires()) {
+            // Because {@link QueryTask} and {@link processResponse} are running in different
+            // threads. We need to synchronize {@link lock} to protect
+            // {@link instanceNameToResponse} won’t be modified at the same time.
+            synchronized (lock) {
+                if (response.isGoodbye()) {
+                    onGoodbyeReceived(response.getServiceInstanceName());
+                } else {
+                    onResponseReceived(response);
+                }
+            }
         } else {
-            onResponseReceived(response);
+            if (response.isGoodbye()) {
+                onGoodbyeReceived(response.getServiceInstanceName());
+            } else {
+                onResponseReceived(response);
+            }
         }
     }
 
@@ -212,6 +236,15 @@
         }
     }
 
+    private boolean shouldRemoveServiceAfterTtlExpires() {
+        if (removeServiceAfterTtlExpires) {
+            return true;
+        }
+        return allowSearchOptionsToRemoveExpiredService
+                && searchOptions != null
+                && searchOptions.removeExpiredService();
+    }
+
     @VisibleForTesting
     MdnsPacketWriter createMdnsPacketWriter() {
         return new MdnsPacketWriter(DEFAULT_MTU);
@@ -359,11 +392,27 @@
                         listener.onDiscoveryQuerySent(result.second, result.first);
                     }
                 }
+                if (shouldRemoveServiceAfterTtlExpires()) {
+                    Iterator<MdnsResponse> iter = instanceNameToResponse.values().iterator();
+                    while (iter.hasNext()) {
+                        MdnsResponse existingResponse = iter.next();
+                        if (existingResponse.isComplete()
+                                && existingResponse
+                                .getServiceRecord()
+                                .getRemainingTTL(SystemClock.elapsedRealtime())
+                                == 0) {
+                            iter.remove();
+                            for (MdnsServiceBrowserListener listener : listeners) {
+                                listener.onServiceRemoved(
+                                        existingResponse.getServiceInstanceName());
+                            }
+                        }
+                    }
+                }
                 QueryTaskConfig config = this.config.getConfigForNextRun();
                 requestTaskFuture =
                         executor.schedule(
-                                new QueryTask(config), config.timeToRunNextTaskInMs,
-                                TimeUnit.MILLISECONDS);
+                                new QueryTask(config), config.timeToRunNextTaskInMs, MILLISECONDS);
             }
         }
     }
diff --git a/service/native/TrafficController.cpp b/service/native/TrafficController.cpp
index fc76ae5..8f6df21 100644
--- a/service/native/TrafficController.cpp
+++ b/service/native/TrafficController.cpp
@@ -595,7 +595,7 @@
     }
 }
 
-void TrafficController::dump(int fd, bool verbose) {
+void TrafficController::dump(int fd, bool verbose __unused) {
     std::lock_guard guard(mMutex);
     DumpWriter dw(fd);
 
@@ -623,31 +623,6 @@
                getMapStatus(mConfigurationMap.getMap(), CONFIGURATION_MAP_PATH).c_str());
     dw.println("mUidOwnerMap status: %s",
                getMapStatus(mUidOwnerMap.getMap(), UID_OWNER_MAP_PATH).c_str());
-
-    if (!verbose) {
-        return;
-    }
-
-    dw.blankline();
-    dw.println("BPF map content:");
-
-    ScopedIndent indentForMapContent(dw);
-
-    // Print CookieTagMap content.
-    // TagSocketTest in CTS was using the output of mCookieTagMap dump.
-    // So, mCookieTagMap dump can not be removed until the previous CTS support period is over.
-    dumpBpfMap("mCookieTagMap", dw, "");
-    const auto printCookieTagInfo = [&dw](const uint64_t& key, const UidTagValue& value,
-                                          const BpfMap<uint64_t, UidTagValue>&) {
-        dw.println("cookie=%" PRIu64 " tag=0x%x uid=%u", key, value.tag, value.uid);
-        return base::Result<void>();
-    };
-    base::Result<void> res = mCookieTagMap.iterateWithValue(printCookieTagInfo);
-    if (!res.ok()) {
-        dw.println("mCookieTagMap print end with error: %s", res.error().message().c_str());
-    }
-
-    dw.blankline();
 }
 
 }  // namespace net
diff --git a/service/native/TrafficControllerTest.cpp b/service/native/TrafficControllerTest.cpp
index 6cb0940..57f32af 100644
--- a/service/native/TrafficControllerTest.cpp
+++ b/service/native/TrafficControllerTest.cpp
@@ -59,7 +59,6 @@
 constexpr uid_t TEST_UID3 = 98765;
 constexpr uint32_t TEST_TAG = 42;
 constexpr uint32_t TEST_COUNTERSET = 1;
-constexpr int TEST_COOKIE = 1;
 constexpr int TEST_IFINDEX = 999;
 constexpr int RXPACKETS = 1;
 constexpr int RXBYTES = 100;
@@ -769,46 +768,6 @@
     expectPrivilegedUserSetEmpty();
 }
 
-TEST_F(TrafficControllerTest, TestDumpsys) {
-    StatsKey tagStatsMapKey;
-    populateFakeStats(TEST_COOKIE, TEST_UID, TEST_TAG, &tagStatsMapKey);
-    populateFakeCounterSet(TEST_UID3, TEST_COUNTERSET);
-
-    // Expect: (part of this depends on hard-code values in populateFakeStats())
-    //
-    // mCookieTagMap:
-    // cookie=1 tag=0x2a uid=10086
-    //
-    // mUidCounterSetMap:
-    // 98765 1
-    //
-    // mAppUidStatsMap::
-    // uid rxBytes rxPackets txBytes txPackets
-    // 10086 100 1 0 0
-    //
-    // mStatsMapA:
-    // ifaceIndex ifaceName tag_hex uid_int cnt_set rxBytes rxPackets txBytes txPackets
-    // 999 test0 0x2a 10086 1 100 1 0 0
-    std::vector<std::string> expectedLines = {
-        "mCookieTagMap:",
-        fmt::format("cookie={} tag={:#x} uid={}", TEST_COOKIE, TEST_TAG, TEST_UID)};
-
-    EXPECT_TRUE(expectDumpsysContains(expectedLines));
-}
-
-TEST_F(TrafficControllerTest, dumpsysInvalidMaps) {
-    makeTrafficControllerMapsInvalid();
-
-    const std::string kErrIterate = "print end with error: Get firstKey map -1 failed: "
-            "Bad file descriptor";
-    const std::string kErrReadRulesConfig = "read ownerMatch configure failed with error: "
-            "Read value of map -1 failed: Bad file descriptor";
-
-    std::vector<std::string> expectedLines = {
-        fmt::format("mCookieTagMap {}", kErrIterate)};
-    EXPECT_TRUE(expectDumpsysContains(expectedLines));
-}
-
 TEST_F(TrafficControllerTest, getFirewallType) {
     static const struct TestConfig {
         ChildChain childChain;
diff --git a/service/proguard.flags b/service/proguard.flags
index 478566c..864a28b 100644
--- a/service/proguard.flags
+++ b/service/proguard.flags
@@ -6,3 +6,12 @@
 -keepclassmembers public class * extends **.com.android.net.module.util.Struct {
     *;
 }
+
+-keepclassmembers class com.android.server.**,android.net.**,com.android.networkstack.** {
+    static final % POLICY_*;
+    static final % NOTIFY_TYPE_*;
+    static final % TRANSPORT_*;
+    static final % CMD_*;
+    static final % EVENT_*;
+}
+
diff --git a/service/src/com/android/server/BpfNetMaps.java b/service/src/com/android/server/BpfNetMaps.java
index d560747..aea2103 100644
--- a/service/src/com/android/server/BpfNetMaps.java
+++ b/service/src/com/android/server/BpfNetMaps.java
@@ -1022,11 +1022,25 @@
         }
         mDeps.nativeDump(fd, verbose);
 
+        pw.println();
+        pw.println("sEnableJavaBpfMap: " + sEnableJavaBpfMap);
         if (verbose) {
+            pw.println();
+            pw.println("BPF map content:");
+            pw.increaseIndent();
+
             dumpOwnerMatchConfig(pw);
             dumpCurrentStatsMapConfig(pw);
             pw.println();
 
+            // TODO: Remove CookieTagMap content dump
+            // NetworkStatsService also dumps CookieTagMap and NetworkStatsService is a right place
+            // to dump CookieTagMap. But the TagSocketTest in CTS depends on this dump so the tests
+            // need to be updated before remove the dump from BpfNetMaps.
+            BpfDump.dumpMap(sCookieTagMap, pw, "sCookieTagMap",
+                    (key, value) -> "cookie=" + key.socketCookie
+                            + " tag=0x" + Long.toHexString(value.tag)
+                            + " uid=" + value.uid);
             BpfDump.dumpMap(sUidOwnerMap, pw, "sUidOwnerMap",
                     (uid, match) -> {
                         if ((match.rule & IIF_MATCH) != 0) {
@@ -1038,6 +1052,7 @@
                     });
             BpfDump.dumpMap(sUidPermissionMap, pw, "sUidPermissionMap",
                     (uid, permission) -> uid.val + " " + permissionToString(permission.val));
+            pw.decreaseIndent();
         }
     }
 
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index 0774e24..195f046 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -249,13 +249,13 @@
 import com.android.modules.utils.BasicShellCommandHandler;
 import com.android.modules.utils.build.SdkLevel;
 import com.android.net.module.util.BaseNetdUnsolicitedEventListener;
+import com.android.net.module.util.BitUtils;
 import com.android.net.module.util.CollectionUtils;
 import com.android.net.module.util.DeviceConfigUtils;
 import com.android.net.module.util.InterfaceParams;
 import com.android.net.module.util.LinkPropertiesUtils.CompareOrUpdateResult;
 import com.android.net.module.util.LinkPropertiesUtils.CompareResult;
 import com.android.net.module.util.LocationPermissionChecker;
-import com.android.net.module.util.NetworkCapabilitiesUtils;
 import com.android.net.module.util.PerUidCounter;
 import com.android.net.module.util.PermissionUtils;
 import com.android.net.module.util.TcUtils;
@@ -3650,8 +3650,18 @@
                     break;
                 }
                 case NetworkAgent.EVENT_UNREGISTER_AFTER_REPLACEMENT: {
-                    // If nai is not yet created, or is already destroyed, ignore.
-                    if (!shouldDestroyNativeNetwork(nai)) break;
+                    if (!nai.isCreated()) {
+                        Log.d(TAG, "unregisterAfterReplacement on uncreated " + nai.toShortString()
+                                + ", tearing down instead");
+                        teardownUnneededNetwork(nai);
+                        break;
+                    }
+
+                    if (nai.isDestroyed()) {
+                        Log.d(TAG, "unregisterAfterReplacement on destroyed " + nai.toShortString()
+                                + ", ignoring");
+                        break;
+                    }
 
                     final int timeoutMs = (int) arg.second;
                     if (timeoutMs < 0 || timeoutMs > NetworkAgent.MAX_TEARDOWN_DELAY_MS) {
@@ -7839,7 +7849,7 @@
             @NonNull NetworkCapabilities agentCaps, @NonNull NetworkCapabilities newNc) {
         underlyingNetworks = underlyingNetworksOrDefault(
                 agentCaps.getOwnerUid(), underlyingNetworks);
-        long transportTypes = NetworkCapabilitiesUtils.packBits(agentCaps.getTransportTypes());
+        long transportTypes = BitUtils.packBits(agentCaps.getTransportTypes());
         int downKbps = NetworkCapabilities.LINK_BANDWIDTH_UNSPECIFIED;
         int upKbps = NetworkCapabilities.LINK_BANDWIDTH_UNSPECIFIED;
         // metered if any underlying is metered, or originally declared metered by the agent.
@@ -7892,7 +7902,7 @@
             suspended = false;
         }
 
-        newNc.setTransportTypes(NetworkCapabilitiesUtils.unpackBits(transportTypes));
+        newNc.setTransportTypes(BitUtils.unpackBits(transportTypes));
         newNc.setLinkDownstreamBandwidthKbps(downKbps);
         newNc.setLinkUpstreamBandwidthKbps(upKbps);
         newNc.setCapability(NET_CAPABILITY_NOT_METERED, !metered);
@@ -8004,6 +8014,10 @@
             @NonNull final NetworkCapabilities nc) {
         NetworkCapabilities newNc = mixInCapabilities(nai, nc);
         if (Objects.equals(nai.networkCapabilities, newNc)) return;
+        final String differences = newNc.describeCapsDifferencesFrom(nai.networkCapabilities);
+        if (null != differences) {
+            Log.i(TAG, "Update capabilities for net " + nai.network + " : " + differences);
+        }
         updateNetworkPermissions(nai, newNc);
         final NetworkCapabilities prevNc = nai.getAndSetNetworkCapabilities(newNc);
 
diff --git a/service/src/com/android/server/connectivity/FullScore.java b/service/src/com/android/server/connectivity/FullScore.java
index aec4a71..2303894 100644
--- a/service/src/com/android/server/connectivity/FullScore.java
+++ b/service/src/com/android/server/connectivity/FullScore.java
@@ -23,7 +23,10 @@
 import static android.net.NetworkScore.KEEP_CONNECTED_NONE;
 import static android.net.NetworkScore.POLICY_YIELD_TO_BAD_WIFI;
 
+import static com.android.net.module.util.BitUtils.appendStringRepresentationOfBitMaskToStringBuilder;
+
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.net.NetworkAgentConfig;
 import android.net.NetworkCapabilities;
 import android.net.NetworkScore;
@@ -333,6 +336,35 @@
                 + 5 * mKeepConnectedReason;
     }
 
+    /**
+     * Returns a short but human-readable string of updates from an older score.
+     * @param old the old capabilities to diff from
+     * @return a string fit for logging differences, or null if no differences.
+     *         this method cannot return the empty string.
+     */
+    @Nullable
+    public String describeDifferencesFrom(@Nullable final FullScore old) {
+        final long oldPolicies = null == old ? 0 : old.mPolicies;
+        final long changed = oldPolicies ^ mPolicies;
+        if (0 == changed) return null;
+        // If the control reaches here, there are changes (additions, removals, or both) so
+        // the code below is guaranteed to add something to the string and can't return "".
+        final long removed = oldPolicies & changed;
+        final long added = mPolicies & changed;
+        final StringBuilder sb = new StringBuilder();
+        if (0 != removed) {
+            sb.append("-");
+            appendStringRepresentationOfBitMaskToStringBuilder(sb, removed,
+                    FullScore::policyNameOf, "-");
+        }
+        if (0 != added) {
+            sb.append("+");
+            appendStringRepresentationOfBitMaskToStringBuilder(sb, added,
+                    FullScore::policyNameOf, "+");
+        }
+        return sb.toString();
+    }
+
     // Example output :
     // Score(Policies : EVER_USER_SELECTED&IS_VALIDATED ; KeepConnected : )
     @Override
diff --git a/service/src/com/android/server/connectivity/NetworkAgentInfo.java b/service/src/com/android/server/connectivity/NetworkAgentInfo.java
index 2e92d43..85282cb 100644
--- a/service/src/com/android/server/connectivity/NetworkAgentInfo.java
+++ b/service/src/com/android/server/connectivity/NetworkAgentInfo.java
@@ -1003,9 +1003,7 @@
             @NonNull final NetworkCapabilities nc) {
         final NetworkCapabilities oldNc = networkCapabilities;
         networkCapabilities = nc;
-        mScore = mScore.mixInScore(networkCapabilities, networkAgentConfig, everValidated(),
-                0L != getAvoidUnvalidated(), yieldToBadWiFi(),
-                0L != mFirstEvaluationConcludedTime, isDestroyed());
+        updateScoreForNetworkAgentUpdate();
         final NetworkMonitorManager nm = mNetworkMonitor;
         if (nm != null) {
             nm.notifyNetworkCapabilitiesChanged(nc);
@@ -1207,9 +1205,11 @@
      * Mix-in the ConnectivityService-managed bits in the score.
      */
     public void setScore(final NetworkScore score) {
+        final FullScore oldScore = mScore;
         mScore = FullScore.fromNetworkScore(score, networkCapabilities, networkAgentConfig,
                 everValidated(), 0L != getAvoidUnvalidated(), yieldToBadWiFi(),
                 0L != mFirstEvaluationConcludedTime, isDestroyed());
+        maybeLogDifferences(oldScore);
     }
 
     /**
@@ -1218,9 +1218,22 @@
      * Call this after changing any data that might affect the score (e.g., agent config).
      */
     public void updateScoreForNetworkAgentUpdate() {
+        final FullScore oldScore = mScore;
         mScore = mScore.mixInScore(networkCapabilities, networkAgentConfig,
                 everValidated(), 0L != getAvoidUnvalidated(), yieldToBadWiFi(),
                 0L != mFirstEvaluationConcludedTime, isDestroyed());
+        maybeLogDifferences(oldScore);
+    }
+
+    /**
+     * Prints score differences to logcat, if any.
+     * @param oldScore the old score. Differences from |oldScore| to |this| are logged, if any.
+     */
+    public void maybeLogDifferences(final FullScore oldScore) {
+        final String differences = mScore.describeDifferencesFrom(oldScore);
+        if (null != differences) {
+            Log.i(TAG, "Update score for net " + network + " : " + differences);
+        }
     }
 
     /**
diff --git a/tests/common/java/android/net/NetworkCapabilitiesTest.java b/tests/common/java/android/net/NetworkCapabilitiesTest.java
index c30e1d3..7b374d2 100644
--- a/tests/common/java/android/net/NetworkCapabilitiesTest.java
+++ b/tests/common/java/android/net/NetworkCapabilitiesTest.java
@@ -36,6 +36,7 @@
 import static android.net.NetworkCapabilities.NET_CAPABILITY_PARTIAL_CONNECTIVITY;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_PRIORITIZE_BANDWIDTH;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_PRIORITIZE_LATENCY;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_SUPL;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_TRUSTED;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_VALIDATED;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_WIFI_P2P;
@@ -1370,4 +1371,23 @@
 
         assertEquals(expectedNcBuilder.build(), restrictedNc);
     }
+
+    @Test
+    public void testDescribeCapsDifferences() throws Exception {
+        final NetworkCapabilities nc1 = new NetworkCapabilities.Builder()
+                .addCapability(NET_CAPABILITY_MMS)
+                .addCapability(NET_CAPABILITY_OEM_PAID)
+                .addCapability(NET_CAPABILITY_INTERNET)
+                .build();
+        final NetworkCapabilities nc2 = new NetworkCapabilities.Builder()
+                .addCapability(NET_CAPABILITY_CAPTIVE_PORTAL)
+                .addCapability(NET_CAPABILITY_SUPL)
+                .addCapability(NET_CAPABILITY_VALIDATED)
+                .addCapability(NET_CAPABILITY_INTERNET)
+                .build();
+        assertEquals("-MMS-OEM_PAID+SUPL+VALIDATED+CAPTIVE_PORTAL",
+                nc2.describeCapsDifferencesFrom(nc1));
+        assertEquals("-SUPL-VALIDATED-CAPTIVE_PORTAL+MMS+OEM_PAID",
+                nc1.describeCapsDifferencesFrom(nc2));
+    }
 }
diff --git a/tests/common/java/android/net/metrics/IpConnectivityLogTest.java b/tests/common/java/android/net/metrics/IpConnectivityLogTest.java
index 93cf748..b6e9b95 100644
--- a/tests/common/java/android/net/metrics/IpConnectivityLogTest.java
+++ b/tests/common/java/android/net/metrics/IpConnectivityLogTest.java
@@ -19,7 +19,7 @@
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 
-import static com.android.net.module.util.NetworkCapabilitiesUtils.unpackBits;
+import static com.android.net.module.util.BitUtils.unpackBits;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index 310d7bf..218eb04 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -270,7 +270,10 @@
     private static final int MIN_KEEPALIVE_INTERVAL = 10;
 
     private static final int NETWORK_CALLBACK_TIMEOUT_MS = 30_000;
-    private static final int LISTEN_ACTIVITY_TIMEOUT_MS = 5_000;
+    // Timeout for waiting network to be validated. Set the timeout to 30s, which is more than
+    // DNS timeout.
+    // TODO(b/252972908): reset the original timer when aosp/2188755 is ramped up.
+    private static final int LISTEN_ACTIVITY_TIMEOUT_MS = 30_000;
     private static final int NO_CALLBACK_TIMEOUT_MS = 100;
     private static final int SOCKET_TIMEOUT_MS = 100;
     private static final int NUM_TRIES_MULTIPATH_PREF_CHECK = 20;
@@ -2618,8 +2621,7 @@
             // the network with the TEST transport. Also wait for validation here, in case there
             // is a bug that's only visible when the network is validated.
             setWifiMeteredStatusAndWait(ssid, true /* isMetered */, true /* waitForValidation */);
-            defaultCallback.expectCallback(CallbackEntry.LOST, wifiNetwork,
-                    NETWORK_CALLBACK_TIMEOUT_MS);
+            defaultCallback.expect(CallbackEntry.LOST, wifiNetwork, NETWORK_CALLBACK_TIMEOUT_MS);
             waitForAvailable(defaultCallback, tnt.getNetwork());
             // Depending on if this device has cellular connectivity or not, multiple available
             // callbacks may be received. Eventually, metered Wi-Fi should be the final available
@@ -2628,10 +2630,9 @@
             waitForAvailable(systemDefaultCallback, TRANSPORT_WIFI);
         }, /* cleanup */ () -> {
             // Validate that removing the test network will fallback to the default network.
-            runWithShellPermissionIdentity(tnt::teardown);
-            defaultCallback.expectCallback(CallbackEntry.LOST, tnt.getNetwork(),
-                    NETWORK_CALLBACK_TIMEOUT_MS);
-            waitForAvailable(defaultCallback);
+                runWithShellPermissionIdentity(tnt::teardown);
+                defaultCallback.expect(CallbackEntry.LOST, tnt, NETWORK_CALLBACK_TIMEOUT_MS);
+                waitForAvailable(defaultCallback);
             }, /* cleanup */ () -> {
                 setWifiMeteredStatusAndWait(ssid, oldMeteredValue, false /* waitForValidation */);
             }, /* cleanup */ () -> {
@@ -2666,8 +2667,7 @@
             waitForAvailable(systemDefaultCallback, wifiNetwork);
         }, /* cleanup */ () -> {
                 runWithShellPermissionIdentity(tnt::teardown);
-                defaultCallback.expectCallback(CallbackEntry.LOST, tnt.getNetwork(),
-                        NETWORK_CALLBACK_TIMEOUT_MS);
+                defaultCallback.expect(CallbackEntry.LOST, tnt, NETWORK_CALLBACK_TIMEOUT_MS);
 
                 // This network preference should only ever use the test network therefore available
                 // should not trigger when the test network goes down (e.g. switch to cellular).
diff --git a/tests/cts/net/src/android/net/cts/EthernetManagerTest.kt b/tests/cts/net/src/android/net/cts/EthernetManagerTest.kt
index 122eb15..7e91478 100644
--- a/tests/cts/net/src/android/net/cts/EthernetManagerTest.kt
+++ b/tests/cts/net/src/android/net/cts/EthernetManagerTest.kt
@@ -60,7 +60,6 @@
 import android.os.SystemProperties
 import android.os.Process
 import android.platform.test.annotations.AppModeFull
-import android.util.ArraySet
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.net.module.util.ArrayTrackRecord
 import com.android.net.module.util.TrackRecord
@@ -75,7 +74,6 @@
 import com.android.testutils.RouterAdvertisementResponder
 import com.android.testutils.TapPacketReader
 import com.android.testutils.TestableNetworkCallback
-import com.android.testutils.anyNetwork
 import com.android.testutils.assertThrows
 import com.android.testutils.runAsShell
 import com.android.testutils.waitForIdle
@@ -103,7 +101,10 @@
 import kotlin.test.fail
 
 private const val TAG = "EthernetManagerTest"
-private const val TIMEOUT_MS = 2000L
+// This timeout does not affect the test duration for passing tests. It needs to be long enough to
+// account for RS delay (and potentially the first retry interval (4s)). There have been failures
+// where the interface did not gain provisioning within the allotted timeout.
+private const val TIMEOUT_MS = 10_000L
 // Timeout used to confirm no callbacks matching given criteria are received. Must be long enough to
 // process all callbacks including ip provisioning when using the updateConfiguration API.
 // Note that increasing this timeout increases the test duration.
@@ -253,7 +254,7 @@
         }
 
         fun <T : CallbackEntry> expectCallback(expected: T): T {
-            val event = pollForNextCallback()
+            val event = pollOrThrow()
             assertEquals(expected, event)
             return event as T
         }
@@ -270,7 +271,7 @@
                 InterfaceStateChanged(iface, state, role,
                         if (state != STATE_ABSENT) DEFAULT_IP_CONFIGURATION else null)
 
-        fun pollForNextCallback(): CallbackEntry {
+        fun pollOrThrow(): CallbackEntry {
             return events.poll(TIMEOUT_MS) ?: fail("Did not receive callback after ${TIMEOUT_MS}ms")
         }
 
@@ -303,8 +304,8 @@
             available.completeExceptionally(IllegalStateException("onUnavailable was called"))
         }
 
-        fun expectOnAvailable(): String {
-            return available.get(TIMEOUT_MS, TimeUnit.MILLISECONDS)
+        fun expectOnAvailable(timeout: Long = TIMEOUT_MS): String {
+            return available.get(timeout, TimeUnit.MILLISECONDS)
         }
 
         fun expectOnUnavailable() {
@@ -406,9 +407,7 @@
     }
 
     private fun addInterfaceStateListener(listener: EthernetStateListener) {
-        runAsShell(CONNECTIVITY_USE_RESTRICTED_NETWORKS) {
-            em.addInterfaceStateListener(handler::post, listener)
-        }
+        em.addInterfaceStateListener(handler::post, listener)
         addedListeners.add(listener)
     }
 
@@ -447,14 +446,18 @@
     }
 
     private fun requestNetwork(request: NetworkRequest): TestableNetworkCallback {
-        return TestableNetworkCallback().also {
+        return TestableNetworkCallback(
+                timeoutMs = TIMEOUT_MS,
+                noCallbackTimeoutMs = NO_CALLBACK_TIMEOUT_MS).also {
             cm.requestNetwork(request, it)
             registeredCallbacks.add(it)
         }
     }
 
     private fun registerNetworkListener(request: NetworkRequest): TestableNetworkCallback {
-        return TestableNetworkCallback().also {
+        return TestableNetworkCallback(
+                timeoutMs = TIMEOUT_MS,
+                noCallbackTimeoutMs = NO_CALLBACK_TIMEOUT_MS).also {
             cm.registerNetworkCallback(request, it)
             registeredCallbacks.add(it)
         }
@@ -521,34 +524,25 @@
         NetworkRequest.Builder(NetworkRequest(ETH_REQUEST))
             .setNetworkSpecifier(EthernetNetworkSpecifier(ifaceName)).build()
 
-    // It can take multiple seconds for the network to become available.
-    private fun TestableNetworkCallback.expectAvailable() =
-            expectCallback<Available>(anyNetwork(), 5000 /* ms timeout */).network
-
-    private fun TestableNetworkCallback.expectLost(n: Network = anyNetwork()) =
-            expectCallback<Lost>(n, 5000 /* ms timeout */)
-
     // b/233534110: eventuallyExpect<Lost>() does not advance ReadHead, use
     // eventuallyExpect(Lost::class) instead.
     private fun TestableNetworkCallback.eventuallyExpectLost(n: Network? = null) =
-        eventuallyExpect(Lost::class, TIMEOUT_MS) { n?.equals(it.network) ?: true }
+        eventuallyExpect(Lost::class) { n?.equals(it.network) ?: true }
 
     private fun TestableNetworkCallback.assertNeverLost(n: Network? = null) =
-        assertNoCallbackThat(NO_CALLBACK_TIMEOUT_MS) {
+        assertNoCallbackThat() {
             it is Lost && (n?.equals(it.network) ?: true)
         }
 
     private fun TestableNetworkCallback.assertNeverAvailable(n: Network? = null) =
-        assertNoCallbackThat() { it is Available && (n?.equals(it.network) ?: true) }
+        assertNoCallbackThat { it is Available && (n?.equals(it.network) ?: true) }
 
     private fun TestableNetworkCallback.expectCapabilitiesWithInterfaceName(name: String) =
-        expectCapabilitiesThat(anyNetwork()) {
-            it.networkSpecifier == EthernetNetworkSpecifier(name)
-        }
+        expect<CapabilitiesChanged> { it.caps.networkSpecifier == EthernetNetworkSpecifier(name) }
 
     private fun TestableNetworkCallback.eventuallyExpectCapabilities(nc: NetworkCapabilities) {
         // b/233534110: eventuallyExpect<CapabilitiesChanged>() does not advance ReadHead.
-        eventuallyExpect(CapabilitiesChanged::class, TIMEOUT_MS) {
+        eventuallyExpect(CapabilitiesChanged::class) {
             // CS may mix in additional capabilities, so NetworkCapabilities#equals cannot be used.
             // Check if all expected capabilities are present instead.
             it is CapabilitiesChanged && nc.capabilities.all { c -> it.caps.hasCapability(c) }
@@ -559,7 +553,7 @@
         config: StaticIpConfiguration
     ) {
         // b/233534110: eventuallyExpect<LinkPropertiesChanged>() does not advance ReadHead.
-        eventuallyExpect(LinkPropertiesChanged::class, TIMEOUT_MS) {
+        eventuallyExpect(LinkPropertiesChanged::class) {
             it is LinkPropertiesChanged && it.lp.linkAddresses.any { la ->
                 la.isSameAddressAs(config.ipAddress)
             }
@@ -568,11 +562,18 @@
 
     @Test
     fun testCallbacks() {
+        // Only run this test when no non-restricted / physical interfaces are present.
+        // This test ensures that interface state listeners function properly, so the assumption
+        // check is explicitly *not* using an interface state listener.
+        // Since restricted interfaces cannot be used for tethering,
+        // assumeNoInterfaceForTetheringAvailable() is an okay proxy.
+        assumeNoInterfaceForTetheringAvailable()
+
         // If an interface exists when the callback is registered, it is reported on registration.
         val iface = createInterface()
         val listener1 = EthernetStateListener()
         addInterfaceStateListener(listener1)
-        validateListenerOnRegistration(listener1)
+        listener1.expectCallback(iface, STATE_LINK_UP, ROLE_CLIENT)
 
         // If an interface appears, existing callbacks see it.
         val iface2 = createInterface()
@@ -582,7 +583,8 @@
         // Register a new listener, it should see state of all existing interfaces immediately.
         val listener2 = EthernetStateListener()
         addInterfaceStateListener(listener2)
-        validateListenerOnRegistration(listener2)
+        listener2.expectCallback(iface, STATE_LINK_UP, ROLE_CLIENT)
+        listener2.expectCallback(iface2, STATE_LINK_UP, ROLE_CLIENT)
 
         // Removing interfaces first sends link down, then STATE_ABSENT/ROLE_NONE.
         removeInterface(iface)
@@ -601,9 +603,10 @@
 
     @Test
     fun testCallbacks_withRunningInterface() {
-        // This test disables ethernet, so check that adb is not connected over ethernet.
         assumeFalse(isAdbOverEthernet())
-        assumeTrue(em.getInterfaceList().isEmpty())
+        // Only run this test when no non-restricted / physical interfaces are present.
+        assumeNoInterfaceForTetheringAvailable()
+
         val iface = createInterface()
         val listener = EthernetStateListener()
         addInterfaceStateListener(listener)
@@ -623,7 +626,7 @@
         // see aosp/2123900.
         try {
             // assumeException does not exist.
-            requestTetheredInterface().expectOnAvailable()
+            requestTetheredInterface().expectOnAvailable(NO_CALLBACK_TIMEOUT_MS)
             // interface used for tethering is available, throw an assumption error.
             assumeTrue(false)
         } catch (e: TimeoutException) {
@@ -659,28 +662,6 @@
         listener.expectCallback(iface, STATE_LINK_UP, ROLE_CLIENT)
     }
 
-    /**
-     * Validate all interfaces are returned for an EthernetStateListener upon registration.
-     */
-    private fun validateListenerOnRegistration(listener: EthernetStateListener) {
-        // Get all tracked interfaces to validate on listener registration. Ordering and interface
-        // state (up/down) can't be validated for interfaces not created as part of testing.
-        val ifaces = em.getInterfaceList()
-        val polledIfaces = ArraySet<String>()
-        for (i in ifaces) {
-            val event = (listener.pollForNextCallback() as InterfaceStateChanged)
-            val iface = event.iface
-            assertTrue(polledIfaces.add(iface), "Duplicate interface $iface returned")
-            assertTrue(ifaces.contains(iface), "Untracked interface $iface returned")
-            // If the event's iface was created in the test, additional criteria can be validated.
-            createdIfaces.find { it.name.equals(iface) }?.let {
-                assertEquals(event, listener.createChangeEvent(it.name, STATE_LINK_UP, ROLE_CLIENT))
-            }
-        }
-        // Assert all callbacks are accounted for.
-        listener.assertNoCallback()
-    }
-
     @Test
     fun testGetInterfaceList() {
         // Create two test interfaces and check the return list contains the interface names.
@@ -711,7 +692,7 @@
         listenerCb.assertNeverAvailable()
 
         val cb = requestNetwork(ETH_REQUEST)
-        val network = cb.expectAvailable()
+        val network = cb.expect<Available>().network
 
         cb.assertNeverLost()
         releaseRequest(cb)
@@ -728,7 +709,7 @@
         // createInterface and the interface actually being properly registered with the ethernet
         // module, so it is extremely unlikely that the CS handler thread has not run until then.
         val iface = createInterface()
-        val network = cb.expectAvailable()
+        val network = cb.expect<Available>().network
 
         // remove interface before network request has been removed
         cb.assertNeverLost()
@@ -743,7 +724,7 @@
 
         val cb = requestNetwork(ETH_REQUEST.copyWithEthernetSpecifier(iface2.name))
 
-        val network = cb.expectAvailable()
+        cb.expect<Available>()
         cb.expectCapabilitiesWithInterfaceName(iface2.name)
 
         removeInterface(iface1)
@@ -757,7 +738,7 @@
         val iface1 = createInterface()
 
         val cb = requestNetwork(ETH_REQUEST)
-        val network = cb.expectAvailable()
+        val network = cb.expect<Available>().network
 
         // create another network and verify the request sticks to the current network
         val iface2 = createInterface()
@@ -766,7 +747,7 @@
         // remove iface1 and verify the request brings up iface2
         removeInterface(iface1)
         cb.eventuallyExpectLost(network)
-        val network2 = cb.expectAvailable()
+        cb.expect<Available>()
     }
 
     @Test
@@ -778,12 +759,12 @@
         val cb2 = requestNetwork(ETH_REQUEST.copyWithEthernetSpecifier(iface2.name))
         val cb3 = requestNetwork(ETH_REQUEST)
 
-        cb1.expectAvailable()
+        cb1.expect<Available>()
         cb1.expectCapabilitiesWithInterfaceName(iface1.name)
-        cb2.expectAvailable()
+        cb2.expect<Available>()
         cb2.expectCapabilitiesWithInterfaceName(iface2.name)
         // this request can be matched by either network.
-        cb3.expectAvailable()
+        cb3.expect<Available>()
 
         cb1.assertNeverLost()
         cb2.assertNeverLost()
@@ -798,10 +779,10 @@
         val cb1 = requestNetwork(ETH_REQUEST)
 
         val iface = createInterface()
-        val network = cb1.expectAvailable()
+        val network = cb1.expect<Available>().network
 
         val cb2 = requestNetwork(ETH_REQUEST)
-        cb2.expectAvailable()
+        cb2.expect<Available>()
 
         // release the first request; this used to trigger b/197548738
         releaseRequest(cb1)
@@ -824,7 +805,7 @@
         // cb.assertNeverAvailable()
 
         iface.setCarrierEnabled(true)
-        cb.expectAvailable()
+        cb.expect<Available>()
 
         iface.setCarrierEnabled(false)
         cb.eventuallyExpectLost()
@@ -836,7 +817,7 @@
         val iface = createInterface()
         val cb = requestNetwork(ETH_REQUEST)
         // b/233534110: eventuallyExpect<LinkPropertiesChanged>() does not advance ReadHead
-        cb.eventuallyExpect(LinkPropertiesChanged::class, TIMEOUT_MS) {
+        cb.eventuallyExpect(LinkPropertiesChanged::class) {
             it is LinkPropertiesChanged && it.lp.addresses.any {
                 address -> iface.onLinkPrefix.contains(address)
             }
@@ -869,14 +850,14 @@
     fun testEnableDisableInterface_withActiveRequest() {
         val iface = createInterface()
         val cb = requestNetwork(ETH_REQUEST)
-        cb.expectAvailable()
+        cb.expect<Available>()
         cb.assertNeverLost()
 
         disableInterface(iface).expectResult(iface.name)
         cb.eventuallyExpectLost()
 
         enableInterface(iface).expectResult(iface.name)
-        cb.expectAvailable()
+        cb.expect<Available>()
     }
 
     @Test
@@ -903,7 +884,7 @@
     fun testUpdateConfiguration_forBothIpConfigAndCapabilities() {
         val iface = createInterface()
         val cb = requestNetwork(ETH_REQUEST.copyWithEthernetSpecifier(iface.name))
-        val network = cb.expectAvailable()
+        cb.expect<Available>()
 
         updateConfiguration(iface, STATIC_IP_CONFIGURATION, TEST_CAPS).expectResult(iface.name)
         cb.eventuallyExpectCapabilities(TEST_CAPS)
@@ -914,7 +895,7 @@
     fun testUpdateConfiguration_forOnlyIpConfig() {
         val iface = createInterface()
         val cb = requestNetwork(ETH_REQUEST.copyWithEthernetSpecifier(iface.name))
-        val network = cb.expectAvailable()
+        cb.expect<Available>()
 
         updateConfiguration(iface, STATIC_IP_CONFIGURATION).expectResult(iface.name)
         cb.eventuallyExpectLpForStaticConfig(STATIC_IP_CONFIGURATION.staticIpConfiguration)
@@ -924,7 +905,7 @@
     fun testUpdateConfiguration_forOnlyCapabilities() {
         val iface = createInterface()
         val cb = requestNetwork(ETH_REQUEST.copyWithEthernetSpecifier(iface.name))
-        val network = cb.expectAvailable()
+        cb.expect<Available>()
 
         updateConfiguration(iface, capabilities = TEST_CAPS).expectResult(iface.name)
         cb.eventuallyExpectCapabilities(TEST_CAPS)
@@ -941,7 +922,7 @@
 
         // Request the restricted network as the shell with CONNECTIVITY_USE_RESTRICTED_NETWORKS.
         val cb = runAsShell(CONNECTIVITY_USE_RESTRICTED_NETWORKS) { requestNetwork(request) }
-        val network = cb.expectAvailable()
+        val network = cb.expect<Available>().network
         cb.assertNeverLost(network)
 
         // The network is restricted therefore binding to it when available will fail.
@@ -957,10 +938,16 @@
 
         // UpdateConfiguration() currently does a restart on the ethernet interface therefore lost
         // will be expected first before available, as part of the restart.
-        cb.expectLost(network)
-        val updatedNetwork = cb.expectAvailable()
+        cb.expect<Lost>(network)
+        val updatedNetwork = cb.expect<Available>().network
         // With the test process UID allowed, binding to a restricted network should be successful.
         Socket().use { socket -> updatedNetwork.bindSocket(socket) }
+
+        // Reset capabilities to not-restricted, otherwise tearDown won't see the interface callback
+        // as ifaceListener does not have the restricted permission.
+        // TODO: Find a better way to do this when there are more tests around restricted
+        // interfaces.
+        updateConfiguration(iface, capabilities = TEST_CAPS).expectResult(iface.name)
     }
 
     // TODO: consider only having this test in MTS as it makes use of the fact that
@@ -978,7 +965,7 @@
 
         setEthernetEnabled(true)
         val cb = requestNetwork(ETH_REQUEST)
-        cb.expectAvailable()
+        cb.expect<Available>()
         cb.eventuallyExpectCapabilities(TEST_CAPS)
         cb.eventuallyExpectLpForStaticConfig(STATIC_IP_CONFIGURATION.staticIpConfiguration)
     }
@@ -989,7 +976,7 @@
         // createInterface without carrier is racy, so create it and then remove carrier.
         val iface = createInterface()
         val cb = requestNetwork(ETH_REQUEST)
-        cb.expectAvailable()
+        cb.expect<Available>()
 
         iface.setCarrierEnabled(false)
         cb.eventuallyExpectLost()
diff --git a/tests/cts/net/src/android/net/cts/Ikev2VpnTest.java b/tests/cts/net/src/android/net/cts/Ikev2VpnTest.java
index 2b1d173..ac50740 100644
--- a/tests/cts/net/src/android/net/cts/Ikev2VpnTest.java
+++ b/tests/cts/net/src/android/net/cts/Ikev2VpnTest.java
@@ -24,7 +24,6 @@
 import static com.android.compatibility.common.util.SystemUtil.runWithShellPermissionIdentity;
 import static com.android.modules.utils.build.SdkLevel.isAtLeastT;
 import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
-import static com.android.testutils.TestableNetworkCallbackKt.anyNetwork;
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
@@ -520,8 +519,7 @@
                 HexDump.hexStringToByteArray(authResp));
 
         // Verify the VPN network came up
-        final Network vpnNetwork = cb.expectCallback(CallbackEntry.AVAILABLE, anyNetwork())
-                .getNetwork();
+        final Network vpnNetwork = cb.expect(CallbackEntry.AVAILABLE).getNetwork();
 
         if (testSessionKey) {
             final VpnProfileStateShim profileState = mVmShim.getProvisionedVpnProfileState();
@@ -536,8 +534,8 @@
                 && caps.hasCapability(NET_CAPABILITY_INTERNET)
                 && !caps.hasCapability(NET_CAPABILITY_VALIDATED)
                 && Process.myUid() == caps.getOwnerUid());
-        cb.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, vpnNetwork);
-        cb.expectCallback(CallbackEntry.BLOCKED_STATUS, vpnNetwork);
+        cb.expect(CallbackEntry.LINK_PROPERTIES_CHANGED, vpnNetwork);
+        cb.expect(CallbackEntry.BLOCKED_STATUS, vpnNetwork);
 
         // A VPN that requires validation is initially not validated, while one that doesn't
         // immediately validate automatically. Because this VPN can't actually access Internet,
diff --git a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
index d2cd7aa..6df71c8 100644
--- a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
+++ b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
@@ -335,6 +335,28 @@
         mFakeConnectivityService.connect(it.registerForTest(Network(FAKE_NET_ID)))
     }
 
+    fun assertLinkPropertiesEventually(
+        n: Network,
+        description: String,
+        condition: (LinkProperties?) -> Boolean
+    ): LinkProperties? {
+        val deadline = SystemClock.elapsedRealtime() + DEFAULT_TIMEOUT_MS
+        do {
+            val lp = mCM.getLinkProperties(n)
+            if (condition(lp)) return lp
+            SystemClock.sleep(10 /* ms */)
+        } while (SystemClock.elapsedRealtime() < deadline)
+        fail("Network $n LinkProperties did not $description after $DEFAULT_TIMEOUT_MS ms")
+    }
+
+    fun assertLinkPropertiesEventuallyNotNull(n: Network) {
+        assertLinkPropertiesEventually(n, "become non-null") { it != null }
+    }
+
+    fun assertLinkPropertiesEventuallyNull(n: Network) {
+        assertLinkPropertiesEventually(n, "become null") { it == null }
+    }
+
     @Test
     fun testSetSubtypeNameAndExtraInfoByAgentConfig() {
         val subtypeLTE = TelephonyManager.NETWORK_TYPE_LTE
@@ -381,7 +403,7 @@
     fun testConnectAndUnregister() {
         val (agent, callback) = createConnectedNetworkAgent()
         unregister(agent)
-        callback.expectCallback<Lost>(agent.network!!)
+        callback.expect<Lost>(agent.network!!)
         assertFailsWith<IllegalStateException>("Must not be able to register an agent twice") {
             agent.register()
         }
@@ -432,7 +454,7 @@
             agent.sendNetworkCapabilities(nc)
             callbacks[0].expectCapabilitiesThat(agent.network!!) { it.signalStrength == 55 }
             callbacks[1].expectCapabilitiesThat(agent.network!!) { it.signalStrength == 55 }
-            callbacks[2].expectCallback<Lost>(agent.network!!)
+            callbacks[2].expect<Lost>(agent.network!!)
         }
         callbacks.forEach {
             mCM.unregisterNetworkCallback(it)
@@ -510,12 +532,12 @@
         agent.markConnected()
 
         // Make sure the UIDs have been ignored.
-        callback.expectCallback<Available>(agent.network!!)
+        callback.expect<Available>(agent.network!!)
         callback.expectCapabilitiesThat(agent.network!!) {
             it.allowedUids.isEmpty() && !it.hasCapability(NET_CAPABILITY_VALIDATED)
         }
-        callback.expectCallback<LinkPropertiesChanged>(agent.network!!)
-        callback.expectCallback<BlockedStatus>(agent.network!!)
+        callback.expect<LinkPropertiesChanged>(agent.network!!)
+        callback.expect<BlockedStatus>(agent.network!!)
         callback.expectCapabilitiesThat(agent.network!!) {
             it.allowedUids.isEmpty() && it.hasCapability(NET_CAPABILITY_VALIDATED)
         }
@@ -555,7 +577,7 @@
                 .setLegacyInt(WORSE_NETWORK_SCORE)
                 .setExiting(true)
                 .build())
-        callback.expectCallback<Available>(agent2.network!!)
+        callback.expect<Available>(agent2.network!!)
 
         // tearDown() will unregister the requests and agents
     }
@@ -639,7 +661,7 @@
         }
 
         unregister(agent)
-        callback.expectCallback<Lost>(agent.network!!)
+        callback.expect<Lost>(agent.network!!)
     }
 
     private fun unregister(agent: TestableNetworkAgent) {
@@ -812,14 +834,14 @@
         agentStronger.register()
         agentStronger.markConnected()
         commonCallback.expectAvailableCallbacks(agentStronger.network!!)
-        callbackWeaker.expectCallback<Losing>(agentWeaker.network!!)
+        callbackWeaker.expect<Losing>(agentWeaker.network!!)
         val expectedRemainingLingerDuration = lingerStart +
                 NetworkAgent.MIN_LINGER_TIMER_MS.toLong() - SystemClock.elapsedRealtime()
         // If the available callback is too late. The remaining duration will be reduced.
         assertTrue(expectedRemainingLingerDuration > 0,
                 "expected remaining linger duration is $expectedRemainingLingerDuration")
         callbackWeaker.assertNoCallback(expectedRemainingLingerDuration)
-        callbackWeaker.expectCallback<Lost>(agentWeaker.network!!)
+        callbackWeaker.expect<Lost>(agentWeaker.network!!)
     }
 
     @Test
@@ -886,7 +908,7 @@
         bestMatchingCb.assertNoCallback()
         agent1.unregister()
         agentsToCleanUp.remove(agent1)
-        bestMatchingCb.expectCallback<Lost>(agent1.network!!)
+        bestMatchingCb.expect<Lost>(agent1.network!!)
 
         // tearDown() will unregister the requests and agents
     }
@@ -1200,7 +1222,7 @@
         // testCallback will not see any events. agent2 is be torn down because it has no requests.
         val (agent2, network2) = connectNetwork()
         matchAllCallback.expectAvailableThenValidatedCallbacks(network2)
-        matchAllCallback.expectCallback<Lost>(network2)
+        matchAllCallback.expect<Lost>(network2)
         agent2.expectCallback<OnNetworkUnwanted>()
         agent2.expectCallback<OnNetworkDestroyed>()
         assertNull(mCM.getLinkProperties(network2))
@@ -1228,7 +1250,7 @@
 
         // As soon as the replacement arrives, network1 is disconnected.
         // Check that this happens before the replacement timeout (5 seconds) fires.
-        matchAllCallback.expectCallback<Lost>(network1, 2_000 /* timeoutMs */)
+        matchAllCallback.expect<Lost>(network1, 2_000 /* timeoutMs */)
         agent1.expectCallback<OnNetworkUnwanted>()
 
         // Test lingering:
@@ -1241,19 +1263,19 @@
         val network4 = agent4.network!!
         matchAllCallback.expectAvailableThenValidatedCallbacks(network4)
         agent4.sendNetworkScore(NetworkScore.Builder().setTransportPrimary(true).build())
-        matchAllCallback.expectCallback<Losing>(network3)
+        matchAllCallback.expect<Losing>(network3)
         testCallback.expectAvailableCallbacks(network4, validated = true)
         mCM.unregisterNetworkCallback(agent4callback)
         agent3.unregisterAfterReplacement(5_000)
         agent3.expectCallback<OnNetworkUnwanted>()
-        matchAllCallback.expectCallback<Lost>(network3, 1000L)
+        matchAllCallback.expect<Lost>(network3, 1000L)
         agent3.expectCallback<OnNetworkDestroyed>()
 
         // Now mark network4 awaiting replacement with a low timeout, and check that if no
         // replacement arrives, it is torn down.
         agent4.unregisterAfterReplacement(100 /* timeoutMillis */)
-        matchAllCallback.expectCallback<Lost>(network4, 1000L /* timeoutMs */)
-        testCallback.expectCallback<Lost>(network4, 1000L /* timeoutMs */)
+        matchAllCallback.expect<Lost>(network4, 1000L /* timeoutMs */)
+        testCallback.expect<Lost>(network4, 1000L /* timeoutMs */)
         agent4.expectCallback<OnNetworkDestroyed>()
         agent4.expectCallback<OnNetworkUnwanted>()
 
@@ -1264,13 +1286,45 @@
         testCallback.expectAvailableThenValidatedCallbacks(network5)
         agent5.unregisterAfterReplacement(5_000 /* timeoutMillis */)
         agent5.unregister()
-        matchAllCallback.expectCallback<Lost>(network5, 1000L /* timeoutMs */)
-        testCallback.expectCallback<Lost>(network5, 1000L /* timeoutMs */)
+        matchAllCallback.expect<Lost>(network5, 1000L /* timeoutMs */)
+        testCallback.expect<Lost>(network5, 1000L /* timeoutMs */)
         agent5.expectCallback<OnNetworkDestroyed>()
         agent5.expectCallback<OnNetworkUnwanted>()
 
+        // If unregisterAfterReplacement is called before markConnected, the network disconnects.
+        val specifier6 = UUID.randomUUID().toString()
+        val callback = TestableNetworkCallback()
+        requestNetwork(makeTestNetworkRequest(specifier = specifier6), callback)
+        val agent6 = createNetworkAgent(specifier = specifier6)
+        val network6 = agent6.register()
+        // No callbacks are sent, so check the LinkProperties to see if the network has connected.
+        assertLinkPropertiesEventuallyNotNull(agent6.network!!)
+
+        // unregisterAfterReplacement tears down the network immediately.
+        // Approximately check that this is the case by picking an unregister timeout that's longer
+        // than the timeout of the expectCallback<OnNetworkUnwanted> below.
+        // TODO: consider adding configurable timeouts to TestableNetworkAgent expectations.
+        val timeoutMs = agent6.DEFAULT_TIMEOUT_MS.toInt() + 1_000
+        agent6.unregisterAfterReplacement(timeoutMs)
+        agent6.expectCallback<OnNetworkUnwanted>()
+        if (!SdkLevel.isAtLeastT()) {
+            // Before T, onNetworkDestroyed is called even if the network was never created.
+            agent6.expectCallback<OnNetworkDestroyed>()
+        }
+        // Poll for LinkProperties becoming null, because when onNetworkUnwanted is called, the
+        // network has not yet been removed from the CS data structures.
+        assertLinkPropertiesEventuallyNull(agent6.network!!)
+        assertFalse(mCM.getAllNetworks().contains(agent6.network!!))
+
+        // After unregisterAfterReplacement is called, the network is no longer usable and
+        // markConnected has no effect.
+        agent6.markConnected()
+        agent6.assertNoCallback()
+        assertNull(mCM.getLinkProperties(agent6.network!!))
+        matchAllCallback.assertNoCallback(200 /* timeoutMs */)
+
         // If wifi is replaced within the timeout, the device does not switch to cellular.
-        val (cellAgent, cellNetwork) = connectNetwork(TRANSPORT_CELLULAR)
+        val (_, cellNetwork) = connectNetwork(TRANSPORT_CELLULAR)
         testCallback.expectAvailableThenValidatedCallbacks(cellNetwork)
         matchAllCallback.expectAvailableThenValidatedCallbacks(cellNetwork)
 
@@ -1280,7 +1334,7 @@
             it.hasCapability(NET_CAPABILITY_VALIDATED)
         }
         matchAllCallback.expectAvailableCallbacks(wifiNetwork, validated = false)
-        matchAllCallback.expectCallback<Losing>(cellNetwork)
+        matchAllCallback.expect<Losing>(cellNetwork)
         matchAllCallback.expectCapabilitiesThat(wifiNetwork) {
             it.hasCapability(NET_CAPABILITY_VALIDATED)
         }
@@ -1309,7 +1363,7 @@
         val (newWifiAgent, newWifiNetwork) = connectNetwork(TRANSPORT_WIFI)
         testCallback.expectAvailableCallbacks(newWifiNetwork, validated = true)
         matchAllCallback.expectAvailableThenValidatedCallbacks(newWifiNetwork)
-        matchAllCallback.expectCallback<Lost>(wifiNetwork)
+        matchAllCallback.expect<Lost>(wifiNetwork)
         wifiAgent.expectCallback<OnNetworkUnwanted>()
     }
 
@@ -1328,7 +1382,7 @@
         // lost callback should be sent still.
         agent.markConnected()
         agent.unregister()
-        callback.expectCallback<Available>(agent.network!!)
+        callback.expect<Available>(agent.network!!)
         callback.eventuallyExpect<Lost> { it.network == agent.network }
     }
 }
diff --git a/tests/unit/Android.bp b/tests/unit/Android.bp
index cb68235..437622b 100644
--- a/tests/unit/Android.bp
+++ b/tests/unit/Android.bp
@@ -101,7 +101,7 @@
     ],
     static_libs: [
         "androidx.test.rules",
-        "androidx.test.uiautomator",
+        "androidx.test.uiautomator_uiautomator",
         "bouncycastle-repackaged-unbundled",
         "core-tests-support",
         "FrameworksNetCommonTests",
diff --git a/tests/unit/java/android/app/usage/NetworkStatsManagerTest.java b/tests/unit/java/android/app/usage/NetworkStatsManagerTest.java
index 8a537be..679427a 100644
--- a/tests/unit/java/android/app/usage/NetworkStatsManagerTest.java
+++ b/tests/unit/java/android/app/usage/NetworkStatsManagerTest.java
@@ -16,7 +16,16 @@
 
 package android.app.usage;
 
+import static android.net.ConnectivityManager.TYPE_MOBILE;
+import static android.net.ConnectivityManager.TYPE_WIFI;
+import static android.net.NetworkStats.DEFAULT_NETWORK_NO;
+import static android.net.NetworkStats.METERED_NO;
 import static android.net.NetworkStats.METERED_YES;
+import static android.net.NetworkStats.ROAMING_NO;
+import static android.net.NetworkStats.SET_ALL;
+import static android.net.NetworkStats.SET_DEFAULT;
+import static android.net.NetworkStats.TAG_NONE;
+import static android.net.NetworkStatsHistory.FIELD_ALL;
 import static android.net.NetworkTemplate.MATCH_MOBILE;
 import static android.net.NetworkTemplate.MATCH_WIFI;
 
@@ -34,7 +43,6 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import android.net.ConnectivityManager;
 import android.net.INetworkStatsService;
 import android.net.INetworkStatsSession;
 import android.net.NetworkStats.Entry;
@@ -86,31 +94,17 @@
         final int uid2 = 10002;
         final int uid3 = 10003;
 
-        Entry uid1Entry1 = new Entry("if1", uid1,
-                android.net.NetworkStats.SET_DEFAULT, android.net.NetworkStats.TAG_NONE,
-                android.net.NetworkStats.METERED_NO, android.net.NetworkStats.ROAMING_NO,
-                android.net.NetworkStats.DEFAULT_NETWORK_NO,
-                100, 10, 200, 20, 0);
+        Entry uid1Entry1 = new Entry("if1", uid1, SET_DEFAULT, TAG_NONE, METERED_NO, ROAMING_NO,
+                DEFAULT_NETWORK_NO, 100, 10, 200, 20, 0);
 
-        Entry uid1Entry2 = new Entry(
-                "if2", uid1,
-                android.net.NetworkStats.SET_DEFAULT, android.net.NetworkStats.TAG_NONE,
-                android.net.NetworkStats.METERED_NO, android.net.NetworkStats.ROAMING_NO,
-                android.net.NetworkStats.DEFAULT_NETWORK_NO,
-                100, 10, 200, 20, 0);
+        Entry uid1Entry2 = new Entry("if2", uid1, SET_DEFAULT, TAG_NONE, METERED_NO, ROAMING_NO,
+                DEFAULT_NETWORK_NO, 100, 10, 200, 20, 0);
 
-        Entry uid2Entry1 = new Entry("if1", uid2,
-                android.net.NetworkStats.SET_DEFAULT, android.net.NetworkStats.TAG_NONE,
-                android.net.NetworkStats.METERED_NO, android.net.NetworkStats.ROAMING_NO,
-                android.net.NetworkStats.DEFAULT_NETWORK_NO,
-                150, 10, 250, 20, 0);
+        Entry uid2Entry1 = new Entry("if1", uid2, SET_DEFAULT, TAG_NONE, METERED_NO, ROAMING_NO,
+                DEFAULT_NETWORK_NO, 150, 10, 250, 20, 0);
 
-        Entry uid2Entry2 = new Entry(
-                "if2", uid2,
-                android.net.NetworkStats.SET_DEFAULT, android.net.NetworkStats.TAG_NONE,
-                android.net.NetworkStats.METERED_NO, android.net.NetworkStats.ROAMING_NO,
-                android.net.NetworkStats.DEFAULT_NETWORK_NO,
-                150, 10, 250, 20, 0);
+        Entry uid2Entry2 = new Entry("if2", uid2, SET_DEFAULT, TAG_NONE, METERED_NO, ROAMING_NO,
+                DEFAULT_NETWORK_NO, 150, 10, 250, 20, 0);
 
         NetworkStatsHistory history1 = new NetworkStatsHistory(10, 2);
         history1.recordData(10, 20, uid1Entry1);
@@ -125,9 +119,8 @@
         when(mStatsSession.getRelevantUids()).thenReturn(new int[] { uid1, uid2, uid3 });
 
         when(mStatsSession.getHistoryIntervalForUid(any(NetworkTemplate.class),
-                eq(uid1), eq(android.net.NetworkStats.SET_ALL),
-                eq(android.net.NetworkStats.TAG_NONE),
-                eq(NetworkStatsHistory.FIELD_ALL), eq(startTime), eq(endTime)))
+                eq(uid1), eq(SET_ALL), eq(TAG_NONE),
+                eq(FIELD_ALL), eq(startTime), eq(endTime)))
                 .then((InvocationOnMock inv) -> {
                     NetworkTemplate template = inv.getArgument(0);
                     assertEquals(MATCH_MOBILE_ALL, template.getMatchRule());
@@ -136,9 +129,8 @@
                 });
 
         when(mStatsSession.getHistoryIntervalForUid(any(NetworkTemplate.class),
-                eq(uid2), eq(android.net.NetworkStats.SET_ALL),
-                eq(android.net.NetworkStats.TAG_NONE),
-                eq(NetworkStatsHistory.FIELD_ALL), eq(startTime), eq(endTime)))
+                eq(uid2), eq(SET_ALL), eq(TAG_NONE),
+                eq(FIELD_ALL), eq(startTime), eq(endTime)))
                 .then((InvocationOnMock inv) -> {
                     NetworkTemplate template = inv.getArgument(0);
                     assertEquals(MATCH_MOBILE_ALL, template.getMatchRule());
@@ -148,7 +140,7 @@
 
 
         NetworkStats stats = mManager.queryDetails(
-                ConnectivityManager.TYPE_MOBILE, TEST_SUBSCRIBER_ID, startTime, endTime);
+                TYPE_MOBILE, TEST_SUBSCRIBER_ID, startTime, endTime);
 
         NetworkStats.Bucket bucket = new NetworkStats.Bucket();
 
@@ -202,35 +194,34 @@
 
         verify(mStatsSession, times(1)).getHistoryIntervalForUid(
                 eq(expectedTemplate),
-                eq(uid1), eq(android.net.NetworkStats.SET_ALL),
-                eq(android.net.NetworkStats.TAG_NONE),
-                eq(NetworkStatsHistory.FIELD_ALL), eq(startTime), eq(endTime));
+                eq(uid1), eq(SET_ALL),
+                eq(TAG_NONE),
+                eq(FIELD_ALL), eq(startTime), eq(endTime));
 
         verify(mStatsSession, times(1)).getHistoryIntervalForUid(
                 eq(expectedTemplate),
-                eq(uid2), eq(android.net.NetworkStats.SET_ALL),
-                eq(android.net.NetworkStats.TAG_NONE),
-                eq(NetworkStatsHistory.FIELD_ALL), eq(startTime), eq(endTime));
+                eq(uid2), eq(SET_ALL),
+                eq(TAG_NONE),
+                eq(FIELD_ALL), eq(startTime), eq(endTime));
 
         assertFalse(stats.hasNextBucket());
     }
 
     @Test
     public void testNetworkTemplateWhenRunningQueryDetails_NoSubscriberId() throws RemoteException {
-        runQueryDetailsAndCheckTemplate(ConnectivityManager.TYPE_MOBILE,
-                null /* subscriberId */, new NetworkTemplate.Builder(MATCH_MOBILE)
-                        .setMeteredness(METERED_YES).build());
-        runQueryDetailsAndCheckTemplate(ConnectivityManager.TYPE_WIFI,
-                "" /* subscriberId */, new NetworkTemplate.Builder(MATCH_WIFI).build());
-        runQueryDetailsAndCheckTemplate(ConnectivityManager.TYPE_WIFI,
-                null /* subscriberId */, new NetworkTemplate.Builder(MATCH_WIFI).build());
+        runQueryDetailsAndCheckTemplate(TYPE_MOBILE, null /* subscriberId */,
+                new NetworkTemplate.Builder(MATCH_MOBILE).setMeteredness(METERED_YES).build());
+        runQueryDetailsAndCheckTemplate(TYPE_WIFI, "" /* subscriberId */,
+                new NetworkTemplate.Builder(MATCH_WIFI).build());
+        runQueryDetailsAndCheckTemplate(TYPE_WIFI, null /* subscriberId */,
+                new NetworkTemplate.Builder(MATCH_WIFI).build());
     }
 
     @Test
     public void testNetworkTemplateWhenRunningQueryDetails_MergedCarrierWifi()
             throws RemoteException {
-        runQueryDetailsAndCheckTemplate(ConnectivityManager.TYPE_WIFI,
-                TEST_SUBSCRIBER_ID, new NetworkTemplate.Builder(MATCH_WIFI)
+        runQueryDetailsAndCheckTemplate(TYPE_WIFI, TEST_SUBSCRIBER_ID,
+                new NetworkTemplate.Builder(MATCH_WIFI)
                         .setSubscriberIds(Set.of(TEST_SUBSCRIBER_ID)).build());
     }
 
@@ -244,7 +235,7 @@
         when(mStatsSession.getTaggedSummaryForAllUid(any(NetworkTemplate.class),
                 anyLong(), anyLong()))
                 .thenReturn(new android.net.NetworkStats(0, 0));
-        final NetworkTemplate template = new NetworkTemplate.Builder(NetworkTemplate.MATCH_MOBILE)
+        final NetworkTemplate template = new NetworkTemplate.Builder(MATCH_MOBILE)
                 .setMeteredness(NetworkStats.Bucket.METERED_YES).build();
         NetworkStats stats = mManager.queryTaggedSummary(template, startTime, endTime);
 
@@ -265,12 +256,12 @@
         when(mStatsSession.getHistoryIntervalForNetwork(any(NetworkTemplate.class),
                 anyInt(), anyLong(), anyLong()))
                 .thenReturn(new NetworkStatsHistory(10, 0));
-        final NetworkTemplate template = new NetworkTemplate.Builder(NetworkTemplate.MATCH_MOBILE)
+        final NetworkTemplate template = new NetworkTemplate.Builder(MATCH_MOBILE)
                 .setMeteredness(NetworkStats.Bucket.METERED_YES).build();
         NetworkStats stats = mManager.queryDetailsForDevice(template, startTime, endTime);
 
         verify(mStatsSession, times(1)).getHistoryIntervalForNetwork(
-                eq(template), eq(NetworkStatsHistory.FIELD_ALL), eq(startTime), eq(endTime));
+                eq(template), eq(FIELD_ALL), eq(startTime), eq(endTime));
 
         assertFalse(stats.hasNextBucket());
     }
diff --git a/tests/unit/java/android/net/NetworkStatsTest.java b/tests/unit/java/android/net/NetworkStatsTest.java
index 709b722..126ad55 100644
--- a/tests/unit/java/android/net/NetworkStatsTest.java
+++ b/tests/unit/java/android/net/NetworkStatsTest.java
@@ -1067,6 +1067,38 @@
         }
     }
 
+    @Test
+    public void testClearInterfaces() {
+        final NetworkStats stats = new NetworkStats(TEST_START, 1);
+        final NetworkStats.Entry entry1 = new NetworkStats.Entry(
+                "test1", 10100, SET_DEFAULT, TAG_NONE, METERED_NO, ROAMING_NO,
+                DEFAULT_NETWORK_NO, 1024L, 50L, 100L, 20L, 0L);
+
+        final NetworkStats.Entry entry2 = new NetworkStats.Entry(
+                "test2", 10101, SET_DEFAULT, 0xF0DD, METERED_NO, ROAMING_NO,
+                DEFAULT_NETWORK_NO, 51200, 25L, 101010L, 50L, 0L);
+
+        stats.insertEntry(entry1);
+        stats.insertEntry(entry2);
+
+        // Verify that the interfaces have indeed been recorded.
+        assertEquals(2, stats.size());
+        assertValues(stats, 0, "test1", 10100, SET_DEFAULT, TAG_NONE, METERED_NO,
+                ROAMING_NO, DEFAULT_NETWORK_NO, 1024L, 50L, 100L, 20L, 0L);
+        assertValues(stats, 1, "test2", 10101, SET_DEFAULT, 0xF0DD, METERED_NO,
+                ROAMING_NO, DEFAULT_NETWORK_NO, 51200, 25L, 101010L, 50L, 0L);
+
+        // Clear interfaces.
+        stats.clearInterfaces();
+
+        // Verify that the interfaces are cleared.
+        assertEquals(2, stats.size());
+        assertValues(stats, 0, null /* iface */, 10100, SET_DEFAULT, TAG_NONE, METERED_NO,
+                ROAMING_NO, DEFAULT_NETWORK_NO, 1024L, 50L, 100L, 20L, 0L);
+        assertValues(stats, 1, null /* iface */, 10101, SET_DEFAULT, 0xF0DD, METERED_NO,
+                ROAMING_NO, DEFAULT_NETWORK_NO, 51200, 25L, 101010L, 50L, 0L);
+    }
+
     private static void assertContains(NetworkStats stats,  String iface, int uid, int set,
             int tag, int metered, int roaming, int defaultNetwork, long rxBytes, long rxPackets,
             long txBytes, long txPackets, long operations) {
diff --git a/tests/unit/java/com/android/server/BpfNetMapsTest.java b/tests/unit/java/com/android/server/BpfNetMapsTest.java
index 0c00bc0..0e17cd7 100644
--- a/tests/unit/java/com/android/server/BpfNetMapsTest.java
+++ b/tests/unit/java/com/android/server/BpfNetMapsTest.java
@@ -1070,4 +1070,11 @@
         doTestDumpOwnerMatchConfig(DOZABLE_MATCH | invalid_match,
                 "DOZABLE_MATCH UNKNOWN_MATCH(" + invalid_match + ")");
     }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testDumpCookieTagMap() throws Exception {
+        mCookieTagMap.updateEntry(new CookieTagMapKey(123), new CookieTagMapValue(456, 0x789));
+        assertDumpContains(getDump(), "cookie=123 tag=0x789 uid=456");
+    }
 }
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index 86b5241..2d8cf80 100755
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -441,8 +441,6 @@
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
-import kotlin.reflect.KClass;
-
 /**
  * Tests for {@link ConnectivityService}.
  *
@@ -2338,7 +2336,7 @@
         b = registerConnectivityBroadcast(1);
         final TestNetworkCallback callback = new TestNetworkCallback();
         mCm.requestNetwork(legacyRequest, callback);
-        callback.expectCallback(CallbackEntry.AVAILABLE, mCellNetworkAgent);
+        callback.expect(CallbackEntry.AVAILABLE, mCellNetworkAgent);
         mCm.unregisterNetworkCallback(callback);
         b.expectNoBroadcast(800);  // 800ms long enough to at least flake if this is sent
 
@@ -2410,7 +2408,7 @@
         // is added in case of flakiness.
         final int nascentTimeoutMs =
                 mService.mNascentDelayMs + mService.mNascentDelayMs / 4;
-        listenCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent, nascentTimeoutMs);
+        listenCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent, nascentTimeoutMs);
 
         // 2. Create a network that is satisfied by a request comes later.
         mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
@@ -2426,7 +2424,7 @@
         // to get disconnected as usual if the request is released after the nascent timer expires.
         listenCallback.assertNoCallback(nascentTimeoutMs);
         mCm.unregisterNetworkCallback(wifiCallback);
-        listenCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        listenCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
 
         // 3. Create a network that is satisfied by a request comes later.
         mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
@@ -2437,7 +2435,7 @@
 
         // Verify that the network will still be torn down after the request gets removed.
         mCm.unregisterNetworkCallback(wifiCallback);
-        listenCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        listenCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
 
         // There is no need to ensure that LOSING is never sent in the common case that the
         // network immediately satisfies a request that was already present, because it is already
@@ -2483,7 +2481,7 @@
         assertFalse(isForegroundNetwork(mCellNetworkAgent));
 
         mCellNetworkAgent.disconnect();
-        bgMobileListenCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        bgMobileListenCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         fgMobileListenCallback.assertNoCallback();
 
         mCm.unregisterNetworkCallback(wifiListenCallback);
@@ -2711,7 +2709,7 @@
         // Make sure the default request goes to net 2
         generalCb.expectAvailableCallbacksUnvalidated(net2);
         if (expectLingering) {
-            generalCb.expectCallback(CallbackEntry.LOSING, net1);
+            generalCb.expectLosing(net1);
         }
         generalCb.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, net2);
         defaultCb.expectAvailableDoubleValidatedCallbacks(net2);
@@ -2728,7 +2726,7 @@
             net1.expectDisconnected(TEST_CALLBACK_TIMEOUT_MS);
         }
         net1.disconnect();
-        generalCb.expectCallback(CallbackEntry.LOST, net1);
+        generalCb.expect(CallbackEntry.LOST, net1);
 
         // Remove primary from net 2
         net2.setScore(new NetworkScore.Builder().build());
@@ -2756,7 +2754,7 @@
             // get LOSING. If the radio can't time share, this is a hard loss, since the last
             // request keeping up this network has been removed and the network isn't lingering
             // for any other request.
-            generalCb.expectCallback(CallbackEntry.LOSING, net2);
+            generalCb.expectLosing(net2);
             net2.assertNotDisconnected(TEST_CALLBACK_TIMEOUT_MS);
             generalCb.assertNoCallback();
             net2.expectDisconnected(UNREASONABLY_LONG_ALARM_WAIT_MS);
@@ -2764,7 +2762,7 @@
             net2.expectDisconnected(TEST_CALLBACK_TIMEOUT_MS);
         }
         net2.disconnect();
-        generalCb.expectCallback(CallbackEntry.LOST, net2);
+        generalCb.expect(CallbackEntry.LOST, net2);
         defaultCb.assertNoCallback();
 
         net3.disconnect();
@@ -2952,7 +2950,7 @@
 
     /**
      * Utility NetworkCallback for testing. The caller must explicitly test for all the callbacks
-     * this class receives, by calling expectCallback() exactly once each time a callback is
+     * this class receives, by calling expect() exactly once each time a callback is
      * received. assertNoCallback may be called at any time.
      */
     private class TestNetworkCallback extends TestableNetworkCallback {
@@ -2967,20 +2965,24 @@
             assertNoCallback(0 /* timeout */);
         }
 
-        @Override
-        public <T extends CallbackEntry> T expectCallback(final KClass<T> type, final HasNetwork n,
-                final long timeoutMs) {
-            final T callback = super.expectCallback(type, n, timeoutMs);
-            if (callback instanceof CallbackEntry.Losing) {
-                // TODO : move this to the specific test(s) needing this rather than here.
-                final CallbackEntry.Losing losing = (CallbackEntry.Losing) callback;
-                final int maxMsToLive = losing.getMaxMsToLive();
-                String msg = String.format(
-                        "Invalid linger time value %d, must be between %d and %d",
-                        maxMsToLive, 0, mService.mLingerDelayMs);
-                assertTrue(msg, 0 <= maxMsToLive && maxMsToLive <= mService.mLingerDelayMs);
+        public CallbackEntry.Losing expectLosing(final HasNetwork n, final long timeoutMs) {
+            final CallbackEntry.Losing losing = expect(CallbackEntry.LOSING, n, timeoutMs);
+            final int maxMsToLive = losing.getMaxMsToLive();
+            if (maxMsToLive < 0 || maxMsToLive > mService.mLingerDelayMs) {
+                // maxMsToLive is the value that was received in the onLosing callback. That must
+                // not be negative, so check that.
+                // Also, maxMsToLive is the remaining time until the network expires.
+                // mService.mLingerDelayMs is how long the network takes from when it's first
+                // detected to be unneeded to when it expires, so maxMsToLive should never
+                // be greater than that.
+                fail(String.format("Invalid linger time value %d, must be between %d and %d",
+                        maxMsToLive, 0, mService.mLingerDelayMs));
             }
-            return callback;
+            return losing;
+        }
+
+        public CallbackEntry.Losing expectLosing(final HasNetwork n) {
+            return expectLosing(n, getDefaultTimeoutMs());
         }
     }
 
@@ -2994,19 +2996,19 @@
 
     static void expectOnLost(TestNetworkAgentWrapper network, TestNetworkCallback ... callbacks) {
         for (TestNetworkCallback c : callbacks) {
-            c.expectCallback(CallbackEntry.LOST, network);
+            c.expect(CallbackEntry.LOST, network);
         }
     }
 
     static void expectAvailableCallbacksUnvalidatedWithSpecifier(TestNetworkAgentWrapper network,
             NetworkSpecifier specifier, TestNetworkCallback ... callbacks) {
         for (TestNetworkCallback c : callbacks) {
-            c.expectCallback(CallbackEntry.AVAILABLE, network);
+            c.expect(CallbackEntry.AVAILABLE, network);
             c.expectCapabilitiesThat(network, (nc) ->
                     !nc.hasCapability(NET_CAPABILITY_VALIDATED)
                             && Objects.equals(specifier, nc.getNetworkSpecifier()));
-            c.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, network);
-            c.expectCallback(CallbackEntry.BLOCKED_STATUS, network);
+            c.expect(CallbackEntry.LINK_PROPERTIES_CHANGED, network);
+            c.expect(CallbackEntry.BLOCKED_STATUS, network);
         }
     }
 
@@ -3089,16 +3091,16 @@
 
         b = registerConnectivityBroadcast(2);
         mWiFiNetworkAgent.disconnect();
-        genericNetworkCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
-        wifiNetworkCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        genericNetworkCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
+        wifiNetworkCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         cellNetworkCallback.assertNoCallback();
         b.expectBroadcast();
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
 
         b = registerConnectivityBroadcast(1);
         mCellNetworkAgent.disconnect();
-        genericNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        genericNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
+        cellNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         b.expectBroadcast();
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
 
@@ -3119,21 +3121,21 @@
         mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
         genericNetworkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        genericNetworkCallback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        genericNetworkCallback.expectLosing(mCellNetworkAgent);
         genericNetworkCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         wifiNetworkCallback.expectAvailableThenValidatedCallbacks(mWiFiNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        cellNetworkCallback.expectLosing(mCellNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
 
         mWiFiNetworkAgent.disconnect();
-        genericNetworkCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
-        wifiNetworkCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        genericNetworkCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
+        wifiNetworkCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
 
         mCellNetworkAgent.disconnect();
-        genericNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        genericNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
+        cellNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
     }
 
@@ -3266,7 +3268,7 @@
         // We then get LOSING when wifi validates and cell is outscored.
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         // TODO: Investigate sending validated before losing.
-        callback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        callback.expectLosing(mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
@@ -3275,15 +3277,15 @@
         mEthernetNetworkAgent.connect(true);
         callback.expectAvailableCallbacksUnvalidated(mEthernetNetworkAgent);
         // TODO: Investigate sending validated before losing.
-        callback.expectCallback(CallbackEntry.LOSING, mWiFiNetworkAgent);
+        callback.expectLosing(mWiFiNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mEthernetNetworkAgent);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mEthernetNetworkAgent);
         assertEquals(mEthernetNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         mEthernetNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mEthernetNetworkAgent);
-        defaultCallback.expectCallback(CallbackEntry.LOST, mEthernetNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mEthernetNetworkAgent);
+        defaultCallback.expect(CallbackEntry.LOST, mEthernetNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
@@ -3299,7 +3301,7 @@
                 newNetwork = mWiFiNetworkAgent;
 
             }
-            callback.expectCallback(CallbackEntry.LOSING, oldNetwork);
+            callback.expectLosing(oldNetwork);
             // TODO: should we send an AVAILABLE callback to newNetwork, to indicate that it is no
             // longer lingering?
             defaultCallback.expectAvailableCallbacksValidated(newNetwork);
@@ -3313,7 +3315,7 @@
         // We expect a notification about the capabilities change, and nothing else.
         defaultCallback.expectCapabilitiesWithout(NET_CAPABILITY_NOT_METERED, mWiFiNetworkAgent);
         defaultCallback.assertNoCallback();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Wifi no longer satisfies our listen, which is for an unmetered network.
@@ -3322,11 +3324,11 @@
 
         // Disconnect our test networks.
         mWiFiNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        defaultCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
         mCellNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        defaultCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         waitForIdle();
         assertEquals(null, mCm.getActiveNetwork());
 
@@ -3357,8 +3359,8 @@
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
-        defaultCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
+        defaultCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
@@ -3369,19 +3371,19 @@
         mWiFiNetworkAgent.connect(true);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         // TODO: Investigate sending validated before losing.
-        callback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        callback.expectLosing(mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         defaultCallback.expectAvailableThenValidatedCallbacks(mWiFiNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
-        defaultCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
+        defaultCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
         mCellNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
-        defaultCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mCellNetworkAgent);
+        defaultCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         waitForIdle();
         assertEquals(null, mCm.getActiveNetwork());
 
@@ -3396,7 +3398,7 @@
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         // TODO: Investigate sending validated before losing.
-        callback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        callback.expectLosing(mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
@@ -3407,13 +3409,13 @@
         // TODO: should this cause an AVAILABLE callback, to indicate that the network is no longer
         // lingering?
         mCm.unregisterNetworkCallback(noopCallback);
-        callback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        callback.expectLosing(mCellNetworkAgent);
 
         // Similar to the above: lingering can start even after the lingered request is removed.
         // Disconnect wifi and switch to cell.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
-        defaultCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
+        defaultCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
@@ -3432,12 +3434,12 @@
         callback.assertNoCallback();
         // Now unregister cellRequest and expect cell to start lingering.
         mCm.unregisterNetworkCallback(noopCallback);
-        callback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        callback.expectLosing(mCellNetworkAgent);
 
         // Let linger run its course.
         callback.assertNoCallback();
         final int lingerTimeoutMs = mService.mLingerDelayMs + mService.mLingerDelayMs / 4;
-        callback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent, lingerTimeoutMs);
+        callback.expect(CallbackEntry.LOST, mCellNetworkAgent, lingerTimeoutMs);
 
         // Register a TRACK_DEFAULT request and check that it does not affect lingering.
         TestNetworkCallback trackDefaultCallback = new TestNetworkCallback();
@@ -3446,20 +3448,20 @@
         mEthernetNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_ETHERNET);
         mEthernetNetworkAgent.connect(true);
         callback.expectAvailableCallbacksUnvalidated(mEthernetNetworkAgent);
-        callback.expectCallback(CallbackEntry.LOSING, mWiFiNetworkAgent);
+        callback.expectLosing(mWiFiNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mEthernetNetworkAgent);
         trackDefaultCallback.expectAvailableDoubleValidatedCallbacks(mEthernetNetworkAgent);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mEthernetNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Let linger run its course.
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent, lingerTimeoutMs);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent, lingerTimeoutMs);
 
         // Clean up.
         mEthernetNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mEthernetNetworkAgent);
-        defaultCallback.expectCallback(CallbackEntry.LOST, mEthernetNetworkAgent);
-        trackDefaultCallback.expectCallback(CallbackEntry.LOST, mEthernetNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mEthernetNetworkAgent);
+        defaultCallback.expect(CallbackEntry.LOST, mEthernetNetworkAgent);
+        trackDefaultCallback.expect(CallbackEntry.LOST, mEthernetNetworkAgent);
 
         mCm.unregisterNetworkCallback(callback);
         mCm.unregisterNetworkCallback(defaultCallback);
@@ -3501,7 +3503,7 @@
         mWiFiNetworkAgent.connect(true);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        callback.expectLosing(mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
 
         // File a request for cellular, then release it.
@@ -3510,7 +3512,7 @@
         NetworkCallback noopCallback = new NetworkCallback();
         mCm.requestNetwork(cellRequest, noopCallback);
         mCm.unregisterNetworkCallback(noopCallback);
-        callback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        callback.expectLosing(mCellNetworkAgent);
 
         // Let linger run its course.
         callback.assertNoCallback();
@@ -3571,7 +3573,7 @@
 
     private void expectDisconnectAndClearNotifications(TestNetworkCallback callback,
             TestNetworkAgentWrapper agent, NotificationType type) {
-        callback.expectCallback(CallbackEntry.LOST, agent);
+        callback.expect(CallbackEntry.LOST, agent);
         expectClearNotification(agent, type);
     }
 
@@ -3710,7 +3712,7 @@
         // If the user chooses yes on the "No Internet access, stay connected?" dialog, we switch to
         // wifi even though it's unvalidated.
         mCm.setAcceptUnvalidated(mWiFiNetworkAgent.getNetwork(), true, false);
-        callback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        callback.expectLosing(mCellNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
 
         // Disconnect wifi, and then reconnect, again with explicitlySelected=true.
@@ -3739,7 +3741,7 @@
         mWiFiNetworkAgent.explicitlySelected(true, false);
         mWiFiNetworkAgent.connect(true);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        callback.expectLosing(mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         expectUnvalidationCheckWillNotNotify(mWiFiNetworkAgent);
@@ -3747,7 +3749,7 @@
         mEthernetNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_ETHERNET);
         mEthernetNetworkAgent.connect(true);
         callback.expectAvailableCallbacksUnvalidated(mEthernetNetworkAgent);
-        callback.expectCallback(CallbackEntry.LOSING, mWiFiNetworkAgent);
+        callback.expectLosing(mWiFiNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mEthernetNetworkAgent);
         assertEquals(mEthernetNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         callback.assertNoCallback();
@@ -3756,21 +3758,21 @@
         // (i.e., with explicitlySelected=true and acceptUnvalidated=true). Expect to switch to
         // wifi immediately.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.explicitlySelected(true, true);
         mWiFiNetworkAgent.connect(false);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackEntry.LOSING, mEthernetNetworkAgent);
+        callback.expectLosing(mEthernetNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         mEthernetNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mEthernetNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mEthernetNetworkAgent);
         expectUnvalidationCheckWillNotNotify(mWiFiNetworkAgent);
 
         // Disconnect and reconnect with explicitlySelected=false and acceptUnvalidated=true.
         // Check that the network is not scored specially and that the device prefers cell data.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
 
         mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.explicitlySelected(false, true);
@@ -3783,8 +3785,8 @@
         mWiFiNetworkAgent.disconnect();
         mCellNetworkAgent.disconnect();
 
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
-        callback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mCellNetworkAgent);
     }
 
     private void doTestFirstEvaluation(
@@ -3802,14 +3804,14 @@
         doConnect.accept(mWiFiNetworkAgent);
         // Expect the available callbacks, but don't require specific values for their arguments
         // since this method doesn't know how the network was connected.
-        callback.expectCallback(CallbackEntry.AVAILABLE, mWiFiNetworkAgent);
-        callback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED, mWiFiNetworkAgent);
-        callback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mWiFiNetworkAgent);
-        callback.expectCallback(CallbackEntry.BLOCKED_STATUS, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.AVAILABLE, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.NETWORK_CAPS_UPDATED, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.BLOCKED_STATUS, mWiFiNetworkAgent);
         if (waitForSecondCaps) {
             // This is necessary because of b/245893397, the same bug that happens where we use
             // expectAvailableDoubleValidatedCallbacks.
-            callback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED, mWiFiNetworkAgent);
+            callback.expect(CallbackEntry.NETWORK_CAPS_UPDATED, mWiFiNetworkAgent);
         }
         final NetworkAgentInfo nai =
                 mService.getNetworkAgentInfoForNetwork(mWiFiNetworkAgent.getNetwork());
@@ -3827,7 +3829,7 @@
             assertNotEquals(0L, nai.getFirstEvaluationConcludedTime());
         }
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
 
         mCm.unregisterNetworkCallback(callback);
     }
@@ -4223,7 +4225,7 @@
         // Need a trigger point to let NetworkMonitor tell ConnectivityService that the network is
         // validated.
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), true);
-        callback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        callback.expectLosing(mCellNetworkAgent);
         NetworkCapabilities nc = callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED,
                 mWiFiNetworkAgent);
         assertTrue(nc.hasCapability(NET_CAPABILITY_PARTIAL_CONNECTIVITY));
@@ -4234,7 +4236,7 @@
 
         // Disconnect and reconnect wifi with partial connectivity again.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
 
         mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connectWithPartialConnectivity();
@@ -4252,7 +4254,7 @@
         // If the user chooses no, disconnect wifi immediately.
         mCm.setAcceptPartialConnectivity(mWiFiNetworkAgent.getNetwork(), false /* accept */,
                 false /* always */);
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         expectClearNotification(mWiFiNetworkAgent, NotificationType.PARTIAL_CONNECTIVITY);
         reset(mNotificationManager);
 
@@ -4271,14 +4273,14 @@
         // ConnectivityService#updateNetworkInfo().
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
-        callback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        callback.expectLosing(mCellNetworkAgent);
         nc = callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         assertFalse(nc.hasCapability(NET_CAPABILITY_PARTIAL_CONNECTIVITY));
 
         // Wifi should be the default network.
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
 
         // The user accepted partial connectivity and selected "don't ask again". Now the user
         // reconnects to the partial connectivity network. Switch to wifi as soon as partial
@@ -4292,7 +4294,7 @@
         // ConnectivityService#updateNetworkInfo().
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
-        callback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        callback.expectLosing(mCellNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         callback.expectCapabilitiesWith(NET_CAPABILITY_PARTIAL_CONNECTIVITY, mWiFiNetworkAgent);
         expectUnvalidationCheckWillNotNotify(mWiFiNetworkAgent);
@@ -4304,7 +4306,7 @@
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), true);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
 
         // If the user accepted partial connectivity, and the device auto-reconnects to the partial
         // connectivity network, it should contain both PARTIAL_CONNECTIVITY and VALIDATED.
@@ -4318,12 +4320,12 @@
         mWiFiNetworkAgent.connectWithPartialValidConnectivity(false /* isStrictMode */);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
-        callback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        callback.expectLosing(mCellNetworkAgent);
         callback.expectCapabilitiesWith(
                 NET_CAPABILITY_PARTIAL_CONNECTIVITY | NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         expectUnvalidationCheckWillNotNotify(mWiFiNetworkAgent);
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         verifyNoMoreInteractions(mNotificationManager);
     }
 
@@ -4406,7 +4408,7 @@
         // Take down network.
         // Expect onLost callback.
         mWiFiNetworkAgent.disconnect();
-        captivePortalCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        captivePortalCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
 
         // Bring up a network with a captive portal.
         // Expect onAvailable callback of listen for NET_CAPABILITY_CAPTIVE_PORTAL.
@@ -4420,7 +4422,7 @@
         // Expect onLost callback because network no longer provides NET_CAPABILITY_CAPTIVE_PORTAL.
         mWiFiNetworkAgent.setNetworkValid(false /* isStrictMode */);
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), true);
-        captivePortalCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        captivePortalCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
 
         // Expect NET_CAPABILITY_VALIDATED onAvailable callback.
         validatedCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
@@ -4429,7 +4431,7 @@
         // Expect NET_CAPABILITY_VALIDATED onLost callback.
         mWiFiNetworkAgent.setNetworkInvalid(false /* isStrictMode */);
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), false);
-        validatedCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        validatedCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
     }
 
     @Test
@@ -4461,11 +4463,11 @@
         mWiFiNetworkAgent.setNetworkPortal("http://example.com", false /* isStrictMode */);
         mCm.reportNetworkConnectivity(wifiNetwork, false);
         captivePortalCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        validatedCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        validatedCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         // This is necessary because of b/245893397, the same bug that happens where we use
         // expectAvailableDoubleValidatedCallbacks.
         // TODO : fix b/245893397 and remove this.
-        captivePortalCallback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED,
+        captivePortalCallback.expect(CallbackEntry.NETWORK_CAPS_UPDATED,
                 mWiFiNetworkAgent);
 
         // Check that startCaptivePortalApp sends the expected command to NetworkMonitor.
@@ -4489,7 +4491,7 @@
         mWiFiNetworkAgent.setNetworkValid(false /* isStrictMode */);
         mWiFiNetworkAgent.mNetworkMonitor.forceReevaluation(Process.myUid());
         validatedCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
-        captivePortalCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        captivePortalCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
 
         mCm.unregisterNetworkCallback(validatedCallback);
         mCm.unregisterNetworkCallback(captivePortalCallback);
@@ -5030,7 +5032,7 @@
 
         // Bring down cell. Expect no default network callback, since it wasn't the default.
         mCellNetworkAgent.disconnect();
-        cellNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        cellNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         defaultNetworkCallback.assertNoCallback();
         systemDefaultCallback.assertNoCallback();
         assertEquals(defaultNetworkCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
@@ -5049,14 +5051,14 @@
         // followed by AVAILABLE cell.
         mWiFiNetworkAgent.disconnect();
         cellNetworkCallback.assertNoCallback();
-        defaultNetworkCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        defaultNetworkCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         defaultNetworkCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
-        systemDefaultCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        systemDefaultCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         systemDefaultCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         mCellNetworkAgent.disconnect();
-        cellNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
-        defaultNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
-        systemDefaultCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        cellNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
+        defaultNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
+        systemDefaultCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         waitForIdle();
         assertEquals(null, mCm.getActiveNetwork());
 
@@ -5068,7 +5070,7 @@
         assertEquals(null, systemDefaultCallback.getLastAvailableNetwork());
 
         mMockVpn.disconnect();
-        defaultNetworkCallback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        defaultNetworkCallback.expect(CallbackEntry.LOST, mMockVpn);
         systemDefaultCallback.assertNoCallback();
         waitForIdle();
         assertEquals(null, mCm.getActiveNetwork());
@@ -5097,7 +5099,7 @@
         lp.setInterfaceName("foonet_data0");
         mCellNetworkAgent.sendLinkProperties(lp);
         // We should get onLinkPropertiesChanged().
-        cellNetworkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED,
+        cellNetworkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
 
@@ -5105,7 +5107,7 @@
         mCellNetworkAgent.suspend();
         cellNetworkCallback.expectCapabilitiesWithout(NET_CAPABILITY_NOT_SUSPENDED,
                 mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackEntry.SUSPENDED, mCellNetworkAgent);
+        cellNetworkCallback.expect(CallbackEntry.SUSPENDED, mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
         assertEquals(NetworkInfo.State.SUSPENDED, mCm.getActiveNetworkInfo().getState());
 
@@ -5121,7 +5123,7 @@
         mCellNetworkAgent.resume();
         cellNetworkCallback.expectCapabilitiesWith(NET_CAPABILITY_NOT_SUSPENDED,
                 mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackEntry.RESUMED, mCellNetworkAgent);
+        cellNetworkCallback.expect(CallbackEntry.RESUMED, mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
         assertEquals(NetworkInfo.State.CONNECTED, mCm.getActiveNetworkInfo().getState());
 
@@ -5192,9 +5194,9 @@
         otherUidCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         includeOtherUidsCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
-        otherUidCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
-        includeOtherUidsCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
+        otherUidCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
+        includeOtherUidsCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
 
         // Only the includeOtherUidsCallback sees a VPN that does not apply to its UID.
         final UidRange range = UidRange.createForUser(UserHandle.of(RESTRICTED_USER));
@@ -5205,7 +5207,7 @@
         otherUidCallback.assertNoCallback();
 
         mMockVpn.disconnect();
-        includeOtherUidsCallback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        includeOtherUidsCallback.expect(CallbackEntry.LOST, mMockVpn);
         callback.assertNoCallback();
         otherUidCallback.assertNoCallback();
     }
@@ -5339,10 +5341,10 @@
 
         // When wifi connects, cell lingers.
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        callback.expectLosing(mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         fgCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        fgCallback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        fgCallback.expectLosing(mCellNetworkAgent);
         fgCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         assertTrue(isForegroundNetwork(mCellNetworkAgent));
         assertTrue(isForegroundNetwork(mWiFiNetworkAgent));
@@ -5350,7 +5352,7 @@
         // When lingering is complete, cell is still there but is now in the background.
         waitForIdle();
         int timeoutMs = TEST_LINGER_DELAY_MS + TEST_LINGER_DELAY_MS / 4;
-        fgCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent, timeoutMs);
+        fgCallback.expect(CallbackEntry.LOST, mCellNetworkAgent, timeoutMs);
         // Expect a network capabilities update sans FOREGROUND.
         callback.expectCapabilitiesWithout(NET_CAPABILITY_FOREGROUND, mCellNetworkAgent);
         assertFalse(isForegroundNetwork(mCellNetworkAgent));
@@ -5373,7 +5375,7 @@
         // Release the request. The network immediately goes into the background, since it was not
         // lingering.
         mCm.unregisterNetworkCallback(cellCallback);
-        fgCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        fgCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         // Expect a network capabilities update sans FOREGROUND.
         callback.expectCapabilitiesWithout(NET_CAPABILITY_FOREGROUND, mCellNetworkAgent);
         assertFalse(isForegroundNetwork(mCellNetworkAgent));
@@ -5381,8 +5383,8 @@
 
         // Disconnect wifi and check that cell is foreground again.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
-        fgCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
+        fgCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         fgCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         assertTrue(isForegroundNetwork(mCellNetworkAgent));
 
@@ -5541,7 +5543,7 @@
             // Cell disconnects. There is still the "mobile data always on" request outstanding,
             // and the test factory should see it now that it isn't hopelessly outscored.
             mCellNetworkAgent.disconnect();
-            cellNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+            cellNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
             // Wait for the network to be removed from internal structures before
             // calling synchronous getter
             waitForIdle();
@@ -5564,7 +5566,7 @@
             testFactory.assertRequestCountEquals(0);
             assertFalse(testFactory.getMyStartRequested());
             // ...  and cell data to be torn down immediately since it is no longer nascent.
-            cellNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+            cellNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
             waitForIdle();
             assertLength(1, mCm.getAllNetworks());
             testFactory.terminate();
@@ -5741,7 +5743,7 @@
 
         // Disconnect wifi and pretend the carrier restricts moving away from bad wifi.
         mWiFiNetworkAgent.disconnect();
-        wifiNetworkCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        wifiNetworkCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         // This has getAvoidBadWifi return false. This test doesn't change the value of the
         // associated setting.
         doReturn(0).when(mResources).getInteger(R.integer.config_networkAvoidBadWifi);
@@ -5861,7 +5863,7 @@
         mWiFiNetworkAgent.setNetworkInvalid(false /* isStrictMode */);
         mCm.reportNetworkConnectivity(wifiNetwork, false);
         defaultCallback.expectCapabilitiesWithout(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
-        validatedWifiCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        validatedWifiCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         expectNotification(mWiFiNetworkAgent, NotificationType.LOST_INTERNET);
 
         // Because avoid bad wifi is off, we don't switch to cellular.
@@ -5912,7 +5914,7 @@
         mWiFiNetworkAgent.setNetworkInvalid(false /* isStrictMode */);
         mCm.reportNetworkConnectivity(wifiNetwork, false);
         defaultCallback.expectCapabilitiesWithout(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
-        validatedWifiCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        validatedWifiCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         expectNotification(mWiFiNetworkAgent, NotificationType.LOST_INTERNET);
 
         // Simulate the user selecting "switch" and checking the don't ask again checkbox.
@@ -5944,7 +5946,7 @@
 
         // If cell goes down, we switch to wifi.
         mCellNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        defaultCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         validatedWifiCallback.assertNoCallback();
         // Notification is cleared yet again because the device switched to wifi.
@@ -6010,7 +6012,7 @@
         networkCallback.expectAvailableCallbacks(mWiFiNetworkAgent, false, false, false,
                 TEST_CALLBACK_TIMEOUT_MS);
         mWiFiNetworkAgent.disconnect();
-        networkCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        networkCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
 
         // Validate that UNAVAILABLE is not called
         networkCallback.assertNoCallback();
@@ -6030,7 +6032,7 @@
         mCm.requestNetwork(nr, networkCallback, timeoutMs);
 
         // pass timeout and validate that UNAVAILABLE is called
-        networkCallback.expectCallback(CallbackEntry.UNAVAILABLE, (Network) null);
+        networkCallback.expect(CallbackEntry.UNAVAILABLE);
 
         // create a network satisfying request - validate that request not triggered
         mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
@@ -6114,7 +6116,7 @@
                 // onUnavailable!
                 testFactory.triggerUnfulfillable(newRequest);
 
-                networkCallback.expectCallback(CallbackEntry.UNAVAILABLE, (Network) null);
+                networkCallback.expect(CallbackEntry.UNAVAILABLE);
 
                 // Declaring a request unfulfillable releases it automatically.
                 testFactory.expectRequestRemove();
@@ -6196,20 +6198,20 @@
             mCallbacks.add(new CallbackValue(CallbackType.ON_ERROR, error));
         }
 
-        private void expectCallback(CallbackValue callbackValue) throws InterruptedException {
+        private void expect(CallbackValue callbackValue) throws InterruptedException {
             assertEquals(callbackValue, mCallbacks.poll(TIMEOUT_MS, TimeUnit.MILLISECONDS));
         }
 
         public void expectStarted() throws Exception {
-            expectCallback(new CallbackValue(CallbackType.ON_STARTED));
+            expect(new CallbackValue(CallbackType.ON_STARTED));
         }
 
         public void expectStopped() throws Exception {
-            expectCallback(new CallbackValue(CallbackType.ON_STOPPED));
+            expect(new CallbackValue(CallbackType.ON_STOPPED));
         }
 
         public void expectError(int error) throws Exception {
-            expectCallback(new CallbackValue(CallbackType.ON_ERROR, error));
+            expect(new CallbackValue(CallbackType.ON_ERROR, error));
         }
     }
 
@@ -6269,21 +6271,21 @@
             mCallbacks.add(new CallbackValue(CallbackType.ON_ERROR, error));
         }
 
-        private void expectCallback(CallbackValue callbackValue) throws InterruptedException {
+        private void expect(CallbackValue callbackValue) throws InterruptedException {
             assertEquals(callbackValue, mCallbacks.poll(TIMEOUT_MS, TimeUnit.MILLISECONDS));
 
         }
 
         public void expectStarted() throws InterruptedException {
-            expectCallback(new CallbackValue(CallbackType.ON_STARTED));
+            expect(new CallbackValue(CallbackType.ON_STARTED));
         }
 
         public void expectStopped() throws InterruptedException {
-            expectCallback(new CallbackValue(CallbackType.ON_STOPPED));
+            expect(new CallbackValue(CallbackType.ON_STOPPED));
         }
 
         public void expectError(int error) throws InterruptedException {
-            expectCallback(new CallbackValue(CallbackType.ON_ERROR, error));
+            expect(new CallbackValue(CallbackType.ON_ERROR, error));
         }
 
         public void assertNoCallback() {
@@ -7079,7 +7081,7 @@
 
         // Disconnect wifi aware network.
         wifiAware.disconnect();
-        callback.expectCallbackThat(TIMEOUT_MS, (info) -> info instanceof CallbackEntry.Lost);
+        callback.expect(CallbackEntry.LOST, TIMEOUT_MS);
         mCm.unregisterNetworkCallback(callback);
 
         verifyNoNetwork();
@@ -7126,12 +7128,12 @@
         // ConnectivityService.
         TestNetworkAgentWrapper networkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI, lp);
         networkAgent.connect(true);
-        networkCallback.expectCallback(CallbackEntry.AVAILABLE, networkAgent);
-        networkCallback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED, networkAgent);
+        networkCallback.expect(CallbackEntry.AVAILABLE, networkAgent);
+        networkCallback.expect(CallbackEntry.NETWORK_CAPS_UPDATED, networkAgent);
         CallbackEntry.LinkPropertiesChanged cbi =
-                networkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED,
+                networkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED,
                 networkAgent);
-        networkCallback.expectCallback(CallbackEntry.BLOCKED_STATUS, networkAgent);
+        networkCallback.expect(CallbackEntry.BLOCKED_STATUS, networkAgent);
         networkCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, networkAgent);
         networkCallback.assertNoCallback();
         checkDirectlyConnectedRoutes(cbi.getLp(), asList(myIpv4Address),
@@ -7146,7 +7148,7 @@
         newLp.addLinkAddress(myIpv6Address1);
         newLp.addLinkAddress(myIpv6Address2);
         networkAgent.sendLinkProperties(newLp);
-        cbi = networkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, networkAgent);
+        cbi = networkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED, networkAgent);
         networkCallback.assertNoCallback();
         checkDirectlyConnectedRoutes(cbi.getLp(),
                 asList(myIpv4Address, myIpv6Address1, myIpv6Address2),
@@ -7154,17 +7156,38 @@
         mCm.unregisterNetworkCallback(networkCallback);
     }
 
-    private void expectNotifyNetworkStatus(List<Network> defaultNetworks, String defaultIface,
-            Integer vpnUid, String vpnIfname, List<String> underlyingIfaces) throws Exception {
+    private void expectNotifyNetworkStatus(List<Network> allNetworks, List<Network> defaultNetworks,
+            String defaultIface, Integer vpnUid, String vpnIfname, List<String> underlyingIfaces)
+            throws Exception {
         ArgumentCaptor<List<Network>> defaultNetworksCaptor = ArgumentCaptor.forClass(List.class);
         ArgumentCaptor<List<UnderlyingNetworkInfo>> vpnInfosCaptor =
                 ArgumentCaptor.forClass(List.class);
+        ArgumentCaptor<List<NetworkStateSnapshot>> snapshotsCaptor =
+                ArgumentCaptor.forClass(List.class);
 
         verify(mStatsManager, atLeastOnce()).notifyNetworkStatus(defaultNetworksCaptor.capture(),
-                any(List.class), eq(defaultIface), vpnInfosCaptor.capture());
+                snapshotsCaptor.capture(), eq(defaultIface), vpnInfosCaptor.capture());
 
         assertSameElements(defaultNetworks, defaultNetworksCaptor.getValue());
 
+        List<Network> snapshotNetworks = new ArrayList<Network>();
+        for (NetworkStateSnapshot ns : snapshotsCaptor.getValue()) {
+            snapshotNetworks.add(ns.getNetwork());
+        }
+        assertSameElements(allNetworks, snapshotNetworks);
+
+        if (defaultIface != null) {
+            assertNotNull(
+                    "Did not find interface " + defaultIface + " in call to notifyNetworkStatus",
+                    CollectionUtils.findFirst(snapshotsCaptor.getValue(), (ns) -> {
+                        final LinkProperties lp = ns.getLinkProperties();
+                        if (lp != null && TextUtils.equals(defaultIface, lp.getInterfaceName())) {
+                            return true;
+                        }
+                        return false;
+                    }));
+        }
+
         List<UnderlyingNetworkInfo> infos = vpnInfosCaptor.getValue();
         if (vpnUid != null) {
             assertEquals("Should have exactly one VPN:", 1, infos.size());
@@ -7179,28 +7202,38 @@
     }
 
     private void expectNotifyNetworkStatus(
-            List<Network> defaultNetworks, String defaultIface) throws Exception {
-        expectNotifyNetworkStatus(defaultNetworks, defaultIface, null, null, List.of());
+            List<Network> allNetworks, List<Network> defaultNetworks, String defaultIface)
+            throws Exception {
+        expectNotifyNetworkStatus(allNetworks, defaultNetworks, defaultIface, null, null,
+                List.of());
+    }
+
+    private List<Network> onlyCell() {
+        return List.of(mCellNetworkAgent.getNetwork());
+    }
+
+    private List<Network> onlyWifi() {
+        return List.of(mWiFiNetworkAgent.getNetwork());
+    }
+
+    private List<Network> cellAndWifi() {
+        return List.of(mCellNetworkAgent.getNetwork(), mWiFiNetworkAgent.getNetwork());
     }
 
     @Test
     public void testStatsIfacesChanged() throws Exception {
-        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
-        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
-
-        final List<Network> onlyCell = List.of(mCellNetworkAgent.getNetwork());
-        final List<Network> onlyWifi = List.of(mWiFiNetworkAgent.getNetwork());
-
         LinkProperties cellLp = new LinkProperties();
         cellLp.setInterfaceName(MOBILE_IFNAME);
         LinkProperties wifiLp = new LinkProperties();
         wifiLp.setInterfaceName(WIFI_IFNAME);
 
-        // Simple connection should have updated ifaces
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR, cellLp);
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
+
+        // Simple connection with initial LP should have updated ifaces.
         mCellNetworkAgent.connect(false);
-        mCellNetworkAgent.sendLinkProperties(cellLp);
         waitForIdle();
-        expectNotifyNetworkStatus(onlyCell, MOBILE_IFNAME);
+        expectNotifyNetworkStatus(onlyCell(), onlyCell(), MOBILE_IFNAME);
         reset(mStatsManager);
 
         // Default network switch should update ifaces.
@@ -7208,37 +7241,56 @@
         mWiFiNetworkAgent.sendLinkProperties(wifiLp);
         waitForIdle();
         assertEquals(wifiLp, mService.getActiveLinkProperties());
-        expectNotifyNetworkStatus(onlyWifi, WIFI_IFNAME);
+        expectNotifyNetworkStatus(cellAndWifi(), onlyWifi(), WIFI_IFNAME);
+        reset(mStatsManager);
+
+        // Disconnecting a network updates ifaces again. The soon-to-be disconnected interface is
+        // still in the list to ensure that stats are counted on that interface.
+        // TODO: this means that if anything else uses that interface for any other reason before
+        // notifyNetworkStatus is called again, traffic on that interface will be accounted to the
+        // disconnected network. This is likely a bug in ConnectivityService; it should probably
+        // call notifyNetworkStatus again without the disconnected network.
+        mCellNetworkAgent.disconnect();
+        waitForIdle();
+        expectNotifyNetworkStatus(cellAndWifi(), onlyWifi(), WIFI_IFNAME);
+        verifyNoMoreInteractions(mStatsManager);
+        reset(mStatsManager);
+
+        // Connecting a network updates ifaces even if the network doesn't become default.
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR, cellLp);
+        mCellNetworkAgent.connect(false);
+        waitForIdle();
+        expectNotifyNetworkStatus(cellAndWifi(), onlyWifi(), WIFI_IFNAME);
         reset(mStatsManager);
 
         // Disconnect should update ifaces.
         mWiFiNetworkAgent.disconnect();
         waitForIdle();
-        expectNotifyNetworkStatus(onlyCell, MOBILE_IFNAME);
+        expectNotifyNetworkStatus(onlyCell(), onlyCell(), MOBILE_IFNAME);
         reset(mStatsManager);
 
         // Metered change should update ifaces
         mCellNetworkAgent.addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_METERED);
         waitForIdle();
-        expectNotifyNetworkStatus(onlyCell, MOBILE_IFNAME);
+        expectNotifyNetworkStatus(onlyCell(), onlyCell(), MOBILE_IFNAME);
         reset(mStatsManager);
 
         mCellNetworkAgent.removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_METERED);
         waitForIdle();
-        expectNotifyNetworkStatus(onlyCell, MOBILE_IFNAME);
+        expectNotifyNetworkStatus(onlyCell(), onlyCell(), MOBILE_IFNAME);
         reset(mStatsManager);
 
         // Temp metered change shouldn't update ifaces
         mCellNetworkAgent.addCapability(NET_CAPABILITY_TEMPORARILY_NOT_METERED);
         waitForIdle();
-        verify(mStatsManager, never()).notifyNetworkStatus(eq(onlyCell),
+        verify(mStatsManager, never()).notifyNetworkStatus(eq(onlyCell()),
                 any(List.class), eq(MOBILE_IFNAME), any(List.class));
         reset(mStatsManager);
 
         // Roaming change should update ifaces
         mCellNetworkAgent.addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_ROAMING);
         waitForIdle();
-        expectNotifyNetworkStatus(onlyCell, MOBILE_IFNAME);
+        expectNotifyNetworkStatus(onlyCell(), onlyCell(), MOBILE_IFNAME);
         reset(mStatsManager);
 
         // Test VPNs.
@@ -7252,8 +7304,8 @@
                 List.of(mCellNetworkAgent.getNetwork(), mMockVpn.getNetwork());
 
         // A VPN with default (null) underlying networks sets the underlying network's interfaces...
-        expectNotifyNetworkStatus(cellAndVpn, MOBILE_IFNAME, Process.myUid(), VPN_IFNAME,
-                List.of(MOBILE_IFNAME));
+        expectNotifyNetworkStatus(cellAndVpn, cellAndVpn, MOBILE_IFNAME, Process.myUid(),
+                VPN_IFNAME, List.of(MOBILE_IFNAME));
 
         // ...and updates them as the default network switches.
         mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
@@ -7262,15 +7314,16 @@
         final Network[] onlyNull = new Network[]{null};
         final List<Network> wifiAndVpn =
                 List.of(mWiFiNetworkAgent.getNetwork(), mMockVpn.getNetwork());
-        final List<Network> cellAndWifi =
-                List.of(mCellNetworkAgent.getNetwork(), mWiFiNetworkAgent.getNetwork());
+        final List<Network> cellWifiAndVpn =
+                List.of(mCellNetworkAgent.getNetwork(), mWiFiNetworkAgent.getNetwork(),
+                        mMockVpn.getNetwork());
         final Network[] cellNullAndWifi =
                 new Network[]{mCellNetworkAgent.getNetwork(), null, mWiFiNetworkAgent.getNetwork()};
 
         waitForIdle();
         assertEquals(wifiLp, mService.getActiveLinkProperties());
-        expectNotifyNetworkStatus(wifiAndVpn, WIFI_IFNAME, Process.myUid(), VPN_IFNAME,
-                List.of(WIFI_IFNAME));
+        expectNotifyNetworkStatus(cellWifiAndVpn, wifiAndVpn, WIFI_IFNAME, Process.myUid(),
+                VPN_IFNAME, List.of(WIFI_IFNAME));
         reset(mStatsManager);
 
         // A VPN that sets its underlying networks passes the underlying interfaces, and influences
@@ -7279,23 +7332,23 @@
         // MOBILE_IFNAME even though the default network is wifi.
         // TODO: fix this to pass in the actual default network interface. Whether or not the VPN
         // applies to the system server UID should not have any bearing on network stats.
-        mMockVpn.setUnderlyingNetworks(onlyCell.toArray(new Network[0]));
+        mMockVpn.setUnderlyingNetworks(onlyCell().toArray(new Network[0]));
         waitForIdle();
-        expectNotifyNetworkStatus(wifiAndVpn, MOBILE_IFNAME, Process.myUid(), VPN_IFNAME,
-                List.of(MOBILE_IFNAME));
+        expectNotifyNetworkStatus(cellWifiAndVpn, wifiAndVpn, MOBILE_IFNAME, Process.myUid(),
+                VPN_IFNAME, List.of(MOBILE_IFNAME));
         reset(mStatsManager);
 
-        mMockVpn.setUnderlyingNetworks(cellAndWifi.toArray(new Network[0]));
+        mMockVpn.setUnderlyingNetworks(cellAndWifi().toArray(new Network[0]));
         waitForIdle();
-        expectNotifyNetworkStatus(wifiAndVpn, MOBILE_IFNAME, Process.myUid(), VPN_IFNAME,
-                List.of(MOBILE_IFNAME, WIFI_IFNAME));
+        expectNotifyNetworkStatus(cellWifiAndVpn, wifiAndVpn, MOBILE_IFNAME, Process.myUid(),
+                VPN_IFNAME,  List.of(MOBILE_IFNAME, WIFI_IFNAME));
         reset(mStatsManager);
 
         // Null underlying networks are ignored.
         mMockVpn.setUnderlyingNetworks(cellNullAndWifi);
         waitForIdle();
-        expectNotifyNetworkStatus(wifiAndVpn, MOBILE_IFNAME, Process.myUid(), VPN_IFNAME,
-                List.of(MOBILE_IFNAME, WIFI_IFNAME));
+        expectNotifyNetworkStatus(cellWifiAndVpn, wifiAndVpn, MOBILE_IFNAME, Process.myUid(),
+                VPN_IFNAME,  List.of(MOBILE_IFNAME, WIFI_IFNAME));
         reset(mStatsManager);
 
         // If an underlying network disconnects, that interface should no longer be underlying.
@@ -7308,8 +7361,8 @@
         mCellNetworkAgent.disconnect();
         waitForIdle();
         assertNull(mService.getLinkProperties(mCellNetworkAgent.getNetwork()));
-        expectNotifyNetworkStatus(wifiAndVpn, MOBILE_IFNAME, Process.myUid(), VPN_IFNAME,
-                List.of(MOBILE_IFNAME, WIFI_IFNAME));
+        expectNotifyNetworkStatus(cellWifiAndVpn, wifiAndVpn, MOBILE_IFNAME, Process.myUid(),
+                VPN_IFNAME, List.of(MOBILE_IFNAME, WIFI_IFNAME));
 
         // Confirm that we never tell NetworkStatsService that cell is no longer the underlying
         // network for the VPN...
@@ -7344,25 +7397,25 @@
         // Also, for the same reason as above, the active interface passed in is null.
         mMockVpn.setUnderlyingNetworks(new Network[0]);
         waitForIdle();
-        expectNotifyNetworkStatus(wifiAndVpn, null);
+        expectNotifyNetworkStatus(wifiAndVpn, wifiAndVpn, null);
         reset(mStatsManager);
 
         // Specifying only a null underlying network is the same as no networks.
         mMockVpn.setUnderlyingNetworks(onlyNull);
         waitForIdle();
-        expectNotifyNetworkStatus(wifiAndVpn, null);
+        expectNotifyNetworkStatus(wifiAndVpn, wifiAndVpn, null);
         reset(mStatsManager);
 
         // Specifying networks that are all disconnected is the same as specifying no networks.
-        mMockVpn.setUnderlyingNetworks(onlyCell.toArray(new Network[0]));
+        mMockVpn.setUnderlyingNetworks(onlyCell().toArray(new Network[0]));
         waitForIdle();
-        expectNotifyNetworkStatus(wifiAndVpn, null);
+        expectNotifyNetworkStatus(wifiAndVpn, wifiAndVpn, null);
         reset(mStatsManager);
 
         // Passing in null again means follow the default network again.
         mMockVpn.setUnderlyingNetworks(null);
         waitForIdle();
-        expectNotifyNetworkStatus(wifiAndVpn, WIFI_IFNAME, Process.myUid(), VPN_IFNAME,
+        expectNotifyNetworkStatus(wifiAndVpn, wifiAndVpn, WIFI_IFNAME, Process.myUid(), VPN_IFNAME,
                 List.of(WIFI_IFNAME));
         reset(mStatsManager);
     }
@@ -7381,7 +7434,7 @@
                 NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK, PERMISSION_GRANTED);
         TestNetworkCallback callback = new TestNetworkCallback();
         mCm.registerDefaultNetworkCallback(callback);
-        callback.expectCallback(CallbackEntry.AVAILABLE, mCellNetworkAgent);
+        callback.expect(CallbackEntry.AVAILABLE, mCellNetworkAgent);
         callback.expectCapabilitiesThat(
                 mCellNetworkAgent, nc -> Arrays.equals(adminUids, nc.getAdministratorUids()));
         mCm.unregisterNetworkCallback(callback);
@@ -7392,7 +7445,7 @@
         mServiceContext.setPermission(NETWORK_STACK, PERMISSION_DENIED);
         callback = new TestNetworkCallback();
         mCm.registerDefaultNetworkCallback(callback);
-        callback.expectCallback(CallbackEntry.AVAILABLE, mCellNetworkAgent);
+        callback.expect(CallbackEntry.AVAILABLE, mCellNetworkAgent);
         callback.expectCapabilitiesThat(
                 mCellNetworkAgent, nc -> nc.getAdministratorUids().length == 0);
     }
@@ -7417,7 +7470,7 @@
         mWiFiNetworkAgent.connect(true /* validated */);
 
         final List<Network> none = List.of();
-        expectNotifyNetworkStatus(none, null);  // Wifi is not the default network
+        expectNotifyNetworkStatus(onlyWifi(), none, null);  // Wifi is not the default network
 
         // Create a virtual network based on the wifi network.
         final int ownerUid = 10042;
@@ -7442,7 +7495,9 @@
         assertFalse(nc.hasTransport(TRANSPORT_WIFI));
         assertFalse(nc.hasCapability(NET_CAPABILITY_NOT_METERED));
         final List<Network> onlyVcn = List.of(vcn.getNetwork());
-        expectNotifyNetworkStatus(onlyVcn, vcnIface, ownerUid, vcnIface, List.of(WIFI_IFNAME));
+        final List<Network> vcnAndWifi = List.of(vcn.getNetwork(), mWiFiNetworkAgent.getNetwork());
+        expectNotifyNetworkStatus(vcnAndWifi, onlyVcn, vcnIface, ownerUid, vcnIface,
+                List.of(WIFI_IFNAME));
 
         // Add NOT_METERED to the underlying network, check that it is not propagated.
         mWiFiNetworkAgent.addCapability(NET_CAPABILITY_NOT_METERED);
@@ -7465,7 +7520,10 @@
         nc = mCm.getNetworkCapabilities(vcn.getNetwork());
         assertFalse(nc.hasTransport(TRANSPORT_WIFI));
         assertFalse(nc.hasCapability(NET_CAPABILITY_NOT_ROAMING));
-        expectNotifyNetworkStatus(onlyVcn, vcnIface, ownerUid, vcnIface, List.of(MOBILE_IFNAME));
+        final List<Network> vcnWifiAndCell = List.of(vcn.getNetwork(),
+                mWiFiNetworkAgent.getNetwork(), mCellNetworkAgent.getNetwork());
+        expectNotifyNetworkStatus(vcnWifiAndCell, onlyVcn, vcnIface, ownerUid, vcnIface,
+                List.of(MOBILE_IFNAME));
     }
 
     @Test
@@ -7646,12 +7704,12 @@
         assertTrue(new ArraySet<>(resolvrParams.tlsServers).containsAll(
                 asList("2001:db8::1", "192.0.2.1")));
         reset(mMockDnsResolver);
-        cellNetworkCallback.expectCallback(CallbackEntry.AVAILABLE, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED,
+        cellNetworkCallback.expect(CallbackEntry.AVAILABLE, mCellNetworkAgent);
+        cellNetworkCallback.expect(CallbackEntry.NETWORK_CAPS_UPDATED,
                 mCellNetworkAgent);
-        CallbackEntry.LinkPropertiesChanged cbi = cellNetworkCallback.expectCallback(
+        CallbackEntry.LinkPropertiesChanged cbi = cellNetworkCallback.expect(
                 CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackEntry.BLOCKED_STATUS, mCellNetworkAgent);
+        cellNetworkCallback.expect(CallbackEntry.BLOCKED_STATUS, mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
         assertFalse(cbi.getLp().isPrivateDnsActive());
         assertNull(cbi.getLp().getPrivateDnsServerName());
@@ -7682,7 +7740,7 @@
         setPrivateDnsSettings(PRIVATE_DNS_MODE_PROVIDER_HOSTNAME, "strict.example.com");
         // Can't test dns configuration for strict mode without properly mocking
         // out the DNS lookups, but can test that LinkProperties is updated.
-        cbi = cellNetworkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED,
+        cbi = cellNetworkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
         assertTrue(cbi.getLp().isPrivateDnsActive());
@@ -7715,12 +7773,12 @@
         mCellNetworkAgent.sendLinkProperties(lp);
         mCellNetworkAgent.connect(false);
         waitForIdle();
-        cellNetworkCallback.expectCallback(CallbackEntry.AVAILABLE, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED,
+        cellNetworkCallback.expect(CallbackEntry.AVAILABLE, mCellNetworkAgent);
+        cellNetworkCallback.expect(CallbackEntry.NETWORK_CAPS_UPDATED,
                 mCellNetworkAgent);
-        CallbackEntry.LinkPropertiesChanged cbi = cellNetworkCallback.expectCallback(
+        CallbackEntry.LinkPropertiesChanged cbi = cellNetworkCallback.expect(
                 CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackEntry.BLOCKED_STATUS, mCellNetworkAgent);
+        cellNetworkCallback.expect(CallbackEntry.BLOCKED_STATUS, mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
         assertFalse(cbi.getLp().isPrivateDnsActive());
         assertNull(cbi.getLp().getPrivateDnsServerName());
@@ -7738,7 +7796,7 @@
         LinkProperties lp2 = new LinkProperties(lp);
         lp2.addDnsServer(InetAddress.getByName("145.100.185.16"));
         mCellNetworkAgent.sendLinkProperties(lp2);
-        cbi = cellNetworkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED,
+        cbi = cellNetworkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
         assertFalse(cbi.getLp().isPrivateDnsActive());
@@ -7765,7 +7823,7 @@
         mService.mResolverUnsolEventCallback.onPrivateDnsValidationEvent(
                 makePrivateDnsValidationEvent(mCellNetworkAgent.getNetwork().netId,
                         "145.100.185.16", "", VALIDATION_RESULT_SUCCESS));
-        cbi = cellNetworkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED,
+        cbi = cellNetworkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
         assertTrue(cbi.getLp().isPrivateDnsActive());
@@ -7777,7 +7835,7 @@
         LinkProperties lp3 = new LinkProperties(lp2);
         lp3.setMtu(1300);
         mCellNetworkAgent.sendLinkProperties(lp3);
-        cbi = cellNetworkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED,
+        cbi = cellNetworkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
         assertTrue(cbi.getLp().isPrivateDnsActive());
@@ -7790,7 +7848,7 @@
         LinkProperties lp4 = new LinkProperties(lp3);
         lp4.removeDnsServer(InetAddress.getByName("145.100.185.16"));
         mCellNetworkAgent.sendLinkProperties(lp4);
-        cbi = cellNetworkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED,
+        cbi = cellNetworkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
         assertFalse(cbi.getLp().isPrivateDnsActive());
@@ -7982,7 +8040,7 @@
             mMockVpn.setUnderlyingNetworks(new Network[]{wifiNetwork});
             // onCapabilitiesChanged() should be called because
             // NetworkCapabilities#mUnderlyingNetworks is updated.
-            CallbackEntry ce = callback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED,
+            CallbackEntry ce = callback.expect(CallbackEntry.NETWORK_CAPS_UPDATED,
                     mMockVpn);
             final NetworkCapabilities vpnNc1 = ((CallbackEntry.CapabilitiesChanged) ce).getCaps();
             // Since the wifi network hasn't brought up,
@@ -8019,7 +8077,7 @@
             // 2. When a network connects, updateNetworkInfo propagates underlying network
             //    capabilities before rematching networks.
             // Given that this scenario can't really happen, this is probably fine for now.
-            ce = callback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED, mMockVpn);
+            ce = callback.expect(CallbackEntry.NETWORK_CAPS_UPDATED, mMockVpn);
             final NetworkCapabilities vpnNc2 = ((CallbackEntry.CapabilitiesChanged) ce).getCaps();
             // The wifi network is brought up, NetworkCapabilities#mUnderlyingNetworks is updated to
             // it.
@@ -8033,7 +8091,7 @@
 
             // Disconnect the network, and expect to see the VPN capabilities change accordingly.
             mWiFiNetworkAgent.disconnect();
-            callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+            callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
             callback.expectCapabilitiesThat(mMockVpn, (nc) ->
                     nc.getTransportTypes().length == 1 && nc.hasTransport(TRANSPORT_VPN));
 
@@ -8079,7 +8137,7 @@
         callback.expectCapabilitiesThat(mMockVpn,
                 nc -> !nc.hasCapability(NET_CAPABILITY_NOT_SUSPENDED)
                         && nc.hasTransport(TRANSPORT_CELLULAR));
-        callback.expectCallback(CallbackEntry.SUSPENDED, mMockVpn);
+        callback.expect(CallbackEntry.SUSPENDED, mMockVpn);
         callback.assertNoCallback();
 
         assertFalse(mCm.getNetworkCapabilities(mMockVpn.getNetwork())
@@ -8097,7 +8155,7 @@
         callback.expectCapabilitiesThat(mMockVpn,
                 nc -> nc.hasCapability(NET_CAPABILITY_NOT_SUSPENDED)
                         && nc.hasTransport(TRANSPORT_WIFI));
-        callback.expectCallback(CallbackEntry.RESUMED, mMockVpn);
+        callback.expect(CallbackEntry.RESUMED, mMockVpn);
         callback.assertNoCallback();
 
         assertTrue(mCm.getNetworkCapabilities(mMockVpn.getNetwork())
@@ -8134,7 +8192,7 @@
         callback.expectCapabilitiesThat(mMockVpn,
                 nc -> !nc.hasCapability(NET_CAPABILITY_NOT_SUSPENDED)
                         && nc.hasTransport(TRANSPORT_CELLULAR));
-        callback.expectCallback(CallbackEntry.SUSPENDED, mMockVpn);
+        callback.expect(CallbackEntry.SUSPENDED, mMockVpn);
         callback.assertNoCallback();
 
         assertFalse(mCm.getNetworkCapabilities(mMockVpn.getNetwork())
@@ -8150,7 +8208,7 @@
         callback.expectCapabilitiesThat(mMockVpn,
                 nc -> nc.hasCapability(NET_CAPABILITY_NOT_SUSPENDED)
                         && nc.hasTransport(TRANSPORT_CELLULAR));
-        callback.expectCallback(CallbackEntry.RESUMED, mMockVpn);
+        callback.expect(CallbackEntry.RESUMED, mMockVpn);
         callback.assertNoCallback();
 
         assertTrue(mCm.getNetworkCapabilities(mMockVpn.getNetwork())
@@ -8230,10 +8288,10 @@
         ranges.clear();
         mMockVpn.setUids(ranges);
 
-        genericNetworkCallback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        genericNetworkCallback.expect(CallbackEntry.LOST, mMockVpn);
         genericNotVpnNetworkCallback.assertNoCallback();
         wifiNetworkCallback.assertNoCallback();
-        vpnNetworkCallback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        vpnNetworkCallback.expect(CallbackEntry.LOST, mMockVpn);
 
         // TODO : The default network callback should actually get a LOST call here (also see the
         // comment below for AVAILABLE). This is because ConnectivityService does not look at UID
@@ -8241,7 +8299,7 @@
         // can't currently update their UIDs without disconnecting, so this does not matter too
         // much, but that is the reason the test here has to check for an update to the
         // capabilities instead of the expected LOST then AVAILABLE.
-        defaultCallback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED, mMockVpn);
+        defaultCallback.expect(CallbackEntry.NETWORK_CAPS_UPDATED, mMockVpn);
         systemDefaultCallback.assertNoCallback();
 
         ranges.add(new UidRange(uid, uid));
@@ -8253,25 +8311,25 @@
         vpnNetworkCallback.expectAvailableCallbacksValidated(mMockVpn);
         // TODO : Here like above, AVAILABLE would be correct, but because this can't actually
         // happen outside of the test, ConnectivityService does not rematch callbacks.
-        defaultCallback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED, mMockVpn);
+        defaultCallback.expect(CallbackEntry.NETWORK_CAPS_UPDATED, mMockVpn);
         systemDefaultCallback.assertNoCallback();
 
         mWiFiNetworkAgent.disconnect();
 
-        genericNetworkCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
-        genericNotVpnNetworkCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
-        wifiNetworkCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        genericNetworkCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
+        genericNotVpnNetworkCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
+        wifiNetworkCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         vpnNetworkCallback.assertNoCallback();
         defaultCallback.assertNoCallback();
-        systemDefaultCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        systemDefaultCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
 
         mMockVpn.disconnect();
 
-        genericNetworkCallback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        genericNetworkCallback.expect(CallbackEntry.LOST, mMockVpn);
         genericNotVpnNetworkCallback.assertNoCallback();
         wifiNetworkCallback.assertNoCallback();
-        vpnNetworkCallback.expectCallback(CallbackEntry.LOST, mMockVpn);
-        defaultCallback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        vpnNetworkCallback.expect(CallbackEntry.LOST, mMockVpn);
+        defaultCallback.expect(CallbackEntry.LOST, mMockVpn);
         systemDefaultCallback.assertNoCallback();
         assertEquals(null, mCm.getActiveNetwork());
 
@@ -8329,7 +8387,7 @@
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         mMockVpn.disconnect();
-        defaultCallback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        defaultCallback.expect(CallbackEntry.LOST, mMockVpn);
         defaultCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
 
         mCm.unregisterNetworkCallback(defaultCallback);
@@ -8377,7 +8435,7 @@
         callback.assertNoCallback();
 
         mMockVpn.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        callback.expect(CallbackEntry.LOST, mMockVpn);
         callback.expectAvailableCallbacksValidated(mEthernetNetworkAgent);
     }
 
@@ -8511,7 +8569,7 @@
                 && caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
                 && !caps.hasCapability(NET_CAPABILITY_NOT_METERED)
                 && !caps.hasCapability(NET_CAPABILITY_NOT_SUSPENDED));
-        vpnNetworkCallback.expectCallback(CallbackEntry.SUSPENDED, mMockVpn);
+        vpnNetworkCallback.expect(CallbackEntry.SUSPENDED, mMockVpn);
 
         // Add NOT_SUSPENDED again and observe VPN is no longer suspended.
         mCellNetworkAgent.addCapability(NET_CAPABILITY_NOT_SUSPENDED);
@@ -8520,7 +8578,7 @@
                 && caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
                 && !caps.hasCapability(NET_CAPABILITY_NOT_METERED)
                 && caps.hasCapability(NET_CAPABILITY_NOT_SUSPENDED));
-        vpnNetworkCallback.expectCallback(CallbackEntry.RESUMED, mMockVpn);
+        vpnNetworkCallback.expect(CallbackEntry.RESUMED, mMockVpn);
 
         // Use Wifi but not cell. Note the VPN is now unmetered and not suspended.
         mMockVpn.setUnderlyingNetworks(
@@ -8556,7 +8614,7 @@
                 && caps.hasTransport(TRANSPORT_CELLULAR)
                 && !caps.hasCapability(NET_CAPABILITY_NOT_METERED)
                 && !caps.hasCapability(NET_CAPABILITY_NOT_SUSPENDED));
-        vpnNetworkCallback.expectCallback(CallbackEntry.SUSPENDED, mMockVpn);
+        vpnNetworkCallback.expect(CallbackEntry.SUSPENDED, mMockVpn);
         assertDefaultNetworkCapabilities(userId, mCellNetworkAgent, mWiFiNetworkAgent);
 
         // Use both again.
@@ -8568,7 +8626,7 @@
                 && caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
                 && !caps.hasCapability(NET_CAPABILITY_NOT_METERED)
                 && caps.hasCapability(NET_CAPABILITY_NOT_SUSPENDED));
-        vpnNetworkCallback.expectCallback(CallbackEntry.RESUMED, mMockVpn);
+        vpnNetworkCallback.expect(CallbackEntry.RESUMED, mMockVpn);
         assertDefaultNetworkCapabilities(userId, mCellNetworkAgent, mWiFiNetworkAgent);
 
         // Disconnect cell. Receive update without even removing the dead network from the
@@ -8709,7 +8767,7 @@
 
         // Change the VPN's capabilities somehow (specifically, disconnect wifi).
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         callback.expectCapabilitiesThat(mMockVpn, (caps)
                 -> caps.getUids().size() == 2
                 && caps.getUids().contains(singleUidRange)
@@ -9133,7 +9191,7 @@
 
         // Switch to METERED network. Restrict the use of the network.
         mWiFiNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        defaultCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidatedAndBlocked(mCellNetworkAgent);
 
         // Network becomes NOT_METERED.
@@ -9147,7 +9205,7 @@
         defaultCallback.assertNoCallback();
 
         mCellNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        defaultCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         defaultCallback.assertNoCallback();
 
         mCm.unregisterNetworkCallback(defaultCallback);
@@ -9211,10 +9269,9 @@
 
         // Expect exactly one blocked callback for each agent.
         for (int i = 0; i < agents.length; i++) {
-            CallbackEntry e = callback.expectCallbackThat(TIMEOUT_MS, (c) ->
-                    c instanceof CallbackEntry.BlockedStatus
-                            && ((CallbackEntry.BlockedStatus) c).getBlocked() == blocked);
-            Network network = e.getNetwork();
+            final CallbackEntry e = callback.expect(CallbackEntry.BLOCKED_STATUS, TIMEOUT_MS,
+                    c -> c.getBlocked() == blocked);
+            final Network network = e.getNetwork();
             assertTrue("Received unexpected blocked callback for network " + network,
                     expectedNetworks.remove(network));
         }
@@ -9419,7 +9476,7 @@
         assertNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
 
         mMockVpn.disconnect();
-        defaultCallback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        defaultCallback.expect(CallbackEntry.LOST, mMockVpn);
         defaultCallback.expectAvailableCallbacksUnvalidatedAndBlocked(mWiFiNetworkAgent);
         vpnUidCallback.assertNoCallback();
         vpnUidDefaultCallback.assertNoCallback();
@@ -9571,9 +9628,9 @@
         cellLp.addLinkAddress(new LinkAddress("192.0.2.2/25"));
         cellLp.addRoute(new RouteInfo(new IpPrefix("0.0.0.0/0"), null, "rmnet0"));
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        callback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
-        defaultCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
-        systemDefaultCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED,
+        callback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        defaultCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        systemDefaultCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         waitForIdle();
         assertNull(mMockVpn.getAgent());
@@ -9582,9 +9639,9 @@
         // Expect lockdown VPN to come up.
         ExpectedBroadcast b1 = expectConnectivityAction(TYPE_MOBILE, DetailedState.DISCONNECTED);
         mCellNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
-        defaultCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
-        systemDefaultCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mCellNetworkAgent);
+        defaultCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
+        systemDefaultCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         b1.expectBroadcast();
 
         // When lockdown VPN is active, the NetworkInfo state in CONNECTIVITY_ACTION is overwritten
@@ -9662,8 +9719,8 @@
         // fact that a VPN is connected should only result in the VPN itself being unblocked, not
         // any other network. Bug in isUidBlockedByVpn?
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackEntry.LOST, mMockVpn);
-        defaultCallback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        callback.expect(CallbackEntry.LOST, mMockVpn);
+        defaultCallback.expect(CallbackEntry.LOST, mMockVpn);
         defaultCallback.expectAvailableCallbacksUnvalidatedAndBlocked(mWiFiNetworkAgent);
         systemDefaultCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
 
@@ -9696,7 +9753,7 @@
 
         // Disconnect cell. Nothing much happens since it's not the default network.
         mCellNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         defaultCallback.assertNoCallback();
         systemDefaultCallback.assertNoCallback();
 
@@ -9709,12 +9766,12 @@
         b1 = expectConnectivityAction(TYPE_WIFI, DetailedState.DISCONNECTED);
         b2 = expectConnectivityAction(TYPE_VPN, DetailedState.DISCONNECTED);
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
-        systemDefaultCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
+        systemDefaultCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         b1.expectBroadcast();
         callback.expectCapabilitiesThat(mMockVpn, nc -> !nc.hasTransport(TRANSPORT_WIFI));
         mMockVpn.expectStopVpnRunnerPrivileged();
-        callback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        callback.expect(CallbackEntry.LOST, mMockVpn);
         b2.expectBroadcast();
 
         VMSHandlerThread.quitSafely();
@@ -9886,7 +9943,7 @@
             reset(mMockNetd);
 
             mCellNetworkAgent.removeCapability(testCap);
-            callbackWithCap.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+            callbackWithCap.expect(CallbackEntry.LOST, mCellNetworkAgent);
             callbackWithoutCap.assertNoCallback();
             verify(mMockNetd).networkClearDefault();
 
@@ -10076,7 +10133,7 @@
         // the NAT64 prefix was removed because one was never discovered.
         cellLp.addLinkAddress(myIpv4);
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        networkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        networkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         assertRoutesAdded(cellNetId, ipv4Subnet);
         verify(mMockDnsResolver, times(1)).stopPrefix64Discovery(cellNetId);
         verify(mMockDnsResolver, atLeastOnce()).setResolverConfiguration(any());
@@ -10099,7 +10156,7 @@
         // Remove IPv4 address. Expect prefix discovery to be started again.
         cellLp.removeLinkAddress(myIpv4);
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        networkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        networkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         verify(mMockDnsResolver, times(1)).startPrefix64Discovery(cellNetId);
         assertRoutesRemoved(cellNetId, ipv4Subnet);
 
@@ -10108,7 +10165,7 @@
         assertNull(mCm.getLinkProperties(mCellNetworkAgent.getNetwork()).getNat64Prefix());
         mService.mResolverUnsolEventCallback.onNat64PrefixEvent(
                 makeNat64PrefixEvent(cellNetId, PREFIX_OPERATION_ADDED, kNat64PrefixString, 96));
-        LinkProperties lpBeforeClat = networkCallback.expectCallback(
+        LinkProperties lpBeforeClat = networkCallback.expect(
                 CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent).getLp();
         assertEquals(0, lpBeforeClat.getStackedLinks().size());
         assertEquals(kNat64Prefix, lpBeforeClat.getNat64Prefix());
@@ -10116,7 +10173,7 @@
 
         // Clat iface comes up. Expect stacked link to be added.
         clat.interfaceLinkStateChanged(CLAT_MOBILE_IFNAME, true);
-        networkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        networkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         List<LinkProperties> stackedLps = mCm.getLinkProperties(mCellNetworkAgent.getNetwork())
                 .getStackedLinks();
         assertEquals(makeClatLinkProperties(myIpv4), stackedLps.get(0));
@@ -10125,7 +10182,7 @@
         // Change trivial linkproperties and see if stacked link is preserved.
         cellLp.addDnsServer(InetAddress.getByName("8.8.8.8"));
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        networkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        networkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
 
         List<LinkProperties> stackedLpsAfterChange =
                 mCm.getLinkProperties(mCellNetworkAgent.getNetwork()).getStackedLinks();
@@ -10174,13 +10231,13 @@
         cellLp.addLinkAddress(myIpv4);
         cellLp.addRoute(ipv4Subnet);
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        networkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        networkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         assertRoutesAdded(cellNetId, ipv4Subnet);
         verifyClatdStop(null /* inOrder */, MOBILE_IFNAME);
         verify(mMockDnsResolver, times(1)).stopPrefix64Discovery(cellNetId);
 
         // As soon as stop is called, the linkproperties lose the stacked interface.
-        networkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        networkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         LinkProperties actualLpAfterIpv4 = mCm.getLinkProperties(mCellNetworkAgent.getNetwork());
         LinkProperties expected = new LinkProperties(cellLp);
         expected.setNat64Prefix(kOtherNat64Prefix);
@@ -10212,12 +10269,12 @@
         cellLp.removeRoute(new RouteInfo(myIpv4, null, MOBILE_IFNAME));
         cellLp.removeDnsServer(InetAddress.getByName("8.8.8.8"));
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        networkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        networkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         assertRoutesRemoved(cellNetId, ipv4Subnet);  // Directly-connected routes auto-added.
         verify(mMockDnsResolver, times(1)).startPrefix64Discovery(cellNetId);
         mService.mResolverUnsolEventCallback.onNat64PrefixEvent(makeNat64PrefixEvent(
                 cellNetId, PREFIX_OPERATION_ADDED, kNat64PrefixString, 96));
-        networkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        networkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         verifyClatdStart(null /* inOrder */, MOBILE_IFNAME, cellNetId, kNat64Prefix.toString());
 
         // Clat iface comes up. Expect stacked link to be added.
@@ -10242,7 +10299,7 @@
         verify(mMockNetd, times(1)).interfaceGetCfg(CLAT_MOBILE_IFNAME);
         // Clean up.
         mCellNetworkAgent.disconnect();
-        networkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        networkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         networkCallback.assertNoCallback();
         verify(mMockNetd, times(1)).idletimerRemoveInterface(eq(MOBILE_IFNAME), anyInt(),
                 eq(Integer.toString(TRANSPORT_CELLULAR)));
@@ -10281,7 +10338,7 @@
 
         // Disconnect the network. clat is stopped and the network is destroyed.
         mCellNetworkAgent.disconnect();
-        networkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        networkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         networkCallback.assertNoCallback();
         verifyClatdStop(null /* inOrder */, MOBILE_IFNAME);
         verify(mMockNetd).idletimerRemoveInterface(eq(MOBILE_IFNAME), anyInt(),
@@ -10455,7 +10512,7 @@
         // clat has been stopped, or the test will be flaky.
         ExpectedBroadcast b = expectConnectivityAction(TYPE_WIFI, DetailedState.DISCONNECTED);
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         b.expectBroadcast();
 
         verifyClatdStop(inOrder, iface);
@@ -10529,7 +10586,7 @@
         // Network switch
         mWiFiNetworkAgent.connect(true);
         networkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        networkCallback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        networkCallback.expectLosing(mCellNetworkAgent);
         networkCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         verify(mMockNetd, times(1)).idletimerAddInterface(eq(WIFI_IFNAME), anyInt(),
                 eq(Integer.toString(TRANSPORT_WIFI)));
@@ -10539,7 +10596,7 @@
         // Disconnect wifi and switch back to cell
         reset(mMockNetd);
         mWiFiNetworkAgent.disconnect();
-        networkCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        networkCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         assertNoCallbacks(networkCallback);
         verify(mMockNetd, times(1)).idletimerRemoveInterface(eq(WIFI_IFNAME), anyInt(),
                 eq(Integer.toString(TRANSPORT_WIFI)));
@@ -10553,7 +10610,7 @@
         mWiFiNetworkAgent.sendLinkProperties(wifiLp);
         mWiFiNetworkAgent.connect(true);
         networkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        networkCallback.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        networkCallback.expectLosing(mCellNetworkAgent);
         networkCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         verify(mMockNetd, times(1)).idletimerAddInterface(eq(WIFI_IFNAME), anyInt(),
                 eq(Integer.toString(TRANSPORT_WIFI)));
@@ -10563,7 +10620,7 @@
         // Disconnect cell
         reset(mMockNetd);
         mCellNetworkAgent.disconnect();
-        networkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        networkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         // LOST callback is triggered earlier than removing idle timer. Broadcast should also be
         // sent as network being switched. Ensure rule removal for cell will not be triggered
         // unexpectedly before network being removed.
@@ -10613,11 +10670,11 @@
         LinkProperties lp = new LinkProperties();
         lp.setTcpBufferSizes(testTcpBufferSizes);
         mCellNetworkAgent.sendLinkProperties(lp);
-        networkCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        networkCallback.expect(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         verifyTcpBufferSizeChange(testTcpBufferSizes);
         // Clean up.
         mCellNetworkAgent.disconnect();
-        networkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        networkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         networkCallback.assertNoCallback();
         mCm.unregisterNetworkCallback(networkCallback);
     }
@@ -12235,9 +12292,9 @@
         // While the default callback doesn't see the network before it's validated, the listen
         // sees the network come up and validate later
         allNetworksCb.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        allNetworksCb.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent);
+        allNetworksCb.expectLosing(mCellNetworkAgent);
         allNetworksCb.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
-        allNetworksCb.expectCallback(CallbackEntry.LOST, mCellNetworkAgent,
+        allNetworksCb.expect(CallbackEntry.LOST, mCellNetworkAgent,
                 TEST_LINGER_DELAY_MS * 2);
 
         // The cell network has disconnected (see LOST above) because it was outscored and
@@ -12249,7 +12306,7 @@
 
         // The cell network gets torn down right away.
         allNetworksCb.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
-        allNetworksCb.expectCallback(CallbackEntry.LOST, mCellNetworkAgent,
+        allNetworksCb.expect(CallbackEntry.LOST, mCellNetworkAgent,
                 TEST_NASCENT_DELAY_MS * 2);
         allNetworksCb.assertNoCallback();
 
@@ -12266,8 +12323,8 @@
 
         mWiFiNetworkAgent.disconnect();
 
-        allNetworksCb.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
-        mDefaultNetworkCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        allNetworksCb.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
+        mDefaultNetworkCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         mDefaultNetworkCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
 
         // Reconnect a WiFi network and make sure the cell network is still not torn down.
@@ -12280,10 +12337,8 @@
         // Now remove the reason to keep connected and make sure the network lingers and is
         // torn down.
         mCellNetworkAgent.setScore(new NetworkScore.Builder().setLegacyInt(30).build());
-        allNetworksCb.expectCallback(CallbackEntry.LOSING, mCellNetworkAgent,
-                TEST_NASCENT_DELAY_MS * 2);
-        allNetworksCb.expectCallback(CallbackEntry.LOST, mCellNetworkAgent,
-                TEST_LINGER_DELAY_MS * 2);
+        allNetworksCb.expectLosing(mCellNetworkAgent, TEST_NASCENT_DELAY_MS * 2);
+        allNetworksCb.expect(CallbackEntry.LOST, mCellNetworkAgent, TEST_LINGER_DELAY_MS * 2);
         mDefaultNetworkCallback.assertNoCallback();
 
         mCm.unregisterNetworkCallback(allNetworksCb);
@@ -13254,7 +13309,7 @@
         setOemNetworkPreferenceAgentConnected(TRANSPORT_ETHERNET, false);
 
         // At this point, with no network is available, the lost callback should trigger
-        defaultNetworkCallback.expectCallback(CallbackEntry.LOST, mEthernetNetworkAgent);
+        defaultNetworkCallback.expect(CallbackEntry.LOST, mEthernetNetworkAgent);
         otherUidDefaultCallback.assertNoCallback();
 
         // Confirm we can unregister without issues.
@@ -13300,7 +13355,7 @@
         otherUidDefaultCallback.assertNoCallback();
 
         // At this point, with no network is available, the lost callback should trigger
-        defaultNetworkCallback.expectCallback(CallbackEntry.LOST, mEthernetNetworkAgent);
+        defaultNetworkCallback.expect(CallbackEntry.LOST, mEthernetNetworkAgent);
         otherUidDefaultCallback.assertNoCallback();
 
         // Confirm we can unregister without issues.
@@ -13362,13 +13417,13 @@
         // Since the callback didn't use the per-app network, only the other UID gets a callback.
         // Because the preference specifies no fallback, it does not switch to cellular.
         defaultNetworkCallback.assertNoCallback();
-        otherUidDefaultCallback.expectCallback(CallbackEntry.LOST, mEthernetNetworkAgent);
+        otherUidDefaultCallback.expect(CallbackEntry.LOST, mEthernetNetworkAgent);
 
         // Now bring down the default network.
         setOemNetworkPreferenceAgentConnected(TRANSPORT_CELLULAR, false);
 
         // As this callback was tracking the default, this should now trigger.
-        defaultNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        defaultNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         otherUidDefaultCallback.assertNoCallback();
 
         // Confirm we can unregister without issues.
@@ -14482,7 +14537,7 @@
         mWiFiNetworkAgent.disconnect();
         bestMatchingCb.assertNoCallback();
         mCellNetworkAgent.disconnect();
-        bestMatchingCb.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        bestMatchingCb.expect(CallbackEntry.LOST, mCellNetworkAgent);
     }
 
     private UidRangeParcel[] uidRangeFor(final UserHandle handle) {
@@ -14627,7 +14682,7 @@
         if (allowFallback && !connectWorkProfileAgentAhead) {
             assertNoCallbacks(profileDefaultNetworkCallback);
         } else if (!connectWorkProfileAgentAhead) {
-            profileDefaultNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+            profileDefaultNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
             if (disAllowProfileDefaultNetworkCallback != null) {
                 assertNoCallbacks(disAllowProfileDefaultNetworkCallback);
             }
@@ -14697,10 +14752,10 @@
         // apps on this network see the appropriate callbacks, and the app on the work profile
         // doesn't because it continues to use the enterprise network.
         mCellNetworkAgent.disconnect();
-        mSystemDefaultNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
-        mDefaultNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        mSystemDefaultNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
+        mDefaultNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         if (disAllowProfileDefaultNetworkCallback != null) {
-            disAllowProfileDefaultNetworkCallback.expectCallback(
+            disAllowProfileDefaultNetworkCallback.expect(
                     CallbackEntry.LOST, mCellNetworkAgent);
         }
         profileDefaultNetworkCallback.assertNoCallback();
@@ -14722,7 +14777,7 @@
         // When the agent disconnects, test that the app on the work profile falls back to the
         // default network.
         workAgent.disconnect();
-        profileDefaultNetworkCallback.expectCallback(CallbackEntry.LOST, workAgent);
+        profileDefaultNetworkCallback.expect(CallbackEntry.LOST, workAgent);
         if (allowFallback) {
             profileDefaultNetworkCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
             if (disAllowProfileDefaultNetworkCallback != null) {
@@ -14739,14 +14794,14 @@
         inOrder.verify(mMockNetd).networkDestroy(workAgent.getNetwork().netId);
 
         mCellNetworkAgent.disconnect();
-        mSystemDefaultNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
-        mDefaultNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        mSystemDefaultNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
+        mDefaultNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         if (disAllowProfileDefaultNetworkCallback != null) {
-            disAllowProfileDefaultNetworkCallback.expectCallback(
+            disAllowProfileDefaultNetworkCallback.expect(
                     CallbackEntry.LOST, mCellNetworkAgent);
         }
         if (allowFallback) {
-            profileDefaultNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+            profileDefaultNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         }
 
         // Waiting for the handler to be idle before checking for networkDestroy is necessary
@@ -14789,7 +14844,7 @@
         // When the agent disconnects, test that the app on the work profile fall back to the
         // default network.
         workAgent2.disconnect();
-        profileDefaultNetworkCallback.expectCallback(CallbackEntry.LOST, workAgent2);
+        profileDefaultNetworkCallback.expect(CallbackEntry.LOST, workAgent2);
         if (disAllowProfileDefaultNetworkCallback != null) {
             assertNoCallbacks(disAllowProfileDefaultNetworkCallback);
         }
@@ -15399,11 +15454,11 @@
         workAgent4.disconnect();
         workAgent5.disconnect();
 
-        appCb1.expectCallback(CallbackEntry.LOST, workAgent1);
-        appCb2.expectCallback(CallbackEntry.LOST, workAgent2);
-        appCb3.expectCallback(CallbackEntry.LOST, workAgent3);
-        appCb4.expectCallback(CallbackEntry.LOST, workAgent4);
-        appCb5.expectCallback(CallbackEntry.LOST, workAgent5);
+        appCb1.expect(CallbackEntry.LOST, workAgent1);
+        appCb2.expect(CallbackEntry.LOST, workAgent2);
+        appCb3.expect(CallbackEntry.LOST, workAgent3);
+        appCb4.expect(CallbackEntry.LOST, workAgent4);
+        appCb5.expect(CallbackEntry.LOST, workAgent5);
 
         appCb1.expectAvailableCallbacksValidated(mCellNetworkAgent);
         appCb2.assertNoCallback();
@@ -15873,7 +15928,7 @@
         }
 
         mEthernetNetworkAgent.disconnect();
-        cb.expectCallback(CallbackEntry.LOST, mEthernetNetworkAgent);
+        cb.expect(CallbackEntry.LOST, mEthernetNetworkAgent);
         mCm.unregisterNetworkCallback(cb);
     }
 
@@ -15943,7 +15998,7 @@
         cb.assertNoCallback(TEST_CALLBACK_TIMEOUT_MS);
 
         mCellNetworkAgent.disconnect();
-        cb.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        cb.expect(CallbackEntry.LOST, mCellNetworkAgent);
         mCm.unregisterNetworkCallback(cb);
 
         // Must be unset before touching the transports, because remove and add transport types
@@ -16253,8 +16308,8 @@
         // callback with wifi network from fallback request.
         mCellNetworkAgent.disconnect();
         mDefaultNetworkCallback.assertNoCallback();
-        cellNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
-        mTestPackageDefaultNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        cellNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
+        mTestPackageDefaultNetworkCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         mTestPackageDefaultNetworkCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetworkForUid(TEST_PACKAGE_UID));
         inorder.verify(mMockNetd, times(1)).networkAddUidRangesParcel(wifiConfig);
@@ -16281,7 +16336,7 @@
         // Wifi network disconnected. mTestPackageDefaultNetworkCallback should not receive
         // any callback.
         mWiFiNetworkAgent.disconnect();
-        mDefaultNetworkCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        mDefaultNetworkCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         mDefaultNetworkCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         mTestPackageDefaultNetworkCallback.assertNoCallback();
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetworkForUid(TEST_PACKAGE_UID));
@@ -16530,7 +16585,7 @@
         // Disconnect wifi
         mWiFiNetworkAgent.disconnect();
         assertNoCallbacks(mProfileDefaultNetworkCallback, mTestPackageDefaultNetworkCallback);
-        mDefaultNetworkCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        mDefaultNetworkCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         mDefaultNetworkCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
     }
 
@@ -16768,9 +16823,9 @@
             mWiFiNetworkAgent.setNetworkCapabilities(wifiNc2, true /* sendToConnectivityService */);
             // The only thing changed in this CAPS is the BSSID, which can't be tested for in this
             // test because it's redacted.
-            wifiNetworkCallback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED,
+            wifiNetworkCallback.expect(CallbackEntry.NETWORK_CAPS_UPDATED,
                     mWiFiNetworkAgent);
-            mDefaultNetworkCallback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED,
+            mDefaultNetworkCallback.expect(CallbackEntry.NETWORK_CAPS_UPDATED,
                     mWiFiNetworkAgent);
             mWiFiNetworkAgent.setNetworkPortal(TEST_REDIRECT_URL, false /* isStrictMode */);
             mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), false);
@@ -16790,9 +16845,9 @@
 
             // Wi-Fi roaming from wifiNc2 to wifiNc1, and the network now has partial connectivity.
             mWiFiNetworkAgent.setNetworkCapabilities(wifiNc1, true);
-            wifiNetworkCallback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED,
+            wifiNetworkCallback.expect(CallbackEntry.NETWORK_CAPS_UPDATED,
                     mWiFiNetworkAgent);
-            mDefaultNetworkCallback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED,
+            mDefaultNetworkCallback.expect(CallbackEntry.NETWORK_CAPS_UPDATED,
                     mWiFiNetworkAgent);
             mWiFiNetworkAgent.setNetworkPartial();
             mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), false);
@@ -16815,7 +16870,7 @@
         // failures after roam are not ignored, this will cause cell to become the default network.
         // If they are ignored, this will not cause a switch until later.
         mWiFiNetworkAgent.setNetworkCapabilities(wifiNc2, true);
-        mDefaultNetworkCallback.expectCallback(CallbackEntry.NETWORK_CAPS_UPDATED,
+        mDefaultNetworkCallback.expect(CallbackEntry.NETWORK_CAPS_UPDATED,
                 mWiFiNetworkAgent);
         mWiFiNetworkAgent.setNetworkInvalid(false /* isStrictMode */);
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), false);
diff --git a/tests/unit/java/com/android/server/connectivity/FullScoreTest.kt b/tests/unit/java/com/android/server/connectivity/FullScoreTest.kt
index b39e960..3520c5b 100644
--- a/tests/unit/java/com/android/server/connectivity/FullScoreTest.kt
+++ b/tests/unit/java/com/android/server/connectivity/FullScoreTest.kt
@@ -152,4 +152,16 @@
         assertNotEquals(ns3, ns1)
         assertFalse(ns1.equals(ns4))
     }
+
+    @Test
+    fun testDescribeDifferences() {
+        val ns1 = FullScore((1L shl POLICY_EVER_EVALUATED) or (1L shl POLICY_IS_VALIDATED),
+                KEEP_CONNECTED_NONE)
+        val ns2 = FullScore((1L shl POLICY_IS_VALIDATED) or (1L shl POLICY_IS_VPN) or
+                (1L shl POLICY_IS_DESTROYED), KEEP_CONNECTED_NONE)
+        assertEquals("-EVER_EVALUATED+IS_DESTROYED+IS_VPN",
+                ns2.describeDifferencesFrom(ns1))
+        assertEquals("-IS_DESTROYED-IS_VPN+EVER_EVALUATED",
+                ns1.describeDifferencesFrom(ns2))
+    }
 }
diff --git a/tests/unit/java/com/android/server/connectivity/VpnTest.java b/tests/unit/java/com/android/server/connectivity/VpnTest.java
index 1c54651..39fd780 100644
--- a/tests/unit/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/unit/java/com/android/server/connectivity/VpnTest.java
@@ -20,6 +20,8 @@
 import static android.Manifest.permission.CONTROL_VPN;
 import static android.content.pm.PackageManager.PERMISSION_DENIED;
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
+import static android.net.ConnectivityDiagnosticsManager.ConnectivityDiagnosticsCallback;
+import static android.net.ConnectivityDiagnosticsManager.DataStallReport;
 import static android.net.ConnectivityManager.NetworkCallback;
 import static android.net.INetd.IF_STATE_DOWN;
 import static android.net.INetd.IF_STATE_UP;
@@ -76,6 +78,7 @@
 import android.content.pm.ServiceInfo;
 import android.content.pm.UserInfo;
 import android.content.res.Resources;
+import android.net.ConnectivityDiagnosticsManager;
 import android.net.ConnectivityManager;
 import android.net.INetd;
 import android.net.Ikev2VpnProfile;
@@ -115,6 +118,7 @@
 import android.os.ConditionVariable;
 import android.os.INetworkManagementService;
 import android.os.ParcelFileDescriptor;
+import android.os.PersistableBundle;
 import android.os.PowerWhitelistManager;
 import android.os.Process;
 import android.os.UserHandle;
@@ -245,6 +249,7 @@
     @Mock private Vpn.Ikev2SessionCreator mIkev2SessionCreator;
     @Mock private Vpn.VpnNetworkAgentWrapper mMockNetworkAgent;
     @Mock private ConnectivityManager mConnectivityManager;
+    @Mock private ConnectivityDiagnosticsManager mCdm;
     @Mock private IpSecService mIpSecService;
     @Mock private VpnProfileStore mVpnProfileStore;
     @Mock private ScheduledThreadPoolExecutor mExecutor;
@@ -283,6 +288,8 @@
         mockService(NotificationManager.class, Context.NOTIFICATION_SERVICE, mNotificationManager);
         mockService(ConnectivityManager.class, Context.CONNECTIVITY_SERVICE, mConnectivityManager);
         mockService(IpSecManager.class, Context.IPSEC_SERVICE, mIpSecManager);
+        mockService(ConnectivityDiagnosticsManager.class, Context.CONNECTIVITY_DIAGNOSTICS_SERVICE,
+                mCdm);
         when(mContext.getString(R.string.config_customVpnAlwaysOnDisconnectedDialogComponent))
                 .thenReturn(Resources.getSystem().getString(
                         R.string.config_customVpnAlwaysOnDisconnectedDialogComponent));
@@ -1925,6 +1932,56 @@
         verifyHandlingNetworkLoss(vpnSnapShot);
     }
 
+    private ConnectivityDiagnosticsCallback getConnectivityDiagCallback() {
+        final ArgumentCaptor<ConnectivityDiagnosticsCallback> cdcCaptor =
+                ArgumentCaptor.forClass(ConnectivityDiagnosticsCallback.class);
+        verify(mCdm).registerConnectivityDiagnosticsCallback(
+                any(), any(), cdcCaptor.capture());
+        return cdcCaptor.getValue();
+    }
+
+    private DataStallReport createDataStallReport() {
+        return new DataStallReport(TEST_NETWORK, 1234 /* reportTimestamp */,
+                1 /* detectionMethod */, new LinkProperties(), new NetworkCapabilities(),
+                new PersistableBundle());
+    }
+
+    private void verifyMobikeTriggered(List<Network> expected) {
+        final ArgumentCaptor<Network> networkCaptor = ArgumentCaptor.forClass(Network.class);
+        verify(mIkeSessionWrapper).setNetwork(networkCaptor.capture());
+        assertEquals(expected, Collections.singletonList(networkCaptor.getValue()));
+    }
+
+    @Test
+    public void testDataStallInIkev2VpnMobikeDisabled() throws Exception {
+        verifySetupPlatformVpn(
+                createIkeConfig(createIkeConnectInfo(), false /* isMobikeEnabled */));
+
+        doReturn(TEST_NETWORK).when(mMockNetworkAgent).getNetwork();
+        final ConnectivityDiagnosticsCallback connectivityDiagCallback =
+                getConnectivityDiagCallback();
+        final DataStallReport report = createDataStallReport();
+        connectivityDiagCallback.onDataStallSuspected(report);
+
+        // Should not trigger MOBIKE if MOBIKE is not enabled
+        verify(mIkeSessionWrapper, never()).setNetwork(any());
+    }
+
+    @Test
+    public void testDataStallInIkev2VpnMobikeEnabled() throws Exception {
+        final PlatformVpnSnapshot vpnSnapShot = verifySetupPlatformVpn(
+                createIkeConfig(createIkeConnectInfo(), true /* isMobikeEnabled */));
+
+        doReturn(TEST_NETWORK).when(mMockNetworkAgent).getNetwork();
+        final ConnectivityDiagnosticsCallback connectivityDiagCallback =
+                getConnectivityDiagCallback();
+        final DataStallReport report = createDataStallReport();
+        connectivityDiagCallback.onDataStallSuspected(report);
+
+        // Verify MOBIKE is triggered
+        verifyMobikeTriggered(vpnSnapShot.vpn.mNetworkCapabilities.getUnderlyingNetworks());
+    }
+
     @Test
     public void testStartRacoonNumericAddress() throws Exception {
         startRacoon("1.2.3.4", "1.2.3.4");
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index 2ceb00a..e8f30d6 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -75,12 +75,14 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.AdditionalMatchers.aryEq;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -393,6 +395,10 @@
         verify(mNetd).registerUnsolicitedEventListener(alertObserver.capture());
         mAlertObserver = alertObserver.getValue();
 
+        // Make augmentWithStackedInterfaces returns the interfaces that was passed to it.
+        doAnswer(inv -> ((String[]) inv.getArgument(0)).clone())
+                .when(mStatsFactory).augmentWithStackedInterfaces(any());
+
         // Catch TetheringEventCallback during systemReady().
         ArgumentCaptor<TetheringManager.TetheringEventCallback> tetheringEventCbCaptor =
                 ArgumentCaptor.forClass(TetheringManager.TetheringEventCallback.class);
@@ -1239,45 +1245,73 @@
 
     @Test
     public void testUidStatsForTransport() throws Exception {
-        // pretend that network comes online
+        // Setup both wifi and mobile networks, and set mobile network as the default interface.
         mockDefaultSettings();
-        NetworkStateSnapshot[] states = new NetworkStateSnapshot[] {buildWifiState()};
-        mockNetworkStatsSummary(buildEmptyStats());
         mockNetworkStatsUidDetail(buildEmptyStats());
 
-        mService.notifyNetworkStatus(NETWORKS_WIFI, states, getActiveIface(states),
+        final NetworkStateSnapshot mobileState = buildStateOfTransport(
+                NetworkCapabilities.TRANSPORT_CELLULAR, TYPE_MOBILE,
+                TEST_IFACE2, IMSI_1, null /* wifiNetworkKey */,
+                false /* isTemporarilyNotMetered */, false /* isRoaming */);
+
+        final NetworkStateSnapshot[] states = new NetworkStateSnapshot[] {
+                mobileState, buildWifiState()};
+        mService.notifyNetworkStatus(NETWORKS_MOBILE, states, getActiveIface(states),
                 new UnderlyingNetworkInfo[0]);
+        setMobileRatTypeAndWaitForIdle(TelephonyManager.NETWORK_TYPE_LTE);
 
-        NetworkStats.Entry entry1 = new NetworkStats.Entry(
+        // Mock traffic on wifi network.
+        final NetworkStats.Entry entry1 = new NetworkStats.Entry(
                 TEST_IFACE, UID_RED, SET_DEFAULT, TAG_NONE, METERED_NO, ROAMING_NO,
-                DEFAULT_NETWORK_NO, 50L, 5L, 50L, 5L, 0L);
-        NetworkStats.Entry entry2 = new NetworkStats.Entry(
+                DEFAULT_NETWORK_NO, 50L, 5L, 50L, 5L, 1L);
+        final NetworkStats.Entry entry2 = new NetworkStats.Entry(
                 TEST_IFACE, UID_RED, SET_DEFAULT, 0xF00D, METERED_NO, ROAMING_NO,
-                DEFAULT_NETWORK_NO, 50L, 5L, 50L, 5L, 0L);
-        NetworkStats.Entry entry3 = new NetworkStats.Entry(
+                DEFAULT_NETWORK_NO, 50L, 5L, 50L, 5L, 1L);
+        final NetworkStats.Entry entry3 = new NetworkStats.Entry(
                 TEST_IFACE, UID_BLUE, SET_DEFAULT, 0xBEEF, METERED_NO, ROAMING_NO,
-                DEFAULT_NETWORK_NO, 1024L, 8L, 512L, 4L, 0L);
+                DEFAULT_NETWORK_NO, 1024L, 8L, 512L, 4L, 2L);
 
+        final TetherStatsParcel[] emptyTetherStats = {};
+        // The interfaces that expect to be used to query the stats.
+        final String[] wifiIfaces = {TEST_IFACE};
         incrementCurrentTime(HOUR_IN_MILLIS);
         mockDefaultSettings();
-        mockNetworkStatsSummary(buildEmptyStats());
         mockNetworkStatsUidDetail(new NetworkStats(getElapsedRealtime(), 3)
                 .insertEntry(entry1)
                 .insertEntry(entry2)
-                .insertEntry(entry3));
+                .insertEntry(entry3), emptyTetherStats, wifiIfaces);
+
+        // getUidStatsForTransport (through getNetworkStatsUidDetail) adds all operation counts
+        // with active interface, and the interface here is mobile interface, so this test makes
+        // sure these operations are not surfaced in getUidStatsForTransport if the transport
+        // doesn't match them.
         mService.incrementOperationCount(UID_RED, 0xF00D, 1);
+        final NetworkStats wifiStats = mService.getUidStatsForTransport(
+                NetworkCapabilities.TRANSPORT_WIFI);
 
-        NetworkStats stats = mService.getUidStatsForTransport(NetworkCapabilities.TRANSPORT_WIFI);
+        assertEquals(3, wifiStats.size());
+        // The iface field of the returned stats should be null because getUidStatsForTransport
+        // clears the interface fields before it returns the result.
+        assertValues(wifiStats, null /* iface */, UID_RED, SET_DEFAULT, TAG_NONE,
+                METERED_NO, ROAMING_NO, METERED_NO, 50L, 5L, 50L, 5L, 1L);
+        assertValues(wifiStats, null /* iface */, UID_RED, SET_DEFAULT, 0xF00D,
+                METERED_NO, ROAMING_NO, METERED_NO, 50L, 5L, 50L, 5L, 1L);
+        assertValues(wifiStats, null /* iface */, UID_BLUE, SET_DEFAULT, 0xBEEF,
+                METERED_NO, ROAMING_NO, METERED_NO, 1024L, 8L, 512L, 4L, 2L);
 
-        assertEquals(3, stats.size());
-        entry1.operations = 1;
-        entry1.iface = null;
-        assertEquals(entry1, stats.getValues(0, null));
-        entry2.operations = 1;
-        entry2.iface = null;
-        assertEquals(entry2, stats.getValues(1, null));
-        entry3.iface = null;
-        assertEquals(entry3, stats.getValues(2, null));
+        final String[] mobileIfaces = {TEST_IFACE2};
+        mockNetworkStatsUidDetail(buildEmptyStats(), emptyTetherStats, mobileIfaces);
+        final NetworkStats mobileStats = mService.getUidStatsForTransport(
+                NetworkCapabilities.TRANSPORT_CELLULAR);
+
+        assertEquals(2, mobileStats.size());
+        // Verify the operation count stats that caused by incrementOperationCount only appears
+        // on the mobile interface since incrementOperationCount attributes them onto the active
+        // interface.
+        assertValues(mobileStats, null /* iface */, UID_RED, SET_DEFAULT, 0xF00D,
+                METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 0L, 0L, 0L, 0L, 1);
+        assertValues(mobileStats, null /* iface */, UID_RED, SET_DEFAULT, TAG_NONE,
+                METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 0L, 0L, 0L, 0L, 1);
     }
 
     @Test
@@ -1468,7 +1502,7 @@
                 {buildTetherStatsParcel(TEST_IFACE, 1408L, 10L, 256L, 1L, 0)};
 
         mockNetworkStatsSummary(swIfaceStats);
-        mockNetworkStatsUidDetail(localUidStats, tetherStatsParcels);
+        mockNetworkStatsUidDetail(localUidStats, tetherStatsParcels, INTERFACES_ALL);
         forcePollAndWaitForIdle();
 
         // verify service recorded history
@@ -2235,13 +2269,14 @@
 
     private void mockNetworkStatsUidDetail(NetworkStats detail) throws Exception {
         final TetherStatsParcel[] tetherStatsParcels = {};
-        mockNetworkStatsUidDetail(detail, tetherStatsParcels);
+        mockNetworkStatsUidDetail(detail, tetherStatsParcels, INTERFACES_ALL);
     }
 
     private void mockNetworkStatsUidDetail(NetworkStats detail,
-            TetherStatsParcel[] tetherStatsParcels) throws Exception {
+            TetherStatsParcel[] tetherStatsParcels, String[] ifaces) throws Exception {
+
         doReturn(detail).when(mStatsFactory)
-                .readNetworkStatsDetail(UID_ALL, INTERFACES_ALL, TAG_ALL);
+                .readNetworkStatsDetail(eq(UID_ALL), aryEq(ifaces), eq(TAG_ALL));
 
         // also include tethering details, since they are folded into UID
         doReturn(tetherStatsParcels).when(mNetd).tetherGetStats();