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/clatd.c b/bpf_progs/clatd.c
index a2214dc..b8c6131 100644
--- a/bpf_progs/clatd.c
+++ b/bpf_progs/clatd.c
@@ -38,8 +38,19 @@
 #include "bpf_shared.h"
 #include "clat_mark.h"
 
-// From kernel:include/net/ip.h
-#define IP_DF 0x4000  // Flag: "Don't Fragment"
+// IP flags. (from kernel's include/net/ip.h)
+#define IP_CE      0x8000  // Flag: "Congestion" (really reserved 'evil bit')
+#define IP_DF      0x4000  // Flag: "Don't Fragment"
+#define IP_MF      0x2000  // Flag: "More Fragments"
+#define IP_OFFSET  0x1FFF  // "Fragment Offset" part
+
+// from kernel's include/net/ipv6.h
+struct frag_hdr {
+    __u8   nexthdr;
+    __u8   reserved;        // always zero
+    __be16 frag_off;        // 13 bit offset, 2 bits zero, 1 bit "More Fragments"
+    __be32 identification;
+};
 
 DEFINE_BPF_MAP_GRW(clat_ingress6_map, HASH, ClatIngress6Key, ClatIngress6Value, 16, AID_SYSTEM)
 
@@ -89,7 +100,32 @@
 
     if (!v) return TC_ACT_PIPE;
 
-    switch (ip6->nexthdr) {
+    __u8 proto = ip6->nexthdr;
+    __be16 ip_id = 0;
+    __be16 frag_off = htons(IP_DF);
+    __u16 tot_len = ntohs(ip6->payload_len) + sizeof(struct iphdr);  // cannot overflow, see above
+
+    if (proto == IPPROTO_FRAGMENT) {
+        // Must have (ethernet and) ipv6 header and ipv6 fragment extension header
+        if (data + l2_header_size + sizeof(*ip6) + sizeof(struct frag_hdr) > data_end)
+            return TC_ACT_PIPE;
+        const struct frag_hdr *frag = (const struct frag_hdr *)(ip6 + 1);
+        proto = frag->nexthdr;
+        // Trivial hash of 32-bit IPv6 ID field into 16-bit IPv4 field.
+        ip_id = (frag->identification) ^ (frag->identification >> 16);
+        // Conversion of 16-bit IPv6 frag offset to 16-bit IPv4 frag offset field.
+        // IPv6 is '13 bits of offset in multiples of 8' + 2 zero bits + more fragment bit
+        // IPv4 is zero bit + don't frag bit + more frag bit + '13 bits of offset in multiples of 8'
+        frag_off = ntohs(frag->frag_off);
+        frag_off = ((frag_off & 1) << 13) | (frag_off >> 3);
+        frag_off = htons(frag_off);
+        // Note that by construction tot_len is guaranteed to not underflow here
+        tot_len -= sizeof(struct frag_hdr);
+        // This is a badly formed IPv6 packet with less payload than the size of an IPv6 Frag EH
+        if (tot_len < sizeof(struct iphdr)) return TC_ACT_PIPE;
+    }
+
+    switch (proto) {
         case IPPROTO_TCP:  // For TCP & UDP the checksum neutrality of the chosen IPv6
         case IPPROTO_UDP:  // address means there is no need to update their checksums.
         case IPPROTO_GRE:  // We do not need to bother looking at GRE/ESP headers,
@@ -114,14 +150,14 @@
             .version = 4,                                                      // u4
             .ihl = sizeof(struct iphdr) / sizeof(__u32),                       // u4
             .tos = (ip6->priority << 4) + (ip6->flow_lbl[0] >> 4),             // u8
-            .tot_len = htons(ntohs(ip6->payload_len) + sizeof(struct iphdr)),  // u16
-            .id = 0,                                                           // u16
-            .frag_off = htons(IP_DF),                                          // u16
+            .tot_len = htons(tot_len),                                         // be16
+            .id = ip_id,                                                       // be16
+            .frag_off = frag_off,                                              // be16
             .ttl = ip6->hop_limit,                                             // u8
-            .protocol = ip6->nexthdr,                                          // u8
+            .protocol = proto,                                                 // u8
             .check = 0,                                                        // u16
-            .saddr = ip6->saddr.in6_u.u6_addr32[3],                            // u32
-            .daddr = v->local4.s_addr,                                         // u32
+            .saddr = ip6->saddr.in6_u.u6_addr32[3],                            // be32
+            .daddr = v->local4.s_addr,                                         // be32
     };
 
     // Calculate the IPv4 one's complement checksum of the IPv4 header.
@@ -171,6 +207,13 @@
     //   return -ENOTSUPP;
     bpf_csum_update(skb, sum6);
 
+    if (frag_off != htons(IP_DF)) {
+        // If we're converting an IPv6 Fragment, we need to trim off 8 more bytes
+        // We're beyond recovery on error here... but hard to imagine how this could fail.
+        if (bpf_skb_adjust_room(skb, -(__s32)sizeof(struct frag_hdr), BPF_ADJ_ROOM_NET, /*flags*/0))
+            return TC_ACT_SHOT;
+    }
+
     // bpf_skb_change_proto() invalidates all pointers - reload them.
     data = (void*)(long)skb->data;
     data_end = (void*)(long)skb->data_end;
@@ -211,11 +254,6 @@
 
 DEFINE_BPF_MAP_GRW(clat_egress4_map, HASH, ClatEgress4Key, ClatEgress4Value, 16, AID_SYSTEM)
 
-DEFINE_BPF_PROG("schedcls/egress4/clat_ether", AID_ROOT, AID_SYSTEM, sched_cls_egress4_clat_ether)
-(struct __sk_buff* skb) {
-    return TC_ACT_PIPE;
-}
-
 DEFINE_BPF_PROG("schedcls/egress4/clat_rawip", AID_ROOT, AID_SYSTEM, sched_cls_egress4_clat_rawip)
 (struct __sk_buff* skb) {
     // Must be meta-ethernet IPv4 frame
diff --git a/service/jni/com_android_server_connectivity_ClatCoordinator.cpp b/service/jni/com_android_server_connectivity_ClatCoordinator.cpp
index de0e20a..5cd6e5d 100644
--- a/service/jni/com_android_server_connectivity_ClatCoordinator.cpp
+++ b/service/jni/com_android_server_connectivity_ClatCoordinator.cpp
@@ -87,7 +87,8 @@
 
 // Picks a random interface ID that is checksum neutral with the IPv4 address and the NAT64 prefix.
 jstring com_android_server_connectivity_ClatCoordinator_generateIpv6Address(
-        JNIEnv* env, jobject clazz, jstring ifaceStr, jstring v4Str, jstring prefix64Str) {
+        JNIEnv* env, jobject clazz, jstring ifaceStr, jstring v4Str, jstring prefix64Str,
+        jint mark) {
     ScopedUtfChars iface(env, ifaceStr);
     ScopedUtfChars addr4(env, v4Str);
     ScopedUtfChars prefix64(env, prefix64Str);
@@ -111,7 +112,7 @@
     }
 
     in6_addr v6;
-    if (net::clat::generateIpv6Address(iface.c_str(), v4, nat64Prefix, &v6)) {
+    if (net::clat::generateIpv6Address(iface.c_str(), v4, nat64Prefix, &v6, mark)) {
         jniThrowExceptionFmt(env, "java/io/IOException",
                              "Unable to find global source address on %s for %s", iface.c_str(),
                              prefix64.c_str());
@@ -447,7 +448,7 @@
         {"native_selectIpv4Address", "(Ljava/lang/String;I)Ljava/lang/String;",
          (void*)com_android_server_connectivity_ClatCoordinator_selectIpv4Address},
         {"native_generateIpv6Address",
-         "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;",
+         "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;I)Ljava/lang/String;",
          (void*)com_android_server_connectivity_ClatCoordinator_generateIpv6Address},
         {"native_createTunInterface", "(Ljava/lang/String;)I",
          (void*)com_android_server_connectivity_ClatCoordinator_createTunInterface},
diff --git a/service/native/libs/libclat/clatutils.cpp b/service/native/libs/libclat/clatutils.cpp
index 4a125ba..be86612 100644
--- a/service/native/libs/libclat/clatutils.cpp
+++ b/service/native/libs/libclat/clatutils.cpp
@@ -126,10 +126,19 @@
 
 // Picks a random interface ID that is checksum neutral with the IPv4 address and the NAT64 prefix.
 int generateIpv6Address(const char* iface, const in_addr v4, const in6_addr& nat64Prefix,
-                        in6_addr* v6) {
+                        in6_addr* v6, uint32_t mark) {
     int s = socket(AF_INET6, SOCK_DGRAM | SOCK_CLOEXEC, 0);
     if (s == -1) return -errno;
 
+    // Socket's mark affects routing decisions (network selection)
+    // An fwmark is necessary for clat to bypass the VPN during initialization.
+    if ((mark != MARK_UNSET) && setsockopt(s, SOL_SOCKET, SO_MARK, &mark, sizeof(mark))) {
+        int ret = errno;
+        ALOGE("setsockopt(SOL_SOCKET, SO_MARK) failed: %s", strerror(errno));
+        close(s);
+        return -ret;
+    }
+
     if (setsockopt(s, SOL_SOCKET, SO_BINDTODEVICE, iface, strlen(iface) + 1) == -1) {
         close(s);
         return -errno;
diff --git a/service/native/libs/libclat/clatutils_test.cpp b/service/native/libs/libclat/clatutils_test.cpp
index 8cca1f4..abd4e81 100644
--- a/service/native/libs/libclat/clatutils_test.cpp
+++ b/service/native/libs/libclat/clatutils_test.cpp
@@ -202,7 +202,8 @@
     in6_addr v6;
     EXPECT_EQ(1, inet_pton(AF_INET6, "::", &v6));  // initialize as zero
 
-    EXPECT_EQ(-ENETUNREACH, generateIpv6Address(tun.name().c_str(), v4, nat64Prefix, &v6));
+    // 0u is MARK_UNSET
+    EXPECT_EQ(-ENETUNREACH, generateIpv6Address(tun.name().c_str(), v4, nat64Prefix, &v6, 0u));
     EXPECT_TRUE(IN6_IS_ADDR_ULA(&v6));
 
     tun.destroy();
diff --git a/service/native/libs/libclat/include/libclat/clatutils.h b/service/native/libs/libclat/include/libclat/clatutils.h
index 812c86e..991b193 100644
--- a/service/native/libs/libclat/include/libclat/clatutils.h
+++ b/service/native/libs/libclat/include/libclat/clatutils.h
@@ -24,7 +24,7 @@
 in_addr_t selectIpv4Address(const in_addr ip, int16_t prefixlen);
 void makeChecksumNeutral(in6_addr* v6, const in_addr v4, const in6_addr& nat64Prefix);
 int generateIpv6Address(const char* iface, const in_addr v4, const in6_addr& nat64Prefix,
-                        in6_addr* v6);
+                        in6_addr* v6, uint32_t mark);
 int detect_mtu(const struct in6_addr* plat_subnet, uint32_t plat_suffix, uint32_t mark);
 int configure_packet_socket(int sock, in6_addr* addr, int ifindex);
 
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index e732001..195f046 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -368,7 +368,7 @@
     private static final int DEFAULT_NASCENT_DELAY_MS = 5_000;
 
     // The maximum value for the blocking validation result, in milliseconds.
-    public static final int MAX_VALIDATION_FAILURE_BLOCKING_TIME_MS = 10000;
+    public static final int MAX_VALIDATION_IGNORE_AFTER_ROAM_TIME_MS = 10000;
 
     // The maximum number of network request allowed per uid before an exception is thrown.
     @VisibleForTesting
@@ -4241,12 +4241,18 @@
         return nai.isCreated() && !nai.isDestroyed();
     }
 
-    private boolean shouldIgnoreValidationFailureAfterRoam(NetworkAgentInfo nai) {
+    @VisibleForTesting
+    boolean shouldIgnoreValidationFailureAfterRoam(NetworkAgentInfo nai) {
         // T+ devices should use unregisterAfterReplacement.
         if (SdkLevel.isAtLeastT()) return false;
+
+        // If the network never roamed, return false. The check below is not sufficient if time
+        // since boot is less than blockTimeOut, though that's extremely unlikely to happen.
+        if (nai.lastRoamTime == 0) return false;
+
         final long blockTimeOut = Long.valueOf(mResources.get().getInteger(
                 R.integer.config_validationFailureAfterRoamIgnoreTimeMillis));
-        if (blockTimeOut <= MAX_VALIDATION_FAILURE_BLOCKING_TIME_MS
+        if (blockTimeOut <= MAX_VALIDATION_IGNORE_AFTER_ROAM_TIME_MS
                 && blockTimeOut >= 0) {
             final long currentTimeMs = SystemClock.elapsedRealtime();
             long timeSinceLastRoam = currentTimeMs - nai.lastRoamTime;
@@ -9774,8 +9780,8 @@
         return ((VpnTransportInfo) ti).getType();
     }
 
-    private void maybeUpdateWifiRoamTimestamp(NetworkAgentInfo nai, NetworkCapabilities nc) {
-        if (nai == null) return;
+    private void maybeUpdateWifiRoamTimestamp(@NonNull NetworkAgentInfo nai,
+            @NonNull NetworkCapabilities nc) {
         final TransportInfo prevInfo = nai.networkCapabilities.getTransportInfo();
         final TransportInfo newInfo = nc.getTransportInfo();
         if (!(prevInfo instanceof WifiInfo) || !(newInfo instanceof WifiInfo)) {
diff --git a/service/src/com/android/server/connectivity/ClatCoordinator.java b/service/src/com/android/server/connectivity/ClatCoordinator.java
index d7c3287..42b3827 100644
--- a/service/src/com/android/server/connectivity/ClatCoordinator.java
+++ b/service/src/com/android/server/connectivity/ClatCoordinator.java
@@ -113,12 +113,12 @@
         return "/sys/fs/bpf/net_shared/map_clatd_clat_" + which + "_map";
     }
 
-    private static String makeProgPath(boolean ingress, boolean ether) {
-        String path = "/sys/fs/bpf/net_shared/prog_clatd_schedcls_"
-                + (ingress ? "ingress6" : "egress4")
-                + "_clat_"
+    private static final String CLAT_EGRESS4_RAWIP_PROG_PATH =
+            "/sys/fs/bpf/net_shared/prog_clatd_schedcls_egress4_clat_rawip";
+
+    private static String makeIngressProgPath(boolean ether) {
+        return "/sys/fs/bpf/net_shared/prog_clatd_schedcls_ingress6_clat_"
                 + (ether ? "ether" : "rawip");
-        return path;
     }
 
     @NonNull
@@ -183,8 +183,8 @@
          */
         @NonNull
         public String generateIpv6Address(@NonNull String iface, @NonNull String v4,
-                @NonNull String prefix64) throws IOException {
-            return native_generateIpv6Address(iface, v4, prefix64);
+                @NonNull String prefix64, int mark) throws IOException {
+            return native_generateIpv6Address(iface, v4, prefix64, mark);
         }
 
         /**
@@ -478,7 +478,7 @@
             // tc filter add dev .. egress prio 4 protocol ip bpf object-pinned /sys/fs/bpf/...
             // direct-action
             mDeps.tcFilterAddDevBpf(tracker.v4ifIndex, EGRESS, (short) PRIO_CLAT, (short) ETH_P_IP,
-                    makeProgPath(EGRESS, RAWIP));
+                    CLAT_EGRESS4_RAWIP_PROG_PATH);
         } catch (IOException e) {
             Log.e(TAG, "tc filter add dev (" + tracker.v4ifIndex + "[" + tracker.v4iface
                     + "]) egress prio PRIO_CLAT protocol ip failure: " + e);
@@ -504,7 +504,7 @@
             // tc filter add dev .. ingress prio 4 protocol ipv6 bpf object-pinned /sys/fs/bpf/...
             // direct-action
             mDeps.tcFilterAddDevBpf(tracker.ifIndex, INGRESS, (short) PRIO_CLAT,
-                    (short) ETH_P_IPV6, makeProgPath(INGRESS, isEthernet));
+                    (short) ETH_P_IPV6, makeIngressProgPath(isEthernet));
         } catch (IOException e) {
             Log.e(TAG, "tc filter add dev (" + tracker.ifIndex + "[" + tracker.iface
                     + "]) ingress prio PRIO_CLAT protocol ipv6 failure: " + e);
@@ -629,10 +629,11 @@
         }
 
         // [2] Generate a checksum-neutral IID.
+        final Integer fwmark = getFwmark(netId);
         final String pfx96Str = nat64Prefix.getAddress().getHostAddress();
         final String v6Str;
         try {
-            v6Str = mDeps.generateIpv6Address(iface, v4Str, pfx96Str);
+            v6Str = mDeps.generateIpv6Address(iface, v4Str, pfx96Str, fwmark);
         } catch (IOException e) {
             throw new IOException("no IPv6 addresses were available for clat: " + e);
         }
@@ -676,7 +677,6 @@
         }
 
         // Detect ipv4 mtu.
-        final Integer fwmark = getFwmark(netId);
         final int detectedMtu;
         try {
             detectedMtu = mDeps.detectMtu(pfx96Str,
@@ -931,7 +931,7 @@
     private static native String native_selectIpv4Address(String v4addr, int prefixlen)
             throws IOException;
     private static native String native_generateIpv6Address(String iface, String v4,
-            String prefix64) throws IOException;
+            String prefix64, int mark) throws IOException;
     private static native int native_createTunInterface(String tuniface) throws IOException;
     private static native int native_detectMtu(String platSubnet, int platSuffix, int mark)
             throws IOException;
diff --git a/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java b/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java
index 6b2a1ee..e2821cb 100644
--- a/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java
@@ -21,12 +21,14 @@
 
 import static androidx.test.InstrumentationRegistry.getContext;
 
+import static com.android.compatibility.common.util.BatteryUtils.hasBattery;
 import static com.android.compatibility.common.util.SystemUtil.runShellCommand;
 import static com.android.testutils.MiscAsserts.assertThrows;
 import static com.android.testutils.TestPermissionUtil.runAsShell;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
+import static org.junit.Assume.assumeTrue;
 
 import android.content.Context;
 import android.content.pm.PackageManager;
@@ -97,6 +99,7 @@
     @RequiresDevice // Virtual hardware does not support wifi battery stats
     public void testReportNetworkInterfaceForTransports() throws Exception {
         try {
+            assumeTrue("Battery is not present. Ignore test.", hasBattery());
             // Simulate the device being unplugged from charging.
             executeShellCommand("cmd battery unplug");
             executeShellCommand("cmd battery set status " + BATTERY_STATUS_DISCHARGING);
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index 4887a78..218eb04 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -2621,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
@@ -2631,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 */ () -> {
@@ -2669,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 c2ab6c0..7e91478 100644
--- a/tests/cts/net/src/android/net/cts/EthernetManagerTest.kt
+++ b/tests/cts/net/src/android/net/cts/EthernetManagerTest.kt
@@ -74,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
@@ -255,7 +254,7 @@
         }
 
         fun <T : CallbackEntry> expectCallback(expected: T): T {
-            val event = pollForNextCallback()
+            val event = pollOrThrow()
             assertEquals(expected, event)
             return event as T
         }
@@ -272,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")
         }
 
@@ -525,13 +524,6 @@
         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>().network
-
-    private fun TestableNetworkCallback.expectLost(n: Network = anyNetwork()) =
-            expectCallback<Lost>(n)
-
     // b/233534110: eventuallyExpect<Lost>() does not advance ReadHead, use
     // eventuallyExpect(Lost::class) instead.
     private fun TestableNetworkCallback.eventuallyExpectLost(n: Network? = null) =
@@ -543,12 +535,10 @@
         }
 
     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.
@@ -702,7 +692,7 @@
         listenerCb.assertNeverAvailable()
 
         val cb = requestNetwork(ETH_REQUEST)
-        val network = cb.expectAvailable()
+        val network = cb.expect<Available>().network
 
         cb.assertNeverLost()
         releaseRequest(cb)
@@ -719,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()
@@ -734,7 +724,7 @@
 
         val cb = requestNetwork(ETH_REQUEST.copyWithEthernetSpecifier(iface2.name))
 
-        val network = cb.expectAvailable()
+        cb.expect<Available>()
         cb.expectCapabilitiesWithInterfaceName(iface2.name)
 
         removeInterface(iface1)
@@ -748,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()
@@ -757,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
@@ -769,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()
@@ -789,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)
@@ -815,7 +805,7 @@
         // cb.assertNeverAvailable()
 
         iface.setCarrierEnabled(true)
-        cb.expectAvailable()
+        cb.expect<Available>()
 
         iface.setCarrierEnabled(false)
         cb.eventuallyExpectLost()
@@ -860,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
@@ -894,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)
@@ -905,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)
@@ -915,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)
@@ -932,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.
@@ -948,8 +938,8 @@
 
         // 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) }
 
@@ -975,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)
     }
@@ -986,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 867d672..6df71c8 100644
--- a/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
+++ b/tests/cts/net/src/android/net/cts/NetworkAgentTest.kt
@@ -403,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()
         }
@@ -454,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)
@@ -532,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)
         }
@@ -577,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
     }
@@ -661,7 +661,7 @@
         }
 
         unregister(agent)
-        callback.expectCallback<Lost>(agent.network!!)
+        callback.expect<Lost>(agent.network!!)
     }
 
     private fun unregister(agent: TestableNetworkAgent) {
@@ -834,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
@@ -908,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
     }
@@ -1222,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))
@@ -1250,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:
@@ -1263,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>()
 
@@ -1286,8 +1286,8 @@
         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>()
 
@@ -1334,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)
         }
@@ -1363,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>()
     }
 
@@ -1382,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/mts/bpf_existence_test.cpp b/tests/mts/bpf_existence_test.cpp
index c7e8b97..bb07a98 100644
--- a/tests/mts/bpf_existence_test.cpp
+++ b/tests/mts/bpf_existence_test.cpp
@@ -100,7 +100,6 @@
     NETD "map_netd_uid_counterset_map",
     NETD "map_netd_uid_owner_map",
     NETD "map_netd_uid_permission_map",
-    SHARED "prog_clatd_schedcls_egress4_clat_ether",
     SHARED "prog_clatd_schedcls_egress4_clat_rawip",
     SHARED "prog_clatd_schedcls_ingress6_clat_ether",
     SHARED "prog_clatd_schedcls_ingress6_clat_rawip",
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index d244d63..2d8cf80 100755
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -2336,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
 
@@ -2408,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);
@@ -2424,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);
@@ -2435,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
@@ -2481,7 +2481,7 @@
         assertFalse(isForegroundNetwork(mCellNetworkAgent));
 
         mCellNetworkAgent.disconnect();
-        bgMobileListenCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        bgMobileListenCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         fgMobileListenCallback.assertNoCallback();
 
         mCm.unregisterNetworkCallback(wifiListenCallback);
@@ -2726,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());
@@ -2762,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();
@@ -2950,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 {
@@ -2966,7 +2966,7 @@
         }
 
         public CallbackEntry.Losing expectLosing(final HasNetwork n, final long timeoutMs) {
-            final CallbackEntry.Losing losing = expectCallback(CallbackEntry.LOSING, n, 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
@@ -2996,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);
         }
     }
 
@@ -3091,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);
 
@@ -3129,13 +3129,13 @@
         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);
     }
 
@@ -3284,8 +3284,8 @@
         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());
 
@@ -3315,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.
@@ -3324,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());
 
@@ -3359,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());
@@ -3378,12 +3378,12 @@
         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());
 
@@ -3414,8 +3414,8 @@
         // 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());
 
@@ -3439,7 +3439,7 @@
         // 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();
@@ -3455,13 +3455,13 @@
         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);
@@ -3573,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);
     }
 
@@ -3758,7 +3758,7 @@
         // (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);
@@ -3766,13 +3766,13 @@
         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);
@@ -3785,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(
@@ -3804,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());
@@ -3829,7 +3829,7 @@
             assertNotEquals(0L, nai.getFirstEvaluationConcludedTime());
         }
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
 
         mCm.unregisterNetworkCallback(callback);
     }
@@ -4236,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();
@@ -4254,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);
 
@@ -4280,7 +4280,7 @@
         // 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
@@ -4306,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.
@@ -4325,7 +4325,7 @@
                 NET_CAPABILITY_PARTIAL_CONNECTIVITY | NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         expectUnvalidationCheckWillNotNotify(mWiFiNetworkAgent);
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        callback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         verifyNoMoreInteractions(mNotificationManager);
     }
 
@@ -4408,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.
@@ -4422,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);
@@ -4431,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
@@ -4463,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.
@@ -4491,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);
@@ -5032,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());
@@ -5051,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());
 
@@ -5070,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());
@@ -5099,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();
 
@@ -5107,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());
 
@@ -5123,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());
 
@@ -5194,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));
@@ -5207,7 +5207,7 @@
         otherUidCallback.assertNoCallback();
 
         mMockVpn.disconnect();
-        includeOtherUidsCallback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        includeOtherUidsCallback.expect(CallbackEntry.LOST, mMockVpn);
         callback.assertNoCallback();
         otherUidCallback.assertNoCallback();
     }
@@ -5352,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));
@@ -5375,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));
@@ -5383,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));
 
@@ -5543,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();
@@ -5566,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();
@@ -5743,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);
@@ -5863,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.
@@ -5914,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.
@@ -5946,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.
@@ -6012,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();
@@ -6032,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);
@@ -6116,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();
@@ -6198,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));
         }
     }
 
@@ -6271,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() {
@@ -7081,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();
@@ -7128,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),
@@ -7148,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),
@@ -7156,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());
@@ -7181,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.
@@ -7210,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.
@@ -7254,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);
@@ -7264,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
@@ -7281,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.
@@ -7310,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...
@@ -7346,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);
     }
@@ -7383,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);
@@ -7394,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);
     }
@@ -7419,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;
@@ -7444,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);
@@ -7467,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
@@ -7648,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());
@@ -7684,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());
@@ -7717,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());
@@ -7740,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());
@@ -7767,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());
@@ -7779,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());
@@ -7792,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());
@@ -7984,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,
@@ -8021,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.
@@ -8035,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));
 
@@ -8081,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())
@@ -8099,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())
@@ -8136,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())
@@ -8152,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())
@@ -8232,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
@@ -8243,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));
@@ -8255,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());
 
@@ -8331,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);
@@ -8379,7 +8435,7 @@
         callback.assertNoCallback();
 
         mMockVpn.disconnect();
-        callback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        callback.expect(CallbackEntry.LOST, mMockVpn);
         callback.expectAvailableCallbacksValidated(mEthernetNetworkAgent);
     }
 
@@ -8513,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);
@@ -8522,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(
@@ -8558,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.
@@ -8570,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
@@ -8711,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)
@@ -9135,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.
@@ -9149,7 +9205,7 @@
         defaultCallback.assertNoCallback();
 
         mCellNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        defaultCallback.expect(CallbackEntry.LOST, mCellNetworkAgent);
         defaultCallback.assertNoCallback();
 
         mCm.unregisterNetworkCallback(defaultCallback);
@@ -9213,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));
         }
@@ -9421,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();
@@ -9573,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());
@@ -9584,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
@@ -9664,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);
 
@@ -9698,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();
 
@@ -9711,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();
@@ -9888,7 +9943,7 @@
             reset(mMockNetd);
 
             mCellNetworkAgent.removeCapability(testCap);
-            callbackWithCap.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+            callbackWithCap.expect(CallbackEntry.LOST, mCellNetworkAgent);
             callbackWithoutCap.assertNoCallback();
             verify(mMockNetd).networkClearDefault();
 
@@ -10078,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());
@@ -10101,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);
 
@@ -10110,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());
@@ -10118,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));
@@ -10127,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();
@@ -10176,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);
@@ -10214,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.
@@ -10244,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)));
@@ -10283,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(),
@@ -10457,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);
@@ -10541,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)));
@@ -10565,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.
@@ -10615,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);
     }
@@ -12239,7 +12294,7 @@
         allNetworksCb.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         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
@@ -12251,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();
 
@@ -12268,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.
@@ -12283,8 +12338,7 @@
         // torn down.
         mCellNetworkAgent.setScore(new NetworkScore.Builder().setLegacyInt(30).build());
         allNetworksCb.expectLosing(mCellNetworkAgent, TEST_NASCENT_DELAY_MS * 2);
-        allNetworksCb.expectCallback(CallbackEntry.LOST, mCellNetworkAgent,
-                TEST_LINGER_DELAY_MS * 2);
+        allNetworksCb.expect(CallbackEntry.LOST, mCellNetworkAgent, TEST_LINGER_DELAY_MS * 2);
         mDefaultNetworkCallback.assertNoCallback();
 
         mCm.unregisterNetworkCallback(allNetworksCb);
@@ -13255,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.
@@ -13301,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.
@@ -13363,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.
@@ -14483,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) {
@@ -14628,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);
             }
@@ -14698,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();
@@ -14723,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) {
@@ -14740,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
@@ -14790,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);
         }
@@ -15400,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();
@@ -15874,7 +15928,7 @@
         }
 
         mEthernetNetworkAgent.disconnect();
-        cb.expectCallback(CallbackEntry.LOST, mEthernetNetworkAgent);
+        cb.expect(CallbackEntry.LOST, mEthernetNetworkAgent);
         mCm.unregisterNetworkCallback(cb);
     }
 
@@ -15944,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
@@ -16254,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);
@@ -16282,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));
@@ -16531,7 +16585,7 @@
         // Disconnect wifi
         mWiFiNetworkAgent.disconnect();
         assertNoCallbacks(mProfileDefaultNetworkCallback, mTestPackageDefaultNetworkCallback);
-        mDefaultNetworkCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        mDefaultNetworkCallback.expect(CallbackEntry.LOST, mWiFiNetworkAgent);
         mDefaultNetworkCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
     }
 
@@ -16728,23 +16782,25 @@
         assertThrows(NullPointerException.class, () -> mService.unofferNetwork(null));
     }
 
-    public void doTestIgnoreValidationAfterRoam(final boolean enabled) throws Exception {
-        assumeFalse(SdkLevel.isAtLeastT());
-        doReturn(enabled ? 5000 : -1).when(mResources)
+    public void doTestIgnoreValidationAfterRoam(int resValue, final boolean enabled)
+            throws Exception {
+        doReturn(resValue).when(mResources)
                 .getInteger(R.integer.config_validationFailureAfterRoamIgnoreTimeMillis);
 
+        final String bssid1 = "AA:AA:AA:AA:AA:AA";
+        final String bssid2 = "BB:BB:BB:BB:BB:BB";
         mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
         NetworkCapabilities wifiNc1 = new NetworkCapabilities()
                 .addCapability(NET_CAPABILITY_INTERNET)
+                .addCapability(NET_CAPABILITY_NOT_VPN)
+                .addCapability(NET_CAPABILITY_NOT_RESTRICTED)
+                .addCapability(NET_CAPABILITY_TRUSTED)
                 .addCapability(NET_CAPABILITY_NOT_VCN_MANAGED)
                 .addTransportType(TRANSPORT_WIFI)
-                .setTransportInfo(new WifiInfo.Builder().setBssid("AA:AA:AA:AA:AA:AA").build());
-        NetworkCapabilities wifiNc2 = new NetworkCapabilities()
-                .addCapability(NET_CAPABILITY_INTERNET)
-                .addCapability(NET_CAPABILITY_NOT_VCN_MANAGED)
-                .addTransportType(TRANSPORT_WIFI)
-                .setTransportInfo(new WifiInfo.Builder().setBssid("BB:BB:BB:BB:BB:BB").build());
+                .setTransportInfo(new WifiInfo.Builder().setBssid(bssid1).build());
+        NetworkCapabilities wifiNc2 = new NetworkCapabilities(wifiNc1)
+                .setTransportInfo(new WifiInfo.Builder().setBssid(bssid2).build());
         final LinkProperties wifiLp = new LinkProperties();
         wifiLp.setInterfaceName(WIFI_IFNAME);
         mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI, wifiLp, wifiNc1);
@@ -16767,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);
@@ -16789,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);
@@ -16808,50 +16864,88 @@
             wifiNetworkCallback.expectCapabilitiesWithout(NET_CAPABILITY_PARTIAL_CONNECTIVITY,
                     mWiFiNetworkAgent);
         }
+        mCm.unregisterNetworkCallback(wifiNetworkCallback);
 
         // Wi-Fi roams from wifiNc1 to wifiNc2, and now becomes really invalid. If validation
         // 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);
 
-        if (!enabled) {
+        if (enabled) {
+            // Network validation failed, but the result will be ignored.
+            assertTrue(mCm.getNetworkCapabilities(mWiFiNetworkAgent.getNetwork()).hasCapability(
+                    NET_CAPABILITY_VALIDATED));
+            mWiFiNetworkAgent.setNetworkValid(false);
+
+            // Behavior of after config_validationFailureAfterRoamIgnoreTimeMillis
+            ConditionVariable waitForValidationBlock = new ConditionVariable();
+            doReturn(50).when(mResources)
+                    .getInteger(R.integer.config_validationFailureAfterRoamIgnoreTimeMillis);
+            // Wi-Fi roaming from wifiNc2 to wifiNc1.
+            mWiFiNetworkAgent.setNetworkCapabilities(wifiNc1, true);
+            mWiFiNetworkAgent.setNetworkInvalid(false);
+            waitForValidationBlock.block(150);
+            mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), false);
             mDefaultNetworkCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
-            return;
+        } else {
+            mDefaultNetworkCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         }
 
-        // Network validation failed, but the result will be ignored.
-        assertTrue(mCm.getNetworkCapabilities(mWiFiNetworkAgent.getNetwork()).hasCapability(
-                NET_CAPABILITY_VALIDATED));
-        mWiFiNetworkAgent.setNetworkValid(false);
+        // Wi-Fi is still connected and would become the default network if cell were to
+        // disconnect. This assertion ensures that the switch to cellular was not caused by
+        // Wi-Fi disconnecting (e.g., because the capability change to wifiNc2 caused it
+        // to stop satisfying the default request).
+        mCellNetworkAgent.disconnect();
+        mDefaultNetworkCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        mDefaultNetworkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
 
-        // Behavior of after config_validationFailureAfterRoamIgnoreTimeMillis
-        ConditionVariable waitForValidationBlock = new ConditionVariable();
-        doReturn(50).when(mResources)
-                .getInteger(R.integer.config_validationFailureAfterRoamIgnoreTimeMillis);
-        // Wi-Fi roaming from wifiNc2 to wifiNc1.
-        mWiFiNetworkAgent.setNetworkCapabilities(wifiNc1, true);
-        mWiFiNetworkAgent.setNetworkInvalid(false);
-        waitForValidationBlock.block(150);
-        mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), false);
-        mDefaultNetworkCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
-
-        mCm.unregisterNetworkCallback(wifiNetworkCallback);
     }
 
     @Test
     public void testIgnoreValidationAfterRoamDisabled() throws Exception {
-        doTestIgnoreValidationAfterRoam(false /* enabled */);
+        doTestIgnoreValidationAfterRoam(-1, false /* enabled */);
     }
+
     @Test
     public void testIgnoreValidationAfterRoamEnabled() throws Exception {
-        doTestIgnoreValidationAfterRoam(true /* enabled */);
+        final boolean enabled = !SdkLevel.isAtLeastT();
+        doTestIgnoreValidationAfterRoam(5_000, enabled);
     }
 
     @Test
+    public void testShouldIgnoreValidationFailureAfterRoam() {
+        // Always disabled on T+.
+        assumeFalse(SdkLevel.isAtLeastT());
+
+        NetworkAgentInfo nai = fakeWifiNai(new NetworkCapabilities());
+
+        // Enabled, but never roamed.
+        doReturn(5_000).when(mResources)
+                .getInteger(R.integer.config_validationFailureAfterRoamIgnoreTimeMillis);
+        assertEquals(0, nai.lastRoamTime);
+        assertFalse(mService.shouldIgnoreValidationFailureAfterRoam(nai));
+
+        // Roamed recently.
+        nai.lastRoamTime = SystemClock.elapsedRealtime() - 500 /* ms */;
+        assertTrue(mService.shouldIgnoreValidationFailureAfterRoam(nai));
+
+        // Disabled due to invalid setting (maximum is 10 seconds).
+        doReturn(15_000).when(mResources)
+                .getInteger(R.integer.config_validationFailureAfterRoamIgnoreTimeMillis);
+        assertFalse(mService.shouldIgnoreValidationFailureAfterRoam(nai));
+
+        // Disabled.
+        doReturn(-1).when(mResources)
+                .getInteger(R.integer.config_validationFailureAfterRoamIgnoreTimeMillis);
+        assertFalse(mService.shouldIgnoreValidationFailureAfterRoam(nai));
+    }
+
+
+    @Test
     public void testLegacyTetheringApiGuardWithProperPermission() throws Exception {
         final String testIface = "test0";
         mServiceContext.setPermission(ACCESS_NETWORK_STATE, PERMISSION_DENIED);
diff --git a/tests/unit/java/com/android/server/connectivity/ClatCoordinatorTest.java b/tests/unit/java/com/android/server/connectivity/ClatCoordinatorTest.java
index 85bc4a9..49e3514 100644
--- a/tests/unit/java/com/android/server/connectivity/ClatCoordinatorTest.java
+++ b/tests/unit/java/com/android/server/connectivity/ClatCoordinatorTest.java
@@ -220,12 +220,12 @@
          */
         @Override
         public String generateIpv6Address(@NonNull String iface, @NonNull String v4,
-                @NonNull String prefix64) throws IOException {
+                @NonNull String prefix64, int mark) throws IOException {
             if (BASE_IFACE.equals(iface) && XLAT_LOCAL_IPV4ADDR_STRING.equals(v4)
                     && NAT64_PREFIX_STRING.equals(prefix64)) {
                 return XLAT_LOCAL_IPV6ADDR_STRING;
             }
-            fail("unsupported args: " + iface + ", " + v4 + ", " + prefix64);
+            fail("unsupported args: " + iface + ", " + v4 + ", " + prefix64 + ", " + mark);
             return null;
         }
 
@@ -417,7 +417,7 @@
 
         // Generate a checksum-neutral IID.
         inOrder.verify(mDeps).generateIpv6Address(eq(BASE_IFACE),
-                eq(XLAT_LOCAL_IPV4ADDR_STRING), eq(NAT64_PREFIX_STRING));
+                eq(XLAT_LOCAL_IPV4ADDR_STRING), eq(NAT64_PREFIX_STRING), eq(MARK));
 
         // Open, configure and bring up the tun interface.
         inOrder.verify(mDeps).createTunInterface(eq(STACKED_IFACE));
@@ -617,7 +617,7 @@
         class FailureDependencies extends TestDependencies {
             @Override
             public String generateIpv6Address(@NonNull String iface, @NonNull String v4,
-                    @NonNull String prefix64) throws IOException {
+                    @NonNull String prefix64, int mark) throws IOException {
                 throw new IOException();
             }
         }
diff --git a/tools/gn2bp/compat.py b/tools/gn2bp/compat.py
new file mode 100755
index 0000000..db587a2
--- /dev/null
+++ b/tools/gn2bp/compat.py
@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+# Copyright (C) 2022 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+try:
+  from shlex import quote
+except ImportError:
+  from pipes import quote
+
+try:
+  from urllib.request import urlretrieve
+except ImportError:
+  from urllib import urlretrieve
+
+try:
+  xrange = xrange
+except NameError:
+  xrange = range
+
+try:
+  basestring = basestring
+except NameError:
+  basestring = str
+
+
+def itervalues(o):
+  try:
+    return o.itervalues()
+  except AttributeError:
+    return o.values()
+
+
+def iteritems(o):
+  try:
+    return o.iteritems()
+  except AttributeError:
+    return o.items()
diff --git a/tools/gn2bp/gen_android_bp b/tools/gn2bp/gen_android_bp
new file mode 100755
index 0000000..4cacf7c
--- /dev/null
+++ b/tools/gn2bp/gen_android_bp
@@ -0,0 +1,1129 @@
+#!/usr/bin/env python3
+# Copyright (C) 2022 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This tool translates a collection of BUILD.gn files into a mostly equivalent
+# Android.bp file for the Android Soong build system. The input to the tool is a
+# JSON description of the GN build definition generated with the following
+# command:
+#
+#   gn desc out --format=json --all-toolchains "//*" > desc.json
+#
+# The tool is then given a list of GN labels for which to generate Android.bp
+# build rules. The dependencies for the GN labels are squashed to the generated
+# Android.bp target, except for actions which get their own genrule. Some
+# libraries are also mapped to their Android equivalents -- see |builtin_deps|.
+
+import argparse
+import collections
+import json
+import os
+import re
+import sys
+
+import gn_utils
+
+from compat import itervalues
+
+ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+# Arguments for the GN output directory.
+gn_args = ' '.join([
+    'is_debug=false',
+    'is_perfetto_build_generator=true',
+    'perfetto_build_with_android=true',
+    'target_cpu="arm"',
+    'target_os="android"',
+])
+
+# Default targets to translate to the blueprint file.
+default_targets = [
+    '//:libperfetto_client_experimental',
+    '//:libperfetto',
+    '//:perfetto_integrationtests',
+    '//:perfetto_unittests',
+    '//protos/perfetto/trace:perfetto_trace_protos',
+    '//src/android_internal:libperfetto_android_internal',
+    '//src/perfetto_cmd:perfetto',
+    '//src/perfetto_cmd:trigger_perfetto',
+    '//src/profiling/memory:heapprofd_client',
+    '//src/profiling/memory:heapprofd_client_api',
+    '//src/profiling/memory:heapprofd_api_noop',
+    '//src/profiling/memory:heapprofd',
+    '//src/profiling/memory:heapprofd_standalone_client',
+    '//src/profiling/perf:traced_perf',
+    '//src/traced/probes:traced_probes',
+    '//src/traced/service:traced',
+    '//src/trace_processor:trace_processor_shell',
+    '//test/cts:perfetto_cts_deps',
+    '//test/cts:perfetto_cts_jni_deps',
+    '//test:perfetto_gtest_logcat_printer',
+    '//test/vts:perfetto_vts_deps',
+]
+
+# Host targets
+ipc_plugin = '//src/ipc/protoc_plugin:ipc_plugin(%s)' % gn_utils.HOST_TOOLCHAIN
+protozero_plugin = '//src/protozero/protoc_plugin:protozero_plugin(%s)' % (
+    gn_utils.HOST_TOOLCHAIN)
+cppgen_plugin = '//src/protozero/protoc_plugin:cppgen_plugin(%s)' % (
+    gn_utils.HOST_TOOLCHAIN)
+
+default_targets += [
+    '//src/traceconv:traceconv(%s)' % gn_utils.HOST_TOOLCHAIN,
+    protozero_plugin,
+    ipc_plugin,
+]
+
+# Defines a custom init_rc argument to be applied to the corresponding output
+# blueprint target.
+target_initrc = {
+    '//src/traced/service:traced': {'perfetto.rc'},
+    '//src/profiling/memory:heapprofd': {'heapprofd.rc'},
+    '//src/profiling/perf:traced_perf': {'traced_perf.rc'},
+}
+
+target_host_supported = [
+    '//:libperfetto',
+    '//:libperfetto_client_experimental',
+    '//protos/perfetto/trace:perfetto_trace_protos',
+    '//src/trace_processor:demangle',
+    '//src/trace_processor:trace_processor_shell',
+]
+
+target_vendor_available = [
+    '//:libperfetto_client_experimental',
+]
+
+# Proto target groups which will be made public.
+proto_groups = {
+    'trace': [
+        '//protos/perfetto/trace:non_minimal_source_set',
+        '//protos/perfetto/trace:minimal_source_set'
+    ],
+}
+
+# All module names are prefixed with this string to avoid collisions.
+module_prefix = 'perfetto_'
+
+# Shared libraries which are directly translated to Android system equivalents.
+shared_library_allowlist = [
+    'android',
+    'android.hardware.atrace@1.0',
+    'android.hardware.health@2.0',
+    'android.hardware.health-V1-ndk',
+    'android.hardware.power.stats@1.0',
+    "android.hardware.power.stats-V1-cpp",
+    'base',
+    'binder',
+    'binder_ndk',
+    'cutils',
+    'hidlbase',
+    'hidltransport',
+    'hwbinder',
+    'incident',
+    'log',
+    'services',
+    'statssocket',
+    "tracingproxy",
+    'utils',
+]
+
+# Static libraries which are directly translated to Android system equivalents.
+static_library_allowlist = [
+    'statslog_perfetto',
+]
+
+# Name of the module which settings such as compiler flags for all other
+# modules.
+defaults_module = module_prefix + 'defaults'
+
+# Location of the project in the Android source tree.
+tree_path = 'external/perfetto'
+
+# Path for the protobuf sources in the standalone build.
+buildtools_protobuf_src = '//buildtools/protobuf/src'
+
+# Location of the protobuf src dir in the Android source tree.
+android_protobuf_src = 'external/protobuf/src'
+
+# Compiler flags which are passed through to the blueprint.
+cflag_allowlist = r'^-DPERFETTO.*$'
+
+# Compiler defines which are passed through to the blueprint.
+define_allowlist = r'^(GOOGLE_PROTO.*)|(ZLIB_.*)|(USE_MMAP)|(HAVE_HIDDEN)$'
+
+# The directory where the generated perfetto_build_flags.h will be copied into.
+buildflags_dir = 'include/perfetto/base/build_configs/android_tree'
+
+
+def enumerate_data_deps():
+  with open(os.path.join(ROOT_DIR, 'tools', 'test_data.txt')) as f:
+    lines = f.readlines()
+  for line in (line.strip() for line in lines if not line.startswith('#')):
+    assert os.path.exists(line), 'file %s should exist' % line
+    if line.startswith('test/data/'):
+      # Skip test data files that require GCS. They are only for benchmarks.
+      # We don't run benchmarks in the android tree.
+      continue
+    if line.endswith('/.'):
+      yield line[:-1] + '**/*'
+    else:
+      yield line
+
+
+# Additional arguments to apply to Android.bp rules.
+additional_args = {
+    'heapprofd_client_api': [
+        ('static_libs', {'libasync_safe'}),
+        # heapprofd_client_api MUST NOT have global constructors. Because it
+        # is loaded in an __attribute__((constructor)) of libc, we cannot
+        # guarantee that the global constructors get run before it is used.
+        ('cflags', {'-Wglobal-constructors', '-Werror=global-constructors'}),
+        ('version_script', 'src/profiling/memory/heapprofd_client_api.map.txt'),
+        ('stubs', {
+            'versions': ['S'],
+            'symbol_file': 'src/profiling/memory/heapprofd_client_api.map.txt',
+        }),
+        ('export_include_dirs', {'src/profiling/memory/include'}),
+    ],
+    'heapprofd_api_noop': [
+        ('version_script', 'src/profiling/memory/heapprofd_client_api.map.txt'),
+        ('stubs', {
+            'versions': ['S'],
+            'symbol_file': 'src/profiling/memory/heapprofd_client_api.map.txt',
+        }),
+        ('export_include_dirs', {'src/profiling/memory/include'}),
+    ],
+    'heapprofd_client': [
+        ('include_dirs', {'bionic/libc'}),
+        ('static_libs', {'libasync_safe'}),
+    ],
+    'heapprofd_standalone_client': [
+        ('static_libs', {'libasync_safe'}),
+        ('version_script', 'src/profiling/memory/heapprofd_client_api.map.txt'),
+        ('export_include_dirs', {'src/profiling/memory/include'}),
+        ('stl', 'libc++_static'),
+    ],
+#    'perfetto_unittests': [
+#        ('data', set(enumerate_data_deps())),
+#        ('include_dirs', {'bionic/libc/kernel'}),
+#    ],
+    'perfetto_integrationtests': [
+        ('test_suites', {'general-tests'}),
+        ('test_config', 'PerfettoIntegrationTests.xml'),
+    ],
+    'traced_probes': [('required', {
+        'libperfetto_android_internal', 'trigger_perfetto', 'traced_perf',
+        'mm_events'
+    }),],
+    'libperfetto_android_internal': [('static_libs', {'libhealthhalutils'}),],
+    'trace_processor_shell': [
+        ('strip', {
+            'all': True
+        }),
+        ('host', {
+            'stl': 'libc++_static',
+            'dist': {
+                'targets': ['sdk_repo']
+            },
+        }),
+    ],
+    'libperfetto_client_experimental': [
+        ('apex_available', {
+            '//apex_available:platform', 'com.android.art',
+            'com.android.art.debug'
+        }),
+        ('min_sdk_version', 'S'),
+        ('shared_libs', {'liblog'}),
+        ('export_include_dirs', {'include', buildflags_dir}),
+    ],
+    'perfetto_trace_protos': [
+        ('apex_available', {
+            '//apex_available:platform', 'com.android.art',
+            'com.android.art.debug'
+        }),
+        ('min_sdk_version', 'S'),
+    ],
+    'libperfetto': [('export_include_dirs', {'include', buildflags_dir}),],
+}
+
+
+def enable_gtest_and_gmock(module):
+  module.static_libs.add('libgmock')
+  module.static_libs.add('libgtest')
+  if module.name != 'perfetto_gtest_logcat_printer':
+    module.whole_static_libs.add('perfetto_gtest_logcat_printer')
+
+
+def enable_protobuf_full(module):
+  if module.type == 'cc_binary_host':
+    module.static_libs.add('libprotobuf-cpp-full')
+  elif module.host_supported:
+    module.host.static_libs.add('libprotobuf-cpp-full')
+    module.android.shared_libs.add('libprotobuf-cpp-full')
+  else:
+    module.shared_libs.add('libprotobuf-cpp-full')
+
+
+def enable_protobuf_lite(module):
+  module.shared_libs.add('libprotobuf-cpp-lite')
+
+
+def enable_protoc_lib(module):
+  if module.type == 'cc_binary_host':
+    module.static_libs.add('libprotoc')
+  else:
+    module.shared_libs.add('libprotoc')
+
+
+def enable_libunwindstack(module):
+  if module.name != 'heapprofd_standalone_client':
+    module.shared_libs.add('libunwindstack')
+    module.shared_libs.add('libprocinfo')
+    module.shared_libs.add('libbase')
+  else:
+    module.static_libs.add('libunwindstack')
+    module.static_libs.add('libprocinfo')
+    module.static_libs.add('libbase')
+    module.static_libs.add('liblzma')
+    module.static_libs.add('libdexfile_support')
+    module.runtime_libs.add('libdexfile')  # libdexfile_support dependency
+
+
+def enable_libunwind(module):
+  # libunwind is disabled on Darwin so we cannot depend on it.
+  pass
+
+
+def enable_sqlite(module):
+  if module.type == 'cc_binary_host':
+    module.static_libs.add('libsqlite')
+    module.static_libs.add('sqlite_ext_percentile')
+  elif module.host_supported:
+    # Copy what the sqlite3 command line tool does.
+    module.android.shared_libs.add('libsqlite')
+    module.android.shared_libs.add('libicu')
+    module.android.shared_libs.add('liblog')
+    module.android.shared_libs.add('libutils')
+    module.android.static_libs.add('sqlite_ext_percentile')
+    module.host.static_libs.add('libsqlite')
+    module.host.static_libs.add('sqlite_ext_percentile')
+  else:
+    module.shared_libs.add('libsqlite')
+    module.shared_libs.add('libicu')
+    module.shared_libs.add('liblog')
+    module.shared_libs.add('libutils')
+    module.static_libs.add('sqlite_ext_percentile')
+
+
+def enable_zlib(module):
+  if module.type == 'cc_binary_host':
+    module.static_libs.add('libz')
+  elif module.host_supported:
+    module.android.shared_libs.add('libz')
+    module.host.static_libs.add('libz')
+  else:
+    module.shared_libs.add('libz')
+
+
+def enable_uapi_headers(module):
+  module.include_dirs.add('bionic/libc/kernel')
+
+
+def enable_bionic_libc_platform_headers_on_android(module):
+  module.header_libs.add('bionic_libc_platform_headers')
+
+
+# Android equivalents for third-party libraries that the upstream project
+# depends on.
+builtin_deps = {
+    '//gn:default_deps':
+        lambda x: None,
+    '//gn:gtest_main':
+        lambda x: None,
+    '//gn:protoc':
+        lambda x: None,
+    '//gn:gtest_and_gmock':
+        enable_gtest_and_gmock,
+    '//gn:libunwind':
+        enable_libunwind,
+    '//gn:protobuf_full':
+        enable_protobuf_full,
+    '//gn:protobuf_lite':
+        enable_protobuf_lite,
+    '//gn:protoc_lib':
+        enable_protoc_lib,
+    '//gn:libunwindstack':
+        enable_libunwindstack,
+    '//gn:sqlite':
+        enable_sqlite,
+    '//gn:zlib':
+        enable_zlib,
+    '//gn:bionic_kernel_uapi_headers':
+        enable_uapi_headers,
+    '//src/profiling/memory:bionic_libc_platform_headers_on_android':
+        enable_bionic_libc_platform_headers_on_android,
+}
+
+# ----------------------------------------------------------------------------
+# End of configuration.
+# ----------------------------------------------------------------------------
+
+
+class Error(Exception):
+  pass
+
+
+class ThrowingArgumentParser(argparse.ArgumentParser):
+
+  def __init__(self, context):
+    super(ThrowingArgumentParser, self).__init__()
+    self.context = context
+
+  def error(self, message):
+    raise Error('%s: %s' % (self.context, message))
+
+
+def write_blueprint_key_value(output, name, value, sort=True):
+  """Writes a Blueprint key-value pair to the output"""
+
+  if isinstance(value, bool):
+    if value:
+      output.append('    %s: true,' % name)
+    else:
+      output.append('    %s: false,' % name)
+    return
+  if not value:
+    return
+  if isinstance(value, set):
+    value = sorted(value)
+  if isinstance(value, list):
+    output.append('    %s: [' % name)
+    for item in sorted(value) if sort else value:
+      output.append('        "%s",' % item)
+    output.append('    ],')
+    return
+  if isinstance(value, Target):
+    value.to_string(output)
+    return
+  if isinstance(value, dict):
+    kv_output = []
+    for k, v in value.items():
+      write_blueprint_key_value(kv_output, k, v)
+
+    output.append('    %s: {' % name)
+    for line in kv_output:
+      output.append('    %s' % line)
+    output.append('    },')
+    return
+  output.append('    %s: "%s",' % (name, value))
+
+
+class Target(object):
+  """A target-scoped part of a module"""
+
+  def __init__(self, name):
+    self.name = name
+    self.shared_libs = set()
+    self.static_libs = set()
+    self.whole_static_libs = set()
+    self.cflags = set()
+    self.dist = dict()
+    self.strip = dict()
+    self.stl = None
+
+  def to_string(self, output):
+    nested_out = []
+    self._output_field(nested_out, 'shared_libs')
+    self._output_field(nested_out, 'static_libs')
+    self._output_field(nested_out, 'whole_static_libs')
+    self._output_field(nested_out, 'cflags')
+    self._output_field(nested_out, 'stl')
+    self._output_field(nested_out, 'dist')
+    self._output_field(nested_out, 'strip')
+
+    if nested_out:
+      output.append('    %s: {' % self.name)
+      for line in nested_out:
+        output.append('    %s' % line)
+      output.append('    },')
+
+  def _output_field(self, output, name, sort=True):
+    value = getattr(self, name)
+    return write_blueprint_key_value(output, name, value, sort)
+
+
+class Module(object):
+  """A single module (e.g., cc_binary, cc_test) in a blueprint."""
+
+  def __init__(self, mod_type, name, gn_target):
+    self.type = mod_type
+    self.gn_target = gn_target
+    self.name = name
+    self.srcs = set()
+    self.comment = 'GN: ' + gn_utils.label_without_toolchain(gn_target)
+    self.shared_libs = set()
+    self.static_libs = set()
+    self.whole_static_libs = set()
+    self.runtime_libs = set()
+    self.tools = set()
+    self.cmd = None
+    self.host_supported = False
+    self.vendor_available = False
+    self.init_rc = set()
+    self.out = set()
+    self.export_include_dirs = set()
+    self.generated_headers = set()
+    self.export_generated_headers = set()
+    self.defaults = set()
+    self.cflags = set()
+    self.include_dirs = set()
+    self.header_libs = set()
+    self.required = set()
+    self.user_debug_flag = False
+    self.tool_files = None
+    self.android = Target('android')
+    self.host = Target('host')
+    self.lto = None
+    self.stl = None
+    self.dist = dict()
+    self.strip = dict()
+    self.data = set()
+    self.apex_available = set()
+    self.min_sdk_version = None
+    self.proto = dict()
+    # The genrule_XXX below are properties that must to be propagated back
+    # on the module(s) that depend on the genrule.
+    self.genrule_headers = set()
+    self.genrule_srcs = set()
+    self.genrule_shared_libs = set()
+    self.version_script = None
+    self.test_suites = set()
+    self.test_config = None
+    self.stubs = {}
+
+  def to_string(self, output):
+    if self.comment:
+      output.append('// %s' % self.comment)
+    output.append('%s {' % self.type)
+    self._output_field(output, 'name')
+    self._output_field(output, 'srcs')
+    self._output_field(output, 'shared_libs')
+    self._output_field(output, 'static_libs')
+    self._output_field(output, 'whole_static_libs')
+    self._output_field(output, 'runtime_libs')
+    self._output_field(output, 'tools')
+    self._output_field(output, 'cmd', sort=False)
+    if self.host_supported:
+      self._output_field(output, 'host_supported')
+    if self.vendor_available:
+      self._output_field(output, 'vendor_available')
+    self._output_field(output, 'init_rc')
+    self._output_field(output, 'out')
+    self._output_field(output, 'export_include_dirs')
+    self._output_field(output, 'generated_headers')
+    self._output_field(output, 'export_generated_headers')
+    self._output_field(output, 'defaults')
+    self._output_field(output, 'cflags')
+    self._output_field(output, 'include_dirs')
+    self._output_field(output, 'header_libs')
+    self._output_field(output, 'required')
+    self._output_field(output, 'dist')
+    self._output_field(output, 'strip')
+    self._output_field(output, 'tool_files')
+    self._output_field(output, 'data')
+    self._output_field(output, 'stl')
+    self._output_field(output, 'apex_available')
+    self._output_field(output, 'min_sdk_version')
+    self._output_field(output, 'version_script')
+    self._output_field(output, 'test_suites')
+    self._output_field(output, 'test_config')
+    self._output_field(output, 'stubs')
+    self._output_field(output, 'proto')
+
+    target_out = []
+    self._output_field(target_out, 'android')
+    self._output_field(target_out, 'host')
+    if target_out:
+      output.append('    target: {')
+      for line in target_out:
+        output.append('    %s' % line)
+      output.append('    },')
+
+    if self.user_debug_flag:
+      output.append('    product_variables: {')
+      output.append('        debuggable: {')
+      output.append(
+          '            cflags: ["-DPERFETTO_BUILD_WITH_ANDROID_USERDEBUG"],')
+      output.append('        },')
+      output.append('    },')
+    if self.lto is not None:
+      output.append('    target: {')
+      output.append('        android: {')
+      output.append('            lto: {')
+      output.append('                thin: %s,' %
+                    'true' if self.lto else 'false')
+      output.append('            },')
+      output.append('        },')
+      output.append('    },')
+    output.append('}')
+    output.append('')
+
+  def add_android_static_lib(self, lib):
+    if self.type == 'cc_binary_host':
+      raise Exception('Adding Android static lib for host tool is unsupported')
+    elif self.host_supported:
+      self.android.static_libs.add(lib)
+    else:
+      self.static_libs.add(lib)
+
+  def add_android_shared_lib(self, lib):
+    if self.type == 'cc_binary_host':
+      raise Exception('Adding Android shared lib for host tool is unsupported')
+    elif self.host_supported:
+      self.android.shared_libs.add(lib)
+    else:
+      self.shared_libs.add(lib)
+
+  def _output_field(self, output, name, sort=True):
+    value = getattr(self, name)
+    return write_blueprint_key_value(output, name, value, sort)
+
+
+class Blueprint(object):
+  """In-memory representation of an Android.bp file."""
+
+  def __init__(self):
+    self.modules = {}
+
+  def add_module(self, module):
+    """Adds a new module to the blueprint, replacing any existing module
+        with the same name.
+
+        Args:
+            module: Module instance.
+        """
+    self.modules[module.name] = module
+
+  def to_string(self, output):
+    for m in sorted(itervalues(self.modules), key=lambda m: m.name):
+      m.to_string(output)
+
+
+def label_to_module_name(label):
+  """Turn a GN label (e.g., //:perfetto_tests) into a module name."""
+  # If the label is explicibly listed in the default target list, don't prefix
+  # its name and return just the target name. This is so tools like
+  # "traceconv" stay as such in the Android tree.
+  label_without_toolchain = gn_utils.label_without_toolchain(label)
+  if label in default_targets or label_without_toolchain in default_targets:
+    return label_without_toolchain.split(':')[-1]
+
+  module = re.sub(r'^//:?', '', label_without_toolchain)
+  module = re.sub(r'[^a-zA-Z0-9_]', '_', module)
+  if not module.startswith(module_prefix):
+    return module_prefix + module
+  return module
+
+
+def is_supported_source_file(name):
+  """Returns True if |name| can appear in a 'srcs' list."""
+  return os.path.splitext(name)[1] in ['.c', '.cc', '.proto']
+
+
+def create_proto_modules(blueprint, gn, target):
+  """Generate genrules for a proto GN target.
+
+    GN actions are used to dynamically generate files during the build. The
+    Soong equivalent is a genrule. This function turns a specific kind of
+    genrule which turns .proto files into source and header files into a pair
+    equivalent genrules.
+
+    Args:
+        blueprint: Blueprint instance which is being generated.
+        target: gn_utils.Target object.
+
+    Returns:
+        The source_genrule module.
+    """
+  assert (target.type == 'proto_library')
+
+  tools = {'aprotoc'}
+  cpp_out_dir = '$(genDir)/%s/' % tree_path
+  target_module_name = label_to_module_name(target.name)
+
+  # In GN builds the proto path is always relative to the output directory
+  # (out/tmp.xxx).
+  cmd = ['mkdir -p %s &&' % cpp_out_dir, '$(location aprotoc)']
+  cmd += ['--proto_path=%s' % tree_path]
+
+  if buildtools_protobuf_src in target.proto_paths:
+    cmd += ['--proto_path=%s' % android_protobuf_src]
+
+  # We don't generate any targets for source_set proto modules because
+  # they will be inlined into other modules if required.
+  if target.proto_plugin == 'source_set':
+    return None
+
+  # Descriptor targets only generate a single target.
+  if target.proto_plugin == 'descriptor':
+    out = '{}.bin'.format(target_module_name)
+
+    cmd += ['--descriptor_set_out=$(out)']
+    cmd += ['$(in)']
+
+    descriptor_module = Module('genrule', target_module_name, target.name)
+    descriptor_module.cmd = ' '.join(cmd)
+    descriptor_module.out = [out]
+    descriptor_module.tools = tools
+    blueprint.add_module(descriptor_module)
+
+    # Recursively extract the .proto files of all the dependencies and
+    # add them to srcs.
+    descriptor_module.srcs.update(
+        gn_utils.label_to_path(src) for src in target.sources)
+    for dep in target.transitive_proto_deps:
+      current_target = gn.get_target(dep)
+      descriptor_module.srcs.update(
+          gn_utils.label_to_path(src) for src in current_target.sources)
+
+    return descriptor_module
+
+  # We create two genrules for each proto target: one for the headers and
+  # another for the sources. This is because the module that depends on the
+  # generated files needs to declare two different types of dependencies --
+  # source files in 'srcs' and headers in 'generated_headers' -- and it's not
+  # valid to generate .h files from a source dependency and vice versa.
+  source_module_name = target_module_name + '_gen'
+  source_module = Module('genrule', source_module_name, target.name)
+  blueprint.add_module(source_module)
+  source_module.srcs.update(
+      gn_utils.label_to_path(src) for src in target.sources)
+
+  header_module = Module('genrule', source_module_name + '_headers',
+                         target.name)
+  blueprint.add_module(header_module)
+  header_module.srcs = set(source_module.srcs)
+
+  # TODO(primiano): at some point we should remove this. This was introduced
+  # by aosp/1108421 when adding "protos/" to .proto include paths, in order to
+  # avoid doing multi-repo changes and allow old clients in the android tree
+  # to still do the old #include "perfetto/..." rather than
+  # #include "protos/perfetto/...".
+  header_module.export_include_dirs = {'.', 'protos'}
+
+  source_module.genrule_srcs.add(':' + source_module.name)
+  source_module.genrule_headers.add(header_module.name)
+
+  if target.proto_plugin == 'proto':
+    suffixes = ['pb']
+    source_module.genrule_shared_libs.add('libprotobuf-cpp-lite')
+    cmd += ['--cpp_out=lite=true:' + cpp_out_dir]
+  elif target.proto_plugin == 'protozero':
+    suffixes = ['pbzero']
+    plugin = create_modules_from_target(blueprint, gn, protozero_plugin)
+    tools.add(plugin.name)
+    cmd += ['--plugin=protoc-gen-plugin=$(location %s)' % plugin.name]
+    cmd += ['--plugin_out=wrapper_namespace=pbzero:' + cpp_out_dir]
+  elif target.proto_plugin == 'cppgen':
+    suffixes = ['gen']
+    plugin = create_modules_from_target(blueprint, gn, cppgen_plugin)
+    tools.add(plugin.name)
+    cmd += ['--plugin=protoc-gen-plugin=$(location %s)' % plugin.name]
+    cmd += ['--plugin_out=wrapper_namespace=gen:' + cpp_out_dir]
+  elif target.proto_plugin == 'ipc':
+    suffixes = ['ipc']
+    plugin = create_modules_from_target(blueprint, gn, ipc_plugin)
+    tools.add(plugin.name)
+    cmd += ['--plugin=protoc-gen-plugin=$(location %s)' % plugin.name]
+    cmd += ['--plugin_out=wrapper_namespace=gen:' + cpp_out_dir]
+  else:
+    raise Error('Unsupported proto plugin: %s' % target.proto_plugin)
+
+  cmd += ['$(in)']
+  source_module.cmd = ' '.join(cmd)
+  header_module.cmd = source_module.cmd
+  source_module.tools = tools
+  header_module.tools = tools
+
+  for sfx in suffixes:
+    source_module.out.update('%s/%s' %
+                             (tree_path, src.replace('.proto', '.%s.cc' % sfx))
+                             for src in source_module.srcs)
+    header_module.out.update('%s/%s' %
+                             (tree_path, src.replace('.proto', '.%s.h' % sfx))
+                             for src in header_module.srcs)
+  return source_module
+
+
+def create_amalgamated_sql_metrics_module(blueprint, target):
+  bp_module_name = label_to_module_name(target.name)
+  module = Module('genrule', bp_module_name, target.name)
+  module.tool_files = [
+      'tools/gen_amalgamated_sql_metrics.py',
+  ]
+  module.cmd = ' '.join([
+      '$(location tools/gen_amalgamated_sql_metrics.py)',
+      '--cpp_out=$(out)',
+      '$(in)',
+  ])
+  module.genrule_headers.add(module.name)
+  module.out.update(target.outputs)
+  module.srcs.update(gn_utils.label_to_path(src) for src in target.inputs)
+  blueprint.add_module(module)
+  return module
+
+
+def create_cc_proto_descriptor_module(blueprint, target):
+  bp_module_name = label_to_module_name(target.name)
+  module = Module('genrule', bp_module_name, target.name)
+  module.tool_files = [
+      'tools/gen_cc_proto_descriptor.py',
+  ]
+  module.cmd = ' '.join([
+      '$(location tools/gen_cc_proto_descriptor.py)', '--gen_dir=$(genDir)',
+      '--cpp_out=$(out)', '$(in)'
+  ])
+  module.genrule_headers.add(module.name)
+  module.srcs.update(
+      ':' + label_to_module_name(dep) for dep in target.proto_deps)
+  module.srcs.update(
+      gn_utils.label_to_path(src)
+      for src in target.inputs
+      if "tmp.gn_utils" not in src)
+  module.out.update(target.outputs)
+  blueprint.add_module(module)
+  return module
+
+
+def create_gen_version_module(blueprint, target, bp_module_name):
+  module = Module('genrule', bp_module_name, gn_utils.GEN_VERSION_TARGET)
+  script_path = gn_utils.label_to_path(target.script)
+  module.genrule_headers.add(bp_module_name)
+  module.tool_files = [script_path]
+  module.out.update(target.outputs)
+  module.srcs.update(gn_utils.label_to_path(src) for src in target.inputs)
+  module.cmd = ' '.join([
+      'python3 $(location %s)' % script_path, '--no_git',
+      '--changelog=$(location CHANGELOG)', '--cpp_out=$(out)'
+  ])
+  blueprint.add_module(module)
+  return module
+
+
+def create_proto_group_modules(blueprint, gn, module_name, target_names):
+  # TODO(lalitm): today, we're only adding a Java lite module because that's
+  # the only one used in practice. In the future, if we need other target types
+  # (e.g. C++, Java full etc.) add them here.
+  bp_module_name = label_to_module_name(module_name) + '_java_protos'
+  module = Module('java_library', bp_module_name, bp_module_name)
+  module.comment = f'''GN: [{', '.join(target_names)}]'''
+  module.proto = {'type': 'lite', 'canonical_path_from_root': False}
+
+  for name in target_names:
+    target = gn.get_target(name)
+    module.srcs.update(gn_utils.label_to_path(src) for src in target.sources)
+    for dep_label in target.transitive_proto_deps:
+      dep = gn.get_target(dep_label)
+      module.srcs.update(gn_utils.label_to_path(src) for src in dep.sources)
+
+  blueprint.add_module(module)
+
+
+def _get_cflags(target):
+  cflags = {flag for flag in target.cflags if re.match(cflag_allowlist, flag)}
+  cflags |= set("-D%s" % define
+                for define in target.defines
+                if re.match(define_allowlist, define))
+  return cflags
+
+
+def create_modules_from_target(blueprint, gn, gn_target_name):
+  """Generate module(s) for a given GN target.
+
+    Given a GN target name, generate one or more corresponding modules into a
+    blueprint. The only case when this generates >1 module is proto libraries.
+
+    Args:
+        blueprint: Blueprint instance which is being generated.
+        gn: gn_utils.GnParser object.
+        gn_target_name: GN target for module generation.
+    """
+  bp_module_name = label_to_module_name(gn_target_name)
+  if bp_module_name in blueprint.modules:
+    return blueprint.modules[bp_module_name]
+  target = gn.get_target(gn_target_name)
+
+  name_without_toolchain = gn_utils.label_without_toolchain(target.name)
+  if target.type == 'executable':
+    if target.toolchain == gn_utils.HOST_TOOLCHAIN:
+      module_type = 'cc_binary_host'
+    elif target.testonly:
+      module_type = 'cc_test'
+    else:
+      module_type = 'cc_binary'
+    module = Module(module_type, bp_module_name, gn_target_name)
+  elif target.type == 'static_library':
+    module = Module('cc_library_static', bp_module_name, gn_target_name)
+  elif target.type == 'shared_library':
+    module = Module('cc_library_shared', bp_module_name, gn_target_name)
+  elif target.type == 'source_set':
+    module = Module('filegroup', bp_module_name, gn_target_name)
+  elif target.type == 'group':
+    # "group" targets are resolved recursively by gn_utils.get_target().
+    # There's nothing we need to do at this level for them.
+    return None
+  elif target.type == 'proto_library':
+    module = create_proto_modules(blueprint, gn, target)
+    if module is None:
+      return None
+  elif target.type == 'action':
+    if 'gen_amalgamated_sql_metrics' in target.name:
+      module = create_amalgamated_sql_metrics_module(blueprint, target)
+    elif re.match('.*gen_cc_.*_descriptor$', name_without_toolchain):
+      module = create_cc_proto_descriptor_module(blueprint, target)
+    elif target.type == 'action' and \
+        name_without_toolchain == gn_utils.GEN_VERSION_TARGET:
+      module = create_gen_version_module(blueprint, target, bp_module_name)
+    else:
+      raise Error('Unhandled action: {}'.format(target.name))
+  else:
+    raise Error('Unknown target %s (%s)' % (target.name, target.type))
+
+  blueprint.add_module(module)
+  module.host_supported = (name_without_toolchain in target_host_supported)
+  module.vendor_available = (name_without_toolchain in target_vendor_available)
+  module.init_rc = target_initrc.get(target.name, [])
+  module.srcs.update(
+      gn_utils.label_to_path(src)
+      for src in target.sources
+      if is_supported_source_file(src))
+
+  if target.type in gn_utils.LINKER_UNIT_TYPES:
+    module.cflags.update(_get_cflags(target))
+
+  module_is_compiled = module.type not in ('genrule', 'filegroup')
+  if module_is_compiled:
+    # Don't try to inject library/source dependencies into genrules or
+    # filegroups because they are not compiled in the traditional sense.
+    module.defaults = [defaults_module]
+    for lib in target.libs:
+      # Generally library names should be mangled as 'libXXX', unless they
+      # are HAL libraries (e.g., android.hardware.health@2.0) or AIDL c++ / NDK
+      # libraries (e.g. "android.hardware.power.stats-V1-cpp")
+      android_lib = lib if '@' in lib or "-cpp" in lib or "-ndk" in lib \
+        else 'lib' + lib
+      if lib in shared_library_allowlist:
+        module.add_android_shared_lib(android_lib)
+      if lib in static_library_allowlist:
+        module.add_android_static_lib(android_lib)
+
+  # If the module is a static library, export all the generated headers.
+  if module.type == 'cc_library_static':
+    module.export_generated_headers = module.generated_headers
+
+  # Merge in additional hardcoded arguments.
+  for key, add_val in additional_args.get(module.name, []):
+    curr = getattr(module, key)
+    if add_val and isinstance(add_val, set) and isinstance(curr, set):
+      curr.update(add_val)
+    elif isinstance(add_val, str) and (not curr or isinstance(curr, str)):
+      setattr(module, key, add_val)
+    elif isinstance(add_val, bool) and (not curr or isinstance(curr, bool)):
+      setattr(module, key, add_val)
+    elif isinstance(add_val, dict) and isinstance(curr, dict):
+      curr.update(add_val)
+    elif isinstance(add_val, dict) and isinstance(curr, Target):
+      curr.__dict__.update(add_val)
+    else:
+      raise Error('Unimplemented type %r of additional_args: %r' %
+                  (type(add_val), key))
+
+  # dep_name is an unmangled GN target name (e.g. //foo:bar(toolchain)).
+  all_deps = target.deps | target.source_set_deps | target.transitive_proto_deps
+  for dep_name in all_deps:
+    # If the dependency refers to a library which we can replace with an
+    # Android equivalent, stop recursing and patch the dependency in.
+    # Don't recurse into //buildtools, builtin_deps are intercepted at
+    # the //gn:xxx level.
+    if dep_name.startswith('//buildtools'):
+      continue
+
+    # Ignore the dependency on the gen_buildflags genrule. That is run
+    # separately in this generator and the generated file is copied over
+    # into the repo (see usage of |buildflags_dir| in this script).
+    if dep_name.startswith(gn_utils.BUILDFLAGS_TARGET):
+      continue
+
+    dep_module = create_modules_from_target(blueprint, gn, dep_name)
+
+    # For filegroups and genrule, recurse but don't apply the deps.
+    if not module_is_compiled:
+      continue
+
+    # |builtin_deps| override GN deps with Android-specific ones. See the
+    # config in the top of this file.
+    if gn_utils.label_without_toolchain(dep_name) in builtin_deps:
+      builtin_deps[gn_utils.label_without_toolchain(dep_name)](module)
+      continue
+
+    # Don't recurse in any other //gn dep if not handled by builtin_deps.
+    if dep_name.startswith('//gn:'):
+      continue
+
+    if dep_module is None:
+      continue
+    if dep_module.type == 'cc_library_shared':
+      module.shared_libs.add(dep_module.name)
+    elif dep_module.type == 'cc_library_static':
+      module.static_libs.add(dep_module.name)
+    elif dep_module.type == 'filegroup':
+      module.srcs.add(':' + dep_module.name)
+    elif dep_module.type == 'genrule':
+      module.generated_headers.update(dep_module.genrule_headers)
+      module.srcs.update(dep_module.genrule_srcs)
+      module.shared_libs.update(dep_module.genrule_shared_libs)
+    elif dep_module.type == 'cc_binary':
+      continue  # Ignore executables deps (used by cmdline integration tests).
+    else:
+      raise Error('Unknown dep %s (%s) for target %s' %
+                  (dep_module.name, dep_module.type, module.name))
+
+  return module
+
+
+def create_blueprint_for_targets(gn, desc, targets):
+  """Generate a blueprint for a list of GN targets."""
+  blueprint = Blueprint()
+
+  # Default settings used by all modules.
+  defaults = Module('cc_defaults', defaults_module, '//gn:default_deps')
+
+  # We have to use include_dirs passing the path relative to the android tree.
+  # This is because: (i) perfetto_cc_defaults is used also by
+  # test/**/Android.bp; (ii) if we use local_include_dirs instead, paths
+  # become relative to the Android.bp that *uses* cc_defaults (not the one
+  # that defines it).s
+  defaults.include_dirs = {
+      tree_path, tree_path + '/include', tree_path + '/' + buildflags_dir,
+      tree_path + '/src/profiling/memory/include'
+  }
+  defaults.cflags = [
+      '-Wno-error=return-type',
+      '-Wno-sign-compare',
+      '-Wno-sign-promo',
+      '-Wno-unused-parameter',
+      '-fvisibility=hidden',
+      '-O2',
+  ]
+  defaults.user_debug_flag = True
+  defaults.lto = True
+
+  blueprint.add_module(defaults)
+  for target in targets:
+    create_modules_from_target(blueprint, gn, target)
+  return blueprint
+
+
+def main():
+  parser = argparse.ArgumentParser(
+      description='Generate Android.bp from a GN description.')
+  parser.add_argument(
+      '--check-only',
+      help='Don\'t keep the generated files',
+      action='store_true')
+  parser.add_argument(
+      '--desc',
+      help='GN description (e.g., gn desc out --format=json --all-toolchains "//*"'
+  )
+  parser.add_argument(
+      '--extras',
+      help='Extra targets to include at the end of the Blueprint file',
+      default=os.path.join(gn_utils.repo_root(), 'Android.bp.extras'),
+  )
+  parser.add_argument(
+      '--output',
+      help='Blueprint file to create',
+      default=os.path.join(gn_utils.repo_root(), 'Android.bp'),
+  )
+  parser.add_argument(
+      'targets',
+      nargs=argparse.REMAINDER,
+      help='Targets to include in the blueprint (e.g., "//:perfetto_tests")')
+  args = parser.parse_args()
+
+  if args.desc:
+    with open(args.desc) as f:
+      desc = json.load(f)
+  else:
+    desc = gn_utils.create_build_description(gn_args)
+
+  gn = gn_utils.GnParser(desc)
+  blueprint = create_blueprint_for_targets(gn, desc, args.targets or
+                                           default_targets)
+  project_root = os.path.abspath(os.path.dirname(os.path.dirname(__file__)))
+  tool_name = os.path.relpath(os.path.abspath(__file__), project_root)
+
+  # TODO(primiano): enable this on Android after the TODO in
+  # perfetto_component.gni is fixed.
+  # Check for ODR violations
+  # for target_name in default_targets:
+  # checker = gn_utils.ODRChecker(gn, target_name)
+
+  # Add any proto groups to the blueprint.
+  for l_name, t_names in proto_groups.items():
+    create_proto_group_modules(blueprint, gn, l_name, t_names)
+
+  output = [
+      """// Copyright (C) 2017 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// This file is automatically generated by %s. Do not edit.
+""" % (tool_name)
+  ]
+  blueprint.to_string(output)
+  with open(args.extras, 'r') as r:
+    for line in r:
+      output.append(line.rstrip("\n\r"))
+
+  out_files = []
+
+  # Generate the Android.bp file.
+  out_files.append(args.output + '.swp')
+  with open(out_files[-1], 'w') as f:
+    f.write('\n'.join(output))
+    # Text files should have a trailing EOL.
+    f.write('\n')
+
+  # Generate the perfetto_build_flags.h file.
+  out_files.append(os.path.join(buildflags_dir, 'perfetto_build_flags.h.swp'))
+  gn_utils.gen_buildflags(gn_args, out_files[-1])
+
+  # Either check the contents or move the files to their final destination.
+  return gn_utils.check_or_commit_generated_files(out_files, args.check_only)
+
+
+if __name__ == '__main__':
+  sys.exit(main())
diff --git a/tools/gn2bp/gn_utils.py b/tools/gn2bp/gn_utils.py
new file mode 100644
index 0000000..91d7a32
--- /dev/null
+++ b/tools/gn2bp/gn_utils.py
@@ -0,0 +1,525 @@
+# Copyright (C) 2022 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# A collection of utilities for extracting build rule information from GN
+# projects.
+
+from __future__ import print_function
+import collections
+import errno
+import filecmp
+import json
+import os
+import re
+import shutil
+import subprocess
+import sys
+from compat import iteritems
+
+BUILDFLAGS_TARGET = '//gn:gen_buildflags'
+GEN_VERSION_TARGET = '//src/base:version_gen_h'
+TARGET_TOOLCHAIN = '//gn/standalone/toolchain:gcc_like_host'
+HOST_TOOLCHAIN = '//gn/standalone/toolchain:gcc_like_host'
+LINKER_UNIT_TYPES = ('executable', 'shared_library', 'static_library')
+
+# TODO(primiano): investigate these, they require further componentization.
+ODR_VIOLATION_IGNORE_TARGETS = {
+    '//test/cts:perfetto_cts_deps',
+    '//:perfetto_integrationtests',
+}
+
+
+def _check_command_output(cmd, cwd):
+  try:
+    output = subprocess.check_output(cmd, stderr=subprocess.STDOUT, cwd=cwd)
+  except subprocess.CalledProcessError as e:
+    print(
+        'Command "{}" failed in {}:'.format(' '.join(cmd), cwd),
+        file=sys.stderr)
+    print(e.output.decode(), file=sys.stderr)
+    sys.exit(1)
+  else:
+    return output.decode()
+
+
+def repo_root():
+  """Returns an absolute path to the repository root."""
+  return os.path.join(
+      os.path.realpath(os.path.dirname(__file__)), os.path.pardir)
+
+
+def _tool_path(name, system_buildtools=False):
+  # Pass-through to use name if the caller requests to use the system
+  # toolchain.
+  if system_buildtools:
+    return [name]
+  wrapper = os.path.abspath(
+      os.path.join(repo_root(), 'tools', 'run_buildtools_binary.py'))
+  return ['python3', wrapper, name]
+
+
+def prepare_out_directory(gn_args,
+                          name,
+                          root=repo_root(),
+                          system_buildtools=False):
+  """Creates the JSON build description by running GN.
+
+    Returns (path, desc) where |path| is the location of the output directory
+    and |desc| is the JSON build description.
+    """
+  out = os.path.join(root, 'out', name)
+  try:
+    os.makedirs(out)
+  except OSError as e:
+    if e.errno != errno.EEXIST:
+      raise
+  _check_command_output(
+      _tool_path('gn', system_buildtools) +
+      ['gen', out, '--args=%s' % gn_args],
+      cwd=repo_root())
+  return out
+
+
+def load_build_description(out, system_buildtools=False):
+  """Creates the JSON build description by running GN."""
+  desc = _check_command_output(
+      _tool_path('gn', system_buildtools) +
+      ['desc', out, '--format=json', '--all-toolchains', '//*'],
+      cwd=repo_root())
+  return json.loads(desc)
+
+
+def create_build_description(gn_args, root=repo_root()):
+  """Prepares a GN out directory and loads the build description from it.
+
+    The temporary out directory is automatically deleted.
+    """
+  out = prepare_out_directory(gn_args, 'tmp.gn_utils', root=root)
+  try:
+    return load_build_description(out)
+  finally:
+    shutil.rmtree(out)
+
+
+def build_targets(out, targets, quiet=False, system_buildtools=False):
+  """Runs ninja to build a list of GN targets in the given out directory.
+
+    Compiling these targets is required so that we can include any generated
+    source files in the amalgamated result.
+    """
+  targets = [t.replace('//', '') for t in targets]
+  with open(os.devnull, 'w') as devnull:
+    stdout = devnull if quiet else None
+    cmd = _tool_path('ninja', system_buildtools) + targets
+    subprocess.check_call(cmd, cwd=os.path.abspath(out), stdout=stdout)
+
+
+def compute_source_dependencies(out, system_buildtools=False):
+  """For each source file, computes a set of headers it depends on."""
+  ninja_deps = _check_command_output(
+      _tool_path('ninja', system_buildtools) + ['-t', 'deps'], cwd=out)
+  deps = {}
+  current_source = None
+  for line in ninja_deps.split('\n'):
+    filename = os.path.relpath(os.path.join(out, line.strip()), repo_root())
+    if not line or line[0] != ' ':
+      current_source = None
+      continue
+    elif not current_source:
+      # We're assuming the source file is always listed before the
+      # headers.
+      assert os.path.splitext(line)[1] in ['.c', '.cc', '.cpp', '.S']
+      current_source = filename
+      deps[current_source] = []
+    else:
+      assert current_source
+      deps[current_source].append(filename)
+  return deps
+
+
+def label_to_path(label):
+  """Turn a GN output label (e.g., //some_dir/file.cc) into a path."""
+  assert label.startswith('//')
+  return label[2:]
+
+
+def label_without_toolchain(label):
+  """Strips the toolchain from a GN label.
+
+    Return a GN label (e.g //buildtools:protobuf(//gn/standalone/toolchain:
+    gcc_like_host) without the parenthesised toolchain part.
+    """
+  return label.split('(')[0]
+
+
+def label_to_target_name_with_path(label):
+  """
+  Turn a GN label into a target name involving the full path.
+  e.g., //src/perfetto:tests -> src_perfetto_tests
+  """
+  name = re.sub(r'^//:?', '', label)
+  name = re.sub(r'[^a-zA-Z0-9_]', '_', name)
+  return name
+
+
+def gen_buildflags(gn_args, target_file):
+  """Generates the perfetto_build_flags.h for the given config.
+
+    target_file: the path, relative to the repo root, where the generated
+        buildflag header will be copied into.
+    """
+  tmp_out = prepare_out_directory(gn_args, 'tmp.gen_buildflags')
+  build_targets(tmp_out, [BUILDFLAGS_TARGET], quiet=True)
+  src = os.path.join(tmp_out, 'gen', 'build_config', 'perfetto_build_flags.h')
+  shutil.copy(src, os.path.join(repo_root(), target_file))
+  shutil.rmtree(tmp_out)
+
+
+def check_or_commit_generated_files(tmp_files, check):
+  """Checks that gen files are unchanged or renames them to the final location
+
+    Takes in input a list of 'xxx.swp' files that have been written.
+    If check == False, it renames xxx.swp -> xxx.
+    If check == True, it just checks that the contents of 'xxx.swp' == 'xxx'.
+    Returns 0 if no diff was detected, 1 otherwise (to be used as exit code).
+    """
+  res = 0
+  for tmp_file in tmp_files:
+    assert (tmp_file.endswith('.swp'))
+    target_file = os.path.relpath(tmp_file[:-4])
+    if check:
+      if not filecmp.cmp(tmp_file, target_file):
+        sys.stderr.write('%s needs to be regenerated\n' % target_file)
+        res = 1
+      os.unlink(tmp_file)
+    else:
+      os.rename(tmp_file, target_file)
+  return res
+
+
+class ODRChecker(object):
+  """Detects ODR violations in linker units
+
+  When we turn GN source sets into Soong & Bazel file groups, there is the risk
+  to create ODR violations by including the same file group into different
+  linker unit (this is because other build systems don't have a concept
+  equivalent to GN's source_set). This class navigates the transitive
+  dependencies (mostly static libraries) of a target and detects if multiple
+  paths end up including the same file group. This is to avoid situations like:
+
+  traced.exe -> base(file group)
+  traced.exe -> libperfetto(static lib) -> base(file group)
+  """
+
+  def __init__(self, gn, target_name):
+    self.gn = gn
+    self.root = gn.get_target(target_name)
+    self.source_sets = collections.defaultdict(set)
+    self.deps_visited = set()
+    self.source_set_hdr_only = {}
+
+    self._visit(target_name)
+    num_violations = 0
+    if target_name in ODR_VIOLATION_IGNORE_TARGETS:
+      return
+    for sset, paths in self.source_sets.items():
+      if self.is_header_only(sset):
+        continue
+      if len(paths) != 1:
+        num_violations += 1
+        print(
+            'ODR violation in target %s, multiple paths include %s:\n  %s' %
+            (target_name, sset, '\n  '.join(paths)),
+            file=sys.stderr)
+    if num_violations > 0:
+      raise Exception('%d ODR violations detected. Build generation aborted' %
+                      num_violations)
+
+  def _visit(self, target_name, parent_path=''):
+    target = self.gn.get_target(target_name)
+    path = ((parent_path + ' > ') if parent_path else '') + target_name
+    if not target:
+      raise Exception('Cannot find target %s' % target_name)
+    for ssdep in target.source_set_deps:
+      name_and_path = '%s (via %s)' % (target_name, path)
+      self.source_sets[ssdep].add(name_and_path)
+    deps = set(target.deps).union(
+        target.transitive_proto_deps) - self.deps_visited
+    for dep_name in deps:
+      dep = self.gn.get_target(dep_name)
+      if dep.type == 'executable':
+        continue  # Execs are strong boundaries and don't cause ODR violations.
+      # static_library dependencies should reset the path. It doesn't matter if
+      # we get to a source file via:
+      # source_set1 > static_lib > source.cc OR
+      # source_set1 > source_set2 > static_lib > source.cc
+      # This is NOT an ODR violation because source.cc is linked from the same
+      # static library
+      next_parent_path = path if dep.type != 'static_library' else ''
+      self.deps_visited.add(dep_name)
+      self._visit(dep_name, next_parent_path)
+
+  def is_header_only(self, source_set_name):
+    cached = self.source_set_hdr_only.get(source_set_name)
+    if cached is not None:
+      return cached
+    target = self.gn.get_target(source_set_name)
+    if target.type != 'source_set':
+      raise TypeError('%s is not a source_set' % source_set_name)
+    res = all(src.endswith('.h') for src in target.sources)
+    self.source_set_hdr_only[source_set_name] = res
+    return res
+
+
+class GnParser(object):
+  """A parser with some cleverness for GN json desc files
+
+    The main goals of this parser are:
+    1) Deal with the fact that other build systems don't have an equivalent
+       notion to GN's source_set. Conversely to Bazel's and Soong's filegroups,
+       GN source_sets expect that dependencies, cflags and other source_set
+       properties propagate up to the linker unit (static_library, executable or
+       shared_library). This parser simulates the same behavior: when a
+       source_set is encountered, some of its variables (cflags and such) are
+       copied up to the dependent targets. This is to allow gen_xxx to create
+       one filegroup for each source_set and then squash all the other flags
+       onto the linker unit.
+    2) Detect and special-case protobuf targets, figuring out the protoc-plugin
+       being used.
+    """
+
+  class Target(object):
+    """Reperesents A GN target.
+
+        Maked properties are propagated up the dependency chain when a
+        source_set dependency is encountered.
+        """
+
+    def __init__(self, name, type):
+      self.name = name  # e.g. //src/ipc:ipc
+
+      VALID_TYPES = ('static_library', 'shared_library', 'executable', 'group',
+                     'action', 'source_set', 'proto_library')
+      assert (type in VALID_TYPES)
+      self.type = type
+      self.testonly = False
+      self.toolchain = None
+
+      # These are valid only for type == proto_library.
+      # This is typically: 'proto', 'protozero', 'ipc'.
+      self.proto_plugin = None
+      self.proto_paths = set()
+      self.proto_exports = set()
+
+      self.sources = set()
+      # TODO(primiano): consider whether the public section should be part of
+      # bubbled-up sources.
+      self.public_headers = set()  # 'public'
+
+      # These are valid only for type == 'action'
+      self.inputs = set()
+      self.outputs = set()
+      self.script = None
+      self.args = []
+
+      # These variables are propagated up when encountering a dependency
+      # on a source_set target.
+      self.cflags = set()
+      self.defines = set()
+      self.deps = set()
+      self.libs = set()
+      self.include_dirs = set()
+      self.ldflags = set()
+      self.source_set_deps = set()  # Transitive set of source_set deps.
+      self.proto_deps = set()
+      self.transitive_proto_deps = set()
+
+      # Deps on //gn:xxx have this flag set to True. These dependencies
+      # are special because they pull third_party code from buildtools/.
+      # We don't want to keep recursing into //buildtools in generators,
+      # this flag is used to stop the recursion and create an empty
+      # placeholder target once we hit //gn:protoc or similar.
+      self.is_third_party_dep_ = False
+
+    def __lt__(self, other):
+      if isinstance(other, self.__class__):
+        return self.name < other.name
+      raise TypeError(
+          '\'<\' not supported between instances of \'%s\' and \'%s\'' %
+          (type(self).__name__, type(other).__name__))
+
+    def __repr__(self):
+      return json.dumps({
+          k: (list(sorted(v)) if isinstance(v, set) else v)
+          for (k, v) in iteritems(self.__dict__)
+      },
+                        indent=4,
+                        sort_keys=True)
+
+    def update(self, other):
+      for key in ('cflags', 'defines', 'deps', 'include_dirs', 'ldflags',
+                  'source_set_deps', 'proto_deps', 'transitive_proto_deps',
+                  'libs', 'proto_paths'):
+        self.__dict__[key].update(other.__dict__.get(key, []))
+
+  def __init__(self, gn_desc):
+    self.gn_desc_ = gn_desc
+    self.all_targets = {}
+    self.linker_units = {}  # Executables, shared or static libraries.
+    self.source_sets = {}
+    self.actions = {}
+    self.proto_libs = {}
+
+  def get_target(self, gn_target_name):
+    """Returns a Target object from the fully qualified GN target name.
+
+        It bubbles up variables from source_set dependencies as described in the
+        class-level comments.
+        """
+    target = self.all_targets.get(gn_target_name)
+    if target is not None:
+      return target  # Target already processed.
+
+    desc = self.gn_desc_[gn_target_name]
+    target = GnParser.Target(gn_target_name, desc['type'])
+    target.testonly = desc.get('testonly', False)
+    target.toolchain = desc.get('toolchain', None)
+    self.all_targets[gn_target_name] = target
+
+    # We should never have GN targets directly depend on buidtools. They
+    # should hop via //gn:xxx, so we can give generators an opportunity to
+    # override them.
+    assert (not gn_target_name.startswith('//buildtools'))
+
+    # Don't descend further into third_party targets. Genrators are supposed
+    # to either ignore them or route to other externally-provided targets.
+    if gn_target_name.startswith('//gn'):
+      target.is_third_party_dep_ = True
+      return target
+
+    proto_target_type, proto_desc = self.get_proto_target_type(target)
+    if proto_target_type is not None:
+      self.proto_libs[target.name] = target
+      target.type = 'proto_library'
+      target.proto_plugin = proto_target_type
+      target.proto_paths.update(self.get_proto_paths(proto_desc))
+      target.proto_exports.update(self.get_proto_exports(proto_desc))
+      target.sources.update(proto_desc.get('sources', []))
+      assert (all(x.endswith('.proto') for x in target.sources))
+    elif target.type == 'source_set':
+      self.source_sets[gn_target_name] = target
+      target.sources.update(desc.get('sources', []))
+    elif target.type in LINKER_UNIT_TYPES:
+      self.linker_units[gn_target_name] = target
+      target.sources.update(desc.get('sources', []))
+    elif target.type == 'action':
+      self.actions[gn_target_name] = target
+      target.inputs.update(desc.get('inputs', []))
+      target.sources.update(desc.get('sources', []))
+      outs = [re.sub('^//out/.+?/gen/', '', x) for x in desc['outputs']]
+      target.outputs.update(outs)
+      target.script = desc['script']
+      # Args are typically relative to the root build dir (../../xxx)
+      # because root build dir is typically out/xxx/).
+      target.args = [re.sub('^../../', '//', x) for x in desc['args']]
+
+    # Default for 'public' is //* - all headers in 'sources' are public.
+    # TODO(primiano): if a 'public' section is specified (even if empty), then
+    # the rest of 'sources' is considered inaccessible by gn. Consider
+    # emulating that, so that generated build files don't end up with overly
+    # accessible headers.
+    public_headers = [x for x in desc.get('public', []) if x != '*']
+    target.public_headers.update(public_headers)
+
+    target.cflags.update(desc.get('cflags', []) + desc.get('cflags_cc', []))
+    target.libs.update(desc.get('libs', []))
+    target.ldflags.update(desc.get('ldflags', []))
+    target.defines.update(desc.get('defines', []))
+    target.include_dirs.update(desc.get('include_dirs', []))
+
+    # Recurse in dependencies.
+    for dep_name in desc.get('deps', []):
+      dep = self.get_target(dep_name)
+      if dep.is_third_party_dep_:
+        target.deps.add(dep_name)
+      elif dep.type == 'proto_library':
+        target.proto_deps.add(dep_name)
+        target.transitive_proto_deps.add(dep_name)
+        target.proto_paths.update(dep.proto_paths)
+        target.transitive_proto_deps.update(dep.transitive_proto_deps)
+      elif dep.type == 'source_set':
+        target.source_set_deps.add(dep_name)
+        target.update(dep)  # Bubble up source set's cflags/ldflags etc.
+      elif dep.type == 'group':
+        target.update(dep)  # Bubble up groups's cflags/ldflags etc.
+      elif dep.type == 'action':
+        if proto_target_type is None:
+          target.deps.add(dep_name)
+      elif dep.type in LINKER_UNIT_TYPES:
+        target.deps.add(dep_name)
+
+    return target
+
+  def get_proto_exports(self, proto_desc):
+    # exports in metadata will be available for source_set targets.
+    metadata = proto_desc.get('metadata', {})
+    return metadata.get('exports', [])
+
+  def get_proto_paths(self, proto_desc):
+    # import_dirs in metadata will be available for source_set targets.
+    metadata = proto_desc.get('metadata', {})
+    return metadata.get('import_dirs', [])
+
+  def get_proto_target_type(self, target):
+    """ Checks if the target is a proto library and return the plugin.
+
+        Returns:
+            (None, None): if the target is not a proto library.
+            (plugin, proto_desc) where |plugin| is 'proto' in the default (lite)
+            case or 'protozero' or 'ipc' or 'descriptor'; |proto_desc| is the GN
+            json desc of the target with the .proto sources (_gen target for
+            non-descriptor types or the target itself for descriptor type).
+        """
+    parts = target.name.split('(', 1)
+    name = parts[0]
+    toolchain = '(' + parts[1] if len(parts) > 1 else ''
+
+    # Descriptor targets don't have a _gen target; instead we look for the
+    # characteristic flag in the args of the target itself.
+    desc = self.gn_desc_.get(target.name)
+    if '--descriptor_set_out' in desc.get('args', []):
+      return 'descriptor', desc
+
+    # Source set proto targets have a non-empty proto_library_sources in the
+    # metadata of the description.
+    metadata = desc.get('metadata', {})
+    if 'proto_library_sources' in metadata:
+      return 'source_set', desc
+
+    # In all other cases, we want to look at the _gen target as that has the
+    # important information.
+    gen_desc = self.gn_desc_.get('%s_gen%s' % (name, toolchain))
+    if gen_desc is None or gen_desc['type'] != 'action':
+      return None, None
+    args = gen_desc.get('args', [])
+    if '/protoc' not in args[0]:
+      return None, None
+    plugin = 'proto'
+    for arg in (arg for arg in args if arg.startswith('--plugin=')):
+      # |arg| at this point looks like:
+      #  --plugin=protoc-gen-plugin=gcc_like_host/protozero_plugin
+      # or
+      #  --plugin=protoc-gen-plugin=protozero_plugin
+      plugin = arg.split('=')[-1].split('/')[-1].replace('_plugin', '')
+    return plugin, gen_desc
diff --git a/tools/gn2bp/run_buildtools_binary.py b/tools/gn2bp/run_buildtools_binary.py
new file mode 100755
index 0000000..7a6fe81
--- /dev/null
+++ b/tools/gn2bp/run_buildtools_binary.py
@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+# Copyright (C) 2022 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" A wrapper to run gn, ninja and other buildtools/ for all platforms. """
+
+from __future__ import print_function
+
+import os
+import subprocess
+import sys
+
+from platform import system, machine
+
+ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+
+def run_buildtools_binary(args):
+  if len(args) < 1:
+    print('Usage %s command [args]\n' % sys.argv[0])
+    return 1
+
+  sys_name = system().lower()
+  os_dir = None
+  ext = ''
+  if sys_name == 'windows':
+    os_dir = 'win'
+    ext = '.exe'
+  elif sys_name == 'darwin':
+    os_dir = 'mac'
+  elif sys_name == 'linux':
+    os_dir = 'linux64'
+  else:
+    print('OS not supported: %s\n' % sys_name)
+    return 1
+
+  cmd = args[0]
+  args = args[1:]
+  exe_path = os.path.join(ROOT_DIR, 'buildtools', os_dir, cmd) + ext
+  if sys_name == 'windows':
+    # execl() behaves oddly on Windows: the spawned process doesn't seem to
+    # receive CTRL+C. Use subprocess instead.
+    return subprocess.call([exe_path] + args)
+  else:
+    os.execl(exe_path, os.path.basename(exe_path), *args)
+
+
+if __name__ == '__main__':
+  sys.exit(run_buildtools_binary(sys.argv[1:]))