diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index 1b2a128..5cc95d9 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -208,6 +208,8 @@
 import com.android.server.net.NetworkPolicyManagerInternal;
 import com.android.testutils.ExceptionUtils;
 import com.android.testutils.HandlerUtilsKt;
+import com.android.testutils.RecorderCallback.CallbackRecord;
+import com.android.testutils.TestableNetworkCallback;
 
 import org.junit.After;
 import org.junit.Before;
@@ -234,7 +236,6 @@
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
@@ -243,7 +244,8 @@
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.function.Predicate;
+
+import kotlin.reflect.KClass;
 
 /**
  * Tests for {@link ConnectivityService}.
@@ -467,7 +469,7 @@
         fail("expected race condition at least once in " + attempts + " attempts");
     }
 
-    private class MockNetworkAgent {
+    private class MockNetworkAgent implements TestableNetworkCallback.HasNetwork {
         private static final int VALIDATION_RESULT_BASE = NETWORK_VALIDATION_PROBE_DNS
                 | NETWORK_VALIDATION_PROBE_HTTP
                 | NETWORK_VALIDATION_PROBE_HTTPS;
@@ -1703,281 +1705,33 @@
                 mCm.getDefaultRequest().networkCapabilities));
     }
 
-    enum CallbackState {
-        NONE,
-        AVAILABLE,
-        NETWORK_CAPABILITIES,
-        LINK_PROPERTIES,
-        SUSPENDED,
-        RESUMED,
-        LOSING,
-        LOST,
-        UNAVAILABLE,
-        BLOCKED_STATUS
-    }
-
-    private static class CallbackInfo {
-        public final CallbackState state;
-        public final Network network;
-        public final Object arg;
-        public CallbackInfo(CallbackState s, Network n, Object o) {
-            state = s; network = n; arg = o;
-        }
-        public String toString() {
-            return String.format("%s (%s) (%s)", state, network, arg);
-        }
-        @Override
-        public boolean equals(Object o) {
-            if (!(o instanceof CallbackInfo)) return false;
-            // Ignore timeMs, since it's unpredictable.
-            CallbackInfo other = (CallbackInfo) o;
-            return (state == other.state) && Objects.equals(network, other.network);
-        }
-        @Override
-        public int hashCode() {
-            return Objects.hash(state, network);
-        }
-    }
-
     /**
      * 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
      * received. assertNoCallback may be called at any time.
      */
-    private class TestNetworkCallback extends NetworkCallback {
-        private final LinkedBlockingQueue<CallbackInfo> mCallbacks = new LinkedBlockingQueue<>();
-        private Network mLastAvailableNetwork;
-
-        protected void setLastCallback(CallbackState state, Network network, Object o) {
-            mCallbacks.offer(new CallbackInfo(state, network, o));
+    private class TestNetworkCallback extends TestableNetworkCallback {
+        @Override
+        public void assertNoCallback() {
+            // TODO: better support this use case in TestableNetworkCallback
+            waitForIdle();
+            assertNoCallback(0 /* timeout */);
         }
 
         @Override
-        public void onAvailable(Network network) {
-            mLastAvailableNetwork = network;
-            setLastCallback(CallbackState.AVAILABLE, network, null);
-        }
-
-        @Override
-        public void onCapabilitiesChanged(Network network, NetworkCapabilities netCap) {
-            setLastCallback(CallbackState.NETWORK_CAPABILITIES, network, netCap);
-        }
-
-        @Override
-        public void onLinkPropertiesChanged(Network network, LinkProperties linkProp) {
-            setLastCallback(CallbackState.LINK_PROPERTIES, network, linkProp);
-        }
-
-        @Override
-        public void onUnavailable() {
-            setLastCallback(CallbackState.UNAVAILABLE, null, null);
-        }
-
-        @Override
-        public void onNetworkSuspended(Network network) {
-            setLastCallback(CallbackState.SUSPENDED, network, null);
-        }
-
-        @Override
-        public void onNetworkResumed(Network network) {
-            setLastCallback(CallbackState.RESUMED, network, null);
-        }
-
-        @Override
-        public void onLosing(Network network, int maxMsToLive) {
-            setLastCallback(CallbackState.LOSING, network, maxMsToLive /* autoboxed int */);
-        }
-
-        @Override
-        public void onLost(Network network) {
-            mLastAvailableNetwork = null;
-            setLastCallback(CallbackState.LOST, network, null);
-        }
-
-        @Override
-        public void onBlockedStatusChanged(Network network, boolean blocked) {
-            setLastCallback(CallbackState.BLOCKED_STATUS, network, blocked);
-        }
-
-        public Network getLastAvailableNetwork() {
-            return mLastAvailableNetwork;
-        }
-
-        CallbackInfo nextCallback(int timeoutMs) throws InterruptedException {
-            CallbackInfo cb = null;
-            cb = mCallbacks.poll(timeoutMs, TimeUnit.MILLISECONDS);
-            if (cb == null) {
-                // LinkedBlockingQueue.poll() returns null if it timeouts.
-                fail("Did not receive callback after " + timeoutMs + "ms");
-            }
-            return cb;
-        }
-
-        CallbackInfo expectCallback(CallbackState state, MockNetworkAgent agent, int timeoutMs)
-                throws Exception {
-            final Network expectedNetwork = (agent != null) ? agent.getNetwork() : null;
-            CallbackInfo expected = new CallbackInfo(state, expectedNetwork, 0);
-            CallbackInfo actual = nextCallback(timeoutMs);
-            assertEquals("Unexpected callback:", expected, actual);
-
-            if (state == CallbackState.LOSING) {
+        public <T extends CallbackRecord> T expectCallback(final KClass<T> type, final HasNetwork n,
+                final long timeoutMs) {
+            final T callback = super.expectCallback(type, n, timeoutMs);
+            if (callback instanceof CallbackRecord.Losing) {
+                // TODO : move this to the specific test(s) needing this rather than here.
+                final CallbackRecord.Losing losing = (CallbackRecord.Losing) callback;
+                final int maxMsToLive = losing.getMaxMsToLive();
                 String msg = String.format(
                         "Invalid linger time value %d, must be between %d and %d",
-                        actual.arg, 0, mService.mLingerDelayMs);
-                int maxMsToLive = (Integer) actual.arg;
+                        maxMsToLive, 0, mService.mLingerDelayMs);
                 assertTrue(msg, 0 <= maxMsToLive && maxMsToLive <= mService.mLingerDelayMs);
             }
-
-            return actual;
-        }
-
-        CallbackInfo expectCallback(CallbackState state, MockNetworkAgent agent) throws Exception {
-            return expectCallback(state, agent, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        CallbackInfo expectCallbackLike(Predicate<CallbackInfo> fn) throws Exception {
-            return expectCallbackLike(fn, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        CallbackInfo expectCallbackLike(Predicate<CallbackInfo> fn, int timeoutMs)
-                throws Exception {
-            int timeLeft = timeoutMs;
-            while (timeLeft > 0) {
-                long start = SystemClock.elapsedRealtime();
-                CallbackInfo info = nextCallback(timeLeft);
-                if (fn.test(info)) {
-                    return info;
-                }
-                timeLeft -= (SystemClock.elapsedRealtime() - start);
-            }
-            fail("Did not receive expected callback after " + timeoutMs + "ms");
-            return null;
-        }
-
-        // Expects onAvailable and the callbacks that follow it. These are:
-        // - onSuspended, iff the network was suspended when the callbacks fire.
-        // - onCapabilitiesChanged.
-        // - onLinkPropertiesChanged.
-        // - onBlockedStatusChanged.
-        //
-        // @param agent the network to expect the callbacks on.
-        // @param expectSuspended whether to expect a SUSPENDED callback.
-        // @param expectValidated the expected value of the VALIDATED capability in the
-        //        onCapabilitiesChanged callback.
-        // @param timeoutMs how long to wait for the callbacks.
-        void expectAvailableCallbacks(MockNetworkAgent agent, boolean expectSuspended,
-                boolean expectValidated, boolean expectBlocked, int timeoutMs) throws Exception {
-            expectCallback(CallbackState.AVAILABLE, agent, timeoutMs);
-            if (expectSuspended) {
-                expectCallback(CallbackState.SUSPENDED, agent, timeoutMs);
-            }
-            if (expectValidated) {
-                expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, agent, timeoutMs);
-            } else {
-                expectCapabilitiesWithout(NET_CAPABILITY_VALIDATED, agent, timeoutMs);
-            }
-            expectCallback(CallbackState.LINK_PROPERTIES, agent, timeoutMs);
-            expectBlockedStatusCallback(expectBlocked, agent);
-        }
-
-        // Expects the available callbacks (validated), plus onSuspended.
-        void expectAvailableAndSuspendedCallbacks(MockNetworkAgent agent, boolean expectValidated)
-                throws Exception {
-            expectAvailableCallbacks(agent, true, expectValidated, false, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        void expectAvailableCallbacksValidated(MockNetworkAgent agent)
-                throws Exception {
-            expectAvailableCallbacks(agent, false, true, false, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        void expectAvailableCallbacksValidatedAndBlocked(MockNetworkAgent agent) throws Exception {
-            expectAvailableCallbacks(agent, false, true, true, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        void expectAvailableCallbacksUnvalidated(MockNetworkAgent agent) throws Exception {
-            expectAvailableCallbacks(agent, false, false, false, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        void expectAvailableCallbacksUnvalidatedAndBlocked(MockNetworkAgent agent)
-                throws Exception {
-            expectAvailableCallbacks(agent, false, false, true, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        // Expects the available callbacks (where the onCapabilitiesChanged must contain the
-        // VALIDATED capability), plus another onCapabilitiesChanged which is identical to the
-        // one we just sent.
-        // TODO: this is likely a bug. Fix it and remove this method.
-        void expectAvailableDoubleValidatedCallbacks(MockNetworkAgent agent) throws Exception {
-            expectCallback(CallbackState.AVAILABLE, agent, TEST_CALLBACK_TIMEOUT_MS);
-            NetworkCapabilities nc1 = expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, agent);
-            expectCallback(CallbackState.LINK_PROPERTIES, agent, TEST_CALLBACK_TIMEOUT_MS);
-            // Implicitly check the network is allowed to use.
-            // TODO: should we need to consider if network is in blocked status in this case?
-            expectBlockedStatusCallback(false, agent);
-            NetworkCapabilities nc2 = expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, agent);
-            assertEquals(nc1, nc2);
-        }
-
-        // Expects the available callbacks where the onCapabilitiesChanged must not have validated,
-        // then expects another onCapabilitiesChanged that has the validated bit set. This is used
-        // when a network connects and satisfies a callback, and then immediately validates.
-        void expectAvailableThenValidatedCallbacks(MockNetworkAgent agent) throws Exception {
-            expectAvailableCallbacksUnvalidated(agent);
-            expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, agent);
-        }
-
-        NetworkCapabilities expectCapabilitiesWith(int capability, MockNetworkAgent agent)
-                throws Exception {
-            return expectCapabilitiesWith(capability, agent, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        NetworkCapabilities expectCapabilitiesWith(int capability, MockNetworkAgent agent,
-                int timeoutMs) throws Exception {
-            CallbackInfo cbi = expectCallback(CallbackState.NETWORK_CAPABILITIES, agent, timeoutMs);
-            NetworkCapabilities nc = (NetworkCapabilities) cbi.arg;
-            assertTrue(nc.hasCapability(capability));
-            return nc;
-        }
-
-        NetworkCapabilities expectCapabilitiesWithout(int capability, MockNetworkAgent agent)
-                throws Exception {
-            return expectCapabilitiesWithout(capability, agent, TEST_CALLBACK_TIMEOUT_MS);
-        }
-
-        NetworkCapabilities expectCapabilitiesWithout(int capability, MockNetworkAgent agent,
-                int timeoutMs) throws Exception {
-            CallbackInfo cbi = expectCallback(CallbackState.NETWORK_CAPABILITIES, agent, timeoutMs);
-            NetworkCapabilities nc = (NetworkCapabilities) cbi.arg;
-            assertFalse(nc.hasCapability(capability));
-            return nc;
-        }
-
-        void expectCapabilitiesLike(Predicate<NetworkCapabilities> fn, MockNetworkAgent agent)
-                throws Exception {
-            CallbackInfo cbi = expectCallback(CallbackState.NETWORK_CAPABILITIES, agent);
-            assertTrue("Received capabilities don't match expectations : " + cbi.arg,
-                    fn.test((NetworkCapabilities) cbi.arg));
-        }
-
-        void expectLinkPropertiesLike(Predicate<LinkProperties> fn, MockNetworkAgent agent)
-                throws Exception {
-            CallbackInfo cbi = expectCallback(CallbackState.LINK_PROPERTIES, agent);
-            assertTrue("Received LinkProperties don't match expectations : " + cbi.arg,
-                    fn.test((LinkProperties) cbi.arg));
-        }
-
-        void expectBlockedStatusCallback(boolean expectBlocked, MockNetworkAgent agent)
-                throws Exception {
-            CallbackInfo cbi = expectCallback(CallbackState.BLOCKED_STATUS, agent);
-            boolean actualBlocked = (boolean) cbi.arg;
-            assertEquals(expectBlocked, actualBlocked);
-        }
-
-        void assertNoCallback() {
-            waitForIdle();
-            CallbackInfo c = mCallbacks.peek();
-            assertNull("Unexpected callback: " + c, c);
+            return callback;
         }
     }
 
@@ -2031,16 +1785,16 @@
 
         cv = waitForConnectivityBroadcasts(2);
         mWiFiNetworkAgent.disconnect();
-        genericNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        wifiNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        wifiNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         cellNetworkCallback.assertNoCallback();
         waitFor(cv);
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
 
         cv = waitForConnectivityBroadcasts(1);
         mCellNetworkAgent.disconnect();
-        genericNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         waitFor(cv);
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
 
@@ -2061,21 +1815,21 @@
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connect(true);
         genericNetworkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        genericNetworkCallback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         genericNetworkCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         wifiNetworkCallback.expectAvailableThenValidatedCallbacks(mWiFiNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
 
         mWiFiNetworkAgent.disconnect();
-        genericNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        wifiNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        wifiNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
 
         mCellNetworkAgent.disconnect();
-        genericNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         assertNoCallbacks(genericNetworkCallback, wifiNetworkCallback, cellNetworkCallback);
     }
 
@@ -2115,7 +1869,7 @@
         // We then get LOSING when wifi validates and cell is outscored.
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         // TODO: Investigate sending validated before losing.
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
@@ -2124,15 +1878,15 @@
         mEthernetNetworkAgent.connect(true);
         callback.expectAvailableCallbacksUnvalidated(mEthernetNetworkAgent);
         // TODO: Investigate sending validated before losing.
-        callback.expectCallback(CallbackState.LOSING, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mWiFiNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mEthernetNetworkAgent);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mEthernetNetworkAgent);
         assertEquals(mEthernetNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         mEthernetNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mEthernetNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.LOST, mEthernetNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mEthernetNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mEthernetNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
@@ -2148,7 +1902,7 @@
                 newNetwork = mWiFiNetworkAgent;
 
             }
-            callback.expectCallback(CallbackState.LOSING, oldNetwork);
+            callback.expectCallback(CallbackRecord.LOSING, oldNetwork);
             // TODO: should we send an AVAILABLE callback to newNetwork, to indicate that it is no
             // longer lingering?
             defaultCallback.expectAvailableCallbacksValidated(newNetwork);
@@ -2162,7 +1916,7 @@
         // We expect a notification about the capabilities change, and nothing else.
         defaultCallback.expectCapabilitiesWithout(NET_CAPABILITY_NOT_METERED, mWiFiNetworkAgent);
         defaultCallback.assertNoCallback();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Wifi no longer satisfies our listen, which is for an unmetered network.
@@ -2171,11 +1925,11 @@
 
         // Disconnect our test networks.
         mWiFiNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
         mCellNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         waitForIdle();
         assertEquals(null, mCm.getActiveNetwork());
 
@@ -2206,8 +1960,8 @@
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
@@ -2218,15 +1972,15 @@
         mWiFiNetworkAgent.adjustScore(50);
         mWiFiNetworkAgent.connect(false);   // Score: 70
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Tear down wifi.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
         assertEquals(mCellNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
@@ -2237,19 +1991,19 @@
         mWiFiNetworkAgent.connect(true);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         // TODO: Investigate sending validated before losing.
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         defaultCallback.expectAvailableThenValidatedCallbacks(mWiFiNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mCellNetworkAgent);
         mCellNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         waitForIdle();
         assertEquals(null, mCm.getActiveNetwork());
 
@@ -2264,7 +2018,7 @@
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         // TODO: Investigate sending validated before losing.
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
@@ -2275,13 +2029,13 @@
         // TODO: should this cause an AVAILABLE callback, to indicate that the network is no longer
         // lingering?
         mCm.unregisterNetworkCallback(noopCallback);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
 
         // Similar to the above: lingering can start even after the lingered request is removed.
         // Disconnect wifi and switch to cell.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
@@ -2300,12 +2054,12 @@
         callback.assertNoCallback();
         // Now unregister cellRequest and expect cell to start lingering.
         mCm.unregisterNetworkCallback(noopCallback);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
 
         // Let linger run its course.
         callback.assertNoCallback();
         final int lingerTimeoutMs = mService.mLingerDelayMs + mService.mLingerDelayMs / 4;
-        callback.expectCallback(CallbackState.LOST, mCellNetworkAgent, lingerTimeoutMs);
+        callback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent, lingerTimeoutMs);
 
         // Register a TRACK_DEFAULT request and check that it does not affect lingering.
         TestNetworkCallback trackDefaultCallback = new TestNetworkCallback();
@@ -2314,20 +2068,20 @@
         mEthernetNetworkAgent = new MockNetworkAgent(TRANSPORT_ETHERNET);
         mEthernetNetworkAgent.connect(true);
         callback.expectAvailableCallbacksUnvalidated(mEthernetNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mWiFiNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mEthernetNetworkAgent);
         trackDefaultCallback.expectAvailableDoubleValidatedCallbacks(mEthernetNetworkAgent);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mEthernetNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         // Let linger run its course.
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent, lingerTimeoutMs);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent, lingerTimeoutMs);
 
         // Clean up.
         mEthernetNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mEthernetNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.LOST, mEthernetNetworkAgent);
-        trackDefaultCallback.expectCallback(CallbackState.LOST, mEthernetNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mEthernetNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mEthernetNetworkAgent);
+        trackDefaultCallback.expectCallback(CallbackRecord.LOST, mEthernetNetworkAgent);
 
         mCm.unregisterNetworkCallback(callback);
         mCm.unregisterNetworkCallback(defaultCallback);
@@ -2357,7 +2111,7 @@
         mWiFiNetworkAgent.connect(true);
         defaultCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
 
         // File a request for cellular, then release it.
@@ -2366,7 +2120,7 @@
         NetworkCallback noopCallback = new NetworkCallback();
         mCm.requestNetwork(cellRequest, noopCallback);
         mCm.unregisterNetworkCallback(noopCallback);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
 
         // Let linger run its course.
         callback.assertNoCallback();
@@ -2410,12 +2164,12 @@
         // If the user chooses yes on the "No Internet access, stay connected?" dialog, we switch to
         // wifi even though it's unvalidated.
         mCm.setAcceptUnvalidated(mWiFiNetworkAgent.getNetwork(), true, false);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
 
         // Disconnect wifi, and then reconnect, again with explicitlySelected=true.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.explicitlySelected(true, false);
         mWiFiNetworkAgent.connect(false);
@@ -2424,14 +2178,14 @@
         // If the user chooses no on the "No Internet access, stay connected?" dialog, we ask the
         // network to disconnect.
         mCm.setAcceptUnvalidated(mWiFiNetworkAgent.getNetwork(), false, false);
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // Reconnect, again with explicitlySelected=true, but this time validate.
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.explicitlySelected(true, false);
         mWiFiNetworkAgent.connect(true);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
 
@@ -2447,20 +2201,20 @@
         // (i.e., with explicitlySelected=true and acceptUnvalidated=true). Expect to switch to
         // wifi immediately.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.explicitlySelected(true, true);
         mWiFiNetworkAgent.connect(false);
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mEthernetNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mEthernetNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         mEthernetNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mEthernetNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mEthernetNetworkAgent);
 
         // 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(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.explicitlySelected(false, true);
         mWiFiNetworkAgent.connect(false);
@@ -2471,8 +2225,8 @@
         mWiFiNetworkAgent.disconnect();
         mCellNetworkAgent.disconnect();
 
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
     }
 
     private int[] makeIntArray(final int size, final int value) {
@@ -2721,7 +2475,7 @@
         // Need a trigger point to let NetworkMonitor tell ConnectivityService that network is
         // validated.
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), true);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         NetworkCapabilities nc = callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED,
                 mWiFiNetworkAgent);
         assertTrue(nc.hasCapability(NET_CAPABILITY_PARTIAL_CONNECTIVITY));
@@ -2729,7 +2483,7 @@
 
         // Disconnect and reconnect wifi with partial connectivity again.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.connectWithPartialConnectivity();
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
@@ -2741,7 +2495,7 @@
         // If the user chooses no, disconnect wifi immediately.
         mCm.setAcceptPartialConnectivity(mWiFiNetworkAgent.getNetwork(), false/* accept */,
                 false /* always */);
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // If user accepted partial connectivity before, and device reconnects to that network
         // again, but now the network has full connectivity. The network shouldn't contain
@@ -2758,14 +2512,14 @@
         waitForIdle();
         verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         nc = callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         assertFalse(nc.hasCapability(NET_CAPABILITY_PARTIAL_CONNECTIVITY));
 
         // Wifi should be the default network.
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.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
@@ -2780,7 +2534,7 @@
         waitForIdle();
         verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
         callback.expectCapabilitiesWith(NET_CAPABILITY_PARTIAL_CONNECTIVITY, mWiFiNetworkAgent);
         mWiFiNetworkAgent.setNetworkValid();
@@ -2790,7 +2544,7 @@
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), true);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.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.
@@ -2805,11 +2559,11 @@
         waitForIdle();
         verify(mWiFiNetworkAgent.mNetworkMonitor, times(1)).setAcceptPartialConnectivity();
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         callback.expectCapabilitiesWith(
                 NET_CAPABILITY_PARTIAL_CONNECTIVITY | NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
     }
 
     @Test
@@ -2851,7 +2605,7 @@
                 false /* always */);
         waitForIdle();
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), true);
-        captivePortalCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        captivePortalCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         validatedCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
         NetworkCapabilities nc =
                 validatedCallback.expectCapabilitiesWith(NET_CAPABILITY_PARTIAL_CONNECTIVITY,
@@ -2884,7 +2638,7 @@
         // Take down network.
         // Expect onLost callback.
         mWiFiNetworkAgent.disconnect();
-        captivePortalCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        captivePortalCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // Bring up a network with a captive portal.
         // Expect onAvailable callback of listen for NET_CAPABILITY_CAPTIVE_PORTAL.
@@ -2898,7 +2652,7 @@
         // Expect onLost callback because network no longer provides NET_CAPABILITY_CAPTIVE_PORTAL.
         mWiFiNetworkAgent.setNetworkValid();
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), true);
-        captivePortalCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        captivePortalCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // Expect NET_CAPABILITY_VALIDATED onAvailable callback.
         validatedCallback.expectAvailableDoubleValidatedCallbacks(mWiFiNetworkAgent);
@@ -2910,7 +2664,7 @@
         // Expect NET_CAPABILITY_VALIDATED onLost callback.
         mWiFiNetworkAgent.setNetworkInvalid();
         mCm.reportNetworkConnectivity(mWiFiNetworkAgent.getNetwork(), false);
-        validatedCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        validatedCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
     }
 
     @Test
@@ -2942,7 +2696,7 @@
         mWiFiNetworkAgent.setNetworkPortal("http://example.com");
         mCm.reportNetworkConnectivity(wifiNetwork, false);
         captivePortalCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        validatedCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        validatedCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // Check that startCaptivePortalApp sends the expected command to NetworkMonitor.
         mCm.startCaptivePortalApp(wifiNetwork);
@@ -2963,7 +2717,7 @@
         mWiFiNetworkAgent.setNetworkValid();
         mWiFiNetworkAgent.mNetworkMonitor.forceReevaluation(Process.myUid());
         validatedCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
-        captivePortalCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        captivePortalCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         verify(mNotificationManager, times(1)).notifyAsUser(anyString(),
                 eq(NotificationType.LOGGED_IN.eventId), any(), eq(UserHandle.ALL));
 
@@ -3103,24 +2857,24 @@
         mWiFiNetworkAgent.setNetworkSpecifier(nsFoo);
         cFoo.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         for (TestNetworkCallback c: emptyCallbacks) {
-            c.expectCapabilitiesLike((caps) -> caps.getNetworkSpecifier().equals(nsFoo),
-                    mWiFiNetworkAgent);
+            c.expectCapabilitiesThat(mWiFiNetworkAgent,
+                    (caps) -> caps.getNetworkSpecifier().equals(nsFoo));
         }
-        cFoo.expectCapabilitiesLike((caps) -> caps.getNetworkSpecifier().equals(nsFoo),
-                mWiFiNetworkAgent);
+        cFoo.expectCapabilitiesThat(mWiFiNetworkAgent,
+                (caps) -> caps.getNetworkSpecifier().equals(nsFoo));
         assertEquals(nsFoo,
                 mCm.getNetworkCapabilities(mWiFiNetworkAgent.getNetwork()).getNetworkSpecifier());
         cFoo.assertNoCallback();
 
         mWiFiNetworkAgent.setNetworkSpecifier(nsBar);
-        cFoo.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        cFoo.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         cBar.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         for (TestNetworkCallback c: emptyCallbacks) {
-            c.expectCapabilitiesLike((caps) -> caps.getNetworkSpecifier().equals(nsBar),
-                    mWiFiNetworkAgent);
+            c.expectCapabilitiesThat(mWiFiNetworkAgent,
+                    (caps) -> caps.getNetworkSpecifier().equals(nsBar));
         }
-        cBar.expectCapabilitiesLike((caps) -> caps.getNetworkSpecifier().equals(nsBar),
-                mWiFiNetworkAgent);
+        cBar.expectCapabilitiesThat(mWiFiNetworkAgent,
+                (caps) -> caps.getNetworkSpecifier().equals(nsBar));
         assertEquals(nsBar,
                 mCm.getNetworkCapabilities(mWiFiNetworkAgent.getNetwork()).getNetworkSpecifier());
         cBar.assertNoCallback();
@@ -3128,23 +2882,23 @@
         mWiFiNetworkAgent.setNetworkSpecifier(new ConfidentialMatchAllNetworkSpecifier());
         cFoo.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         for (TestNetworkCallback c : emptyCallbacks) {
-            c.expectCapabilitiesLike((caps) -> caps.getNetworkSpecifier() == null,
-                    mWiFiNetworkAgent);
+            c.expectCapabilitiesThat(mWiFiNetworkAgent,
+                    (caps) -> caps.getNetworkSpecifier() == null);
         }
-        cFoo.expectCapabilitiesLike((caps) -> caps.getNetworkSpecifier() == null,
-                mWiFiNetworkAgent);
-        cBar.expectCapabilitiesLike((caps) -> caps.getNetworkSpecifier() == null,
-                mWiFiNetworkAgent);
+        cFoo.expectCapabilitiesThat(mWiFiNetworkAgent,
+                (caps) -> caps.getNetworkSpecifier() == null);
+        cBar.expectCapabilitiesThat(mWiFiNetworkAgent,
+                (caps) -> caps.getNetworkSpecifier() == null);
         assertNull(
                 mCm.getNetworkCapabilities(mWiFiNetworkAgent.getNetwork()).getNetworkSpecifier());
         cFoo.assertNoCallback();
         cBar.assertNoCallback();
 
         mWiFiNetworkAgent.setNetworkSpecifier(null);
-        cFoo.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        cBar.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        cFoo.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        cBar.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         for (TestNetworkCallback c: emptyCallbacks) {
-            c.expectCallback(CallbackState.NETWORK_CAPABILITIES, mWiFiNetworkAgent);
+            c.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED, mWiFiNetworkAgent);
         }
 
         assertNoCallbacks(cEmpty1, cEmpty2, cEmpty3, cEmpty4, cFoo, cBar);
@@ -3287,7 +3041,7 @@
 
         // Bring down cell. Expect no default network callback, since it wasn't the default.
         mCellNetworkAgent.disconnect();
-        cellNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         defaultNetworkCallback.assertNoCallback();
         assertEquals(defaultNetworkCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
@@ -3302,11 +3056,11 @@
         // followed by AVAILABLE cell.
         mWiFiNetworkAgent.disconnect();
         cellNetworkCallback.assertNoCallback();
-        defaultNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        defaultNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         defaultNetworkCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         mCellNetworkAgent.disconnect();
-        cellNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
-        defaultNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
+        defaultNetworkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         waitForIdle();
         assertEquals(null, mCm.getActiveNetwork());
 
@@ -3322,7 +3076,7 @@
         assertEquals(defaultNetworkCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         vpnNetworkAgent.disconnect();
-        defaultNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        defaultNetworkCallback.expectCallback(CallbackRecord.LOST, vpnNetworkAgent);
         waitForIdle();
         assertEquals(null, mCm.getActiveNetwork());
     }
@@ -3350,14 +3104,15 @@
         lp.setInterfaceName("foonet_data0");
         mCellNetworkAgent.sendLinkProperties(lp);
         // We should get onLinkPropertiesChanged().
-        cellNetworkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED,
+                mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
 
         // Suspend the network.
         mCellNetworkAgent.suspend();
         cellNetworkCallback.expectCapabilitiesWithout(NET_CAPABILITY_NOT_SUSPENDED,
                 mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.SUSPENDED, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.SUSPENDED, mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
 
         // Register a garden variety default network request.
@@ -3372,7 +3127,7 @@
         mCellNetworkAgent.resume();
         cellNetworkCallback.expectCapabilitiesWith(NET_CAPABILITY_NOT_SUSPENDED,
                 mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.RESUMED, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.RESUMED, mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
 
         dfltNetworkCallback = new TestNetworkCallback();
@@ -3435,10 +3190,10 @@
 
         // When wifi connects, cell lingers.
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         callback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         fgCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        fgCallback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        fgCallback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         fgCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         assertTrue(isForegroundNetwork(mCellNetworkAgent));
         assertTrue(isForegroundNetwork(mWiFiNetworkAgent));
@@ -3446,7 +3201,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(CallbackState.LOST, mCellNetworkAgent, timeoutMs);
+        fgCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent, timeoutMs);
         // Expect a network capabilities update sans FOREGROUND.
         callback.expectCapabilitiesWithout(NET_CAPABILITY_FOREGROUND, mCellNetworkAgent);
         assertFalse(isForegroundNetwork(mCellNetworkAgent));
@@ -3472,7 +3227,7 @@
         // Release the request. The network immediately goes into the background, since it was not
         // lingering.
         mCm.unregisterNetworkCallback(cellCallback);
-        fgCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        fgCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         // Expect a network capabilities update sans FOREGROUND.
         callback.expectCapabilitiesWithout(NET_CAPABILITY_FOREGROUND, mCellNetworkAgent);
         assertFalse(isForegroundNetwork(mCellNetworkAgent));
@@ -3480,8 +3235,8 @@
 
         // Disconnect wifi and check that cell is foreground again.
         mWiFiNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        fgCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        fgCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         fgCallback.expectAvailableCallbacksValidated(mCellNetworkAgent);
         assertTrue(isForegroundNetwork(mCellNetworkAgent));
 
@@ -3618,7 +3373,7 @@
         testFactory.waitForNetworkRequests(1);
 
         // ...  and cell data to be torn down.
-        cellNetworkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         assertLength(1, mCm.getAllNetworks());
 
         testFactory.unregister();
@@ -3709,7 +3464,7 @@
         mWiFiNetworkAgent.setNetworkInvalid();
         mCm.reportNetworkConnectivity(wifiNetwork, false);
         defaultCallback.expectCapabilitiesWithout(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
-        validatedWifiCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        validatedWifiCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // Because avoid bad wifi is off, we don't switch to cellular.
         defaultCallback.assertNoCallback();
@@ -3753,7 +3508,7 @@
         mWiFiNetworkAgent.setNetworkInvalid();
         mCm.reportNetworkConnectivity(wifiNetwork, false);
         defaultCallback.expectCapabilitiesWithout(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
-        validatedWifiCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        validatedWifiCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // Simulate the user selecting "switch" and checking the don't ask again checkbox.
         Settings.Global.putInt(cr, Settings.Global.NETWORK_AVOID_BAD_WIFI, 1);
@@ -3780,7 +3535,7 @@
 
         // If cell goes down, we switch to wifi.
         mCellNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
         validatedWifiCallback.assertNoCallback();
 
@@ -3845,7 +3600,7 @@
         networkCallback.expectAvailableCallbacks(mWiFiNetworkAgent, false, false, false,
                 TEST_CALLBACK_TIMEOUT_MS);
         mWiFiNetworkAgent.disconnect();
-        networkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
 
         // Validate that UNAVAILABLE is not called
         networkCallback.assertNoCallback();
@@ -3865,7 +3620,7 @@
         mCm.requestNetwork(nr, networkCallback, timeoutMs);
 
         // pass timeout and validate that UNAVAILABLE is called
-        networkCallback.expectCallback(CallbackState.UNAVAILABLE, null);
+        networkCallback.expectCallback(CallbackRecord.UNAVAILABLE, null);
 
         // create a network satisfying request - validate that request not triggered
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
@@ -3956,7 +3711,7 @@
             // Simulate the factory releasing the request as unfulfillable and expect onUnavailable!
             testFactory.triggerUnfulfillable(requests.get(newRequestId));
 
-            networkCallback.expectCallback(CallbackState.UNAVAILABLE, null);
+            networkCallback.expectCallback(CallbackRecord.UNAVAILABLE, null);
             testFactory.waitForRequests();
 
             // unregister network callback - a no-op (since already freed by the
@@ -4797,7 +4552,7 @@
 
         // Disconnect wifi aware network.
         wifiAware.disconnect();
-        callback.expectCallbackLike((info) -> info.state == CallbackState.LOST, TIMEOUT_MS);
+        callback.expectCallbackThat(TIMEOUT_MS, (info) -> info instanceof CallbackRecord.Lost);
         mCm.unregisterNetworkCallback(callback);
 
         verifyNoNetwork();
@@ -4846,14 +4601,15 @@
         // ConnectivityService.
         MockNetworkAgent networkAgent = new MockNetworkAgent(TRANSPORT_WIFI, lp);
         networkAgent.connect(true);
-        networkCallback.expectCallback(CallbackState.AVAILABLE, networkAgent);
-        networkCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES, networkAgent);
-        CallbackInfo cbi = networkCallback.expectCallback(CallbackState.LINK_PROPERTIES,
+        networkCallback.expectCallback(CallbackRecord.AVAILABLE, networkAgent);
+        networkCallback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED, networkAgent);
+        CallbackRecord.LinkPropertiesChanged cbi =
+                networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED,
                 networkAgent);
-        networkCallback.expectCallback(CallbackState.BLOCKED_STATUS, networkAgent);
+        networkCallback.expectCallback(CallbackRecord.BLOCKED_STATUS, networkAgent);
         networkCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, networkAgent);
         networkCallback.assertNoCallback();
-        checkDirectlyConnectedRoutes(cbi.arg, Arrays.asList(myIpv4Address),
+        checkDirectlyConnectedRoutes(cbi.getLp(), Arrays.asList(myIpv4Address),
                 Arrays.asList(myIpv4DefaultRoute));
         checkDirectlyConnectedRoutes(mCm.getLinkProperties(networkAgent.getNetwork()),
                 Arrays.asList(myIpv4Address), Arrays.asList(myIpv4DefaultRoute));
@@ -4865,9 +4621,9 @@
         newLp.addLinkAddress(myIpv6Address1);
         newLp.addLinkAddress(myIpv6Address2);
         networkAgent.sendLinkProperties(newLp);
-        cbi = networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, networkAgent);
+        cbi = networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, networkAgent);
         networkCallback.assertNoCallback();
-        checkDirectlyConnectedRoutes(cbi.arg,
+        checkDirectlyConnectedRoutes(cbi.getLp(),
                 Arrays.asList(myIpv4Address, myIpv6Address1, myIpv6Address2),
                 Arrays.asList(myIpv4DefaultRoute));
         mCm.unregisterNetworkCallback(networkCallback);
@@ -5071,15 +4827,15 @@
         assertTrue(ArrayUtils.containsAll(resolvrParams.tlsServers,
                 new String[] { "2001:db8::1", "192.0.2.1" }));
         reset(mMockDnsResolver);
-        cellNetworkCallback.expectCallback(CallbackState.AVAILABLE, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES,
+        cellNetworkCallback.expectCallback(CallbackRecord.AVAILABLE, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED,
                 mCellNetworkAgent);
-        CallbackInfo cbi = cellNetworkCallback.expectCallback(
-                CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.BLOCKED_STATUS, mCellNetworkAgent);
+        CallbackRecord.LinkPropertiesChanged cbi = cellNetworkCallback.expectCallback(
+                CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.BLOCKED_STATUS, mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
-        assertFalse(((LinkProperties)cbi.arg).isPrivateDnsActive());
-        assertNull(((LinkProperties)cbi.arg).getPrivateDnsServerName());
+        assertFalse(cbi.getLp().isPrivateDnsActive());
+        assertNull(cbi.getLp().getPrivateDnsServerName());
 
         setPrivateDnsSettings(PRIVATE_DNS_MODE_OFF, "ignored.example.com");
         verify(mMockDnsResolver, times(1)).setResolverConfiguration(
@@ -5107,11 +4863,11 @@
         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(CallbackState.LINK_PROPERTIES,
+        cbi = cellNetworkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
-        assertTrue(((LinkProperties)cbi.arg).isPrivateDnsActive());
-        assertEquals("strict.example.com", ((LinkProperties)cbi.arg).getPrivateDnsServerName());
+        assertTrue(cbi.getLp().isPrivateDnsActive());
+        assertEquals("strict.example.com", cbi.getLp().getPrivateDnsServerName());
     }
 
     @Test
@@ -5130,17 +4886,17 @@
         mCellNetworkAgent.sendLinkProperties(lp);
         mCellNetworkAgent.connect(false);
         waitForIdle();
-        cellNetworkCallback.expectCallback(CallbackState.AVAILABLE, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES,
+        cellNetworkCallback.expectCallback(CallbackRecord.AVAILABLE, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED,
                 mCellNetworkAgent);
-        CallbackInfo cbi = cellNetworkCallback.expectCallback(
-                CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
-        cellNetworkCallback.expectCallback(CallbackState.BLOCKED_STATUS, mCellNetworkAgent);
+        CallbackRecord.LinkPropertiesChanged cbi = cellNetworkCallback.expectCallback(
+                CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        cellNetworkCallback.expectCallback(CallbackRecord.BLOCKED_STATUS, mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
-        assertFalse(((LinkProperties)cbi.arg).isPrivateDnsActive());
-        assertNull(((LinkProperties)cbi.arg).getPrivateDnsServerName());
+        assertFalse(cbi.getLp().isPrivateDnsActive());
+        assertNull(cbi.getLp().getPrivateDnsServerName());
         Set<InetAddress> dnsServers = new HashSet<>();
-        checkDnsServers(cbi.arg, dnsServers);
+        checkDnsServers(cbi.getLp(), dnsServers);
 
         // Send a validation event for a server that is not part of the current
         // resolver config. The validation event should be ignored.
@@ -5152,13 +4908,13 @@
         LinkProperties lp2 = new LinkProperties(lp);
         lp2.addDnsServer(InetAddress.getByName("145.100.185.16"));
         mCellNetworkAgent.sendLinkProperties(lp2);
-        cbi = cellNetworkCallback.expectCallback(CallbackState.LINK_PROPERTIES,
+        cbi = cellNetworkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
-        assertFalse(((LinkProperties)cbi.arg).isPrivateDnsActive());
-        assertNull(((LinkProperties)cbi.arg).getPrivateDnsServerName());
+        assertFalse(cbi.getLp().isPrivateDnsActive());
+        assertNull(cbi.getLp().getPrivateDnsServerName());
         dnsServers.add(InetAddress.getByName("145.100.185.16"));
-        checkDnsServers(cbi.arg, dnsServers);
+        checkDnsServers(cbi.getLp(), dnsServers);
 
         // Send a validation event containing a hostname that is not part of
         // the current resolver config. The validation event should be ignored.
@@ -5176,39 +4932,39 @@
         // private dns fields should be sent.
         mService.mNetdEventCallback.onPrivateDnsValidationEvent(
                 mCellNetworkAgent.getNetwork().netId, "145.100.185.16", "", true);
-        cbi = cellNetworkCallback.expectCallback(CallbackState.LINK_PROPERTIES,
+        cbi = cellNetworkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
-        assertTrue(((LinkProperties)cbi.arg).isPrivateDnsActive());
-        assertNull(((LinkProperties)cbi.arg).getPrivateDnsServerName());
-        checkDnsServers(cbi.arg, dnsServers);
+        assertTrue(cbi.getLp().isPrivateDnsActive());
+        assertNull(cbi.getLp().getPrivateDnsServerName());
+        checkDnsServers(cbi.getLp(), dnsServers);
 
         // The private dns fields in LinkProperties should be preserved when
         // the network agent sends unrelated changes.
         LinkProperties lp3 = new LinkProperties(lp2);
         lp3.setMtu(1300);
         mCellNetworkAgent.sendLinkProperties(lp3);
-        cbi = cellNetworkCallback.expectCallback(CallbackState.LINK_PROPERTIES,
+        cbi = cellNetworkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
-        assertTrue(((LinkProperties)cbi.arg).isPrivateDnsActive());
-        assertNull(((LinkProperties)cbi.arg).getPrivateDnsServerName());
-        checkDnsServers(cbi.arg, dnsServers);
-        assertEquals(1300, ((LinkProperties)cbi.arg).getMtu());
+        assertTrue(cbi.getLp().isPrivateDnsActive());
+        assertNull(cbi.getLp().getPrivateDnsServerName());
+        checkDnsServers(cbi.getLp(), dnsServers);
+        assertEquals(1300, cbi.getLp().getMtu());
 
         // Removing the only validated server should affect the private dns
         // fields in LinkProperties.
         LinkProperties lp4 = new LinkProperties(lp3);
         lp4.removeDnsServer(InetAddress.getByName("145.100.185.16"));
         mCellNetworkAgent.sendLinkProperties(lp4);
-        cbi = cellNetworkCallback.expectCallback(CallbackState.LINK_PROPERTIES,
+        cbi = cellNetworkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED,
                 mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
-        assertFalse(((LinkProperties)cbi.arg).isPrivateDnsActive());
-        assertNull(((LinkProperties)cbi.arg).getPrivateDnsServerName());
+        assertFalse(cbi.getLp().isPrivateDnsActive());
+        assertNull(cbi.getLp().getPrivateDnsServerName());
         dnsServers.remove(InetAddress.getByName("145.100.185.16"));
-        checkDnsServers(cbi.arg, dnsServers);
-        assertEquals(1300, ((LinkProperties)cbi.arg).getMtu());
+        checkDnsServers(cbi.getLp(), dnsServers);
+        assertEquals(1300, cbi.getLp().getMtu());
     }
 
     private void checkDirectlyConnectedRoutes(Object callbackObj,
@@ -5290,19 +5046,19 @@
         defaultCallback.expectAvailableCallbacksUnvalidated(vpnNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
-        genericNetworkCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES, vpnNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED, vpnNetworkAgent);
         genericNotVpnNetworkCallback.assertNoCallback();
-        vpnNetworkCallback.expectCapabilitiesLike(nc -> null == nc.getUids(), vpnNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES, vpnNetworkAgent);
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent, nc -> null == nc.getUids());
+        defaultCallback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED, vpnNetworkAgent);
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         ranges.clear();
         vpnNetworkAgent.setUids(ranges);
 
-        genericNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOST, vpnNetworkAgent);
         genericNotVpnNetworkCallback.assertNoCallback();
         wifiNetworkCallback.assertNoCallback();
-        vpnNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        vpnNetworkCallback.expectCallback(CallbackRecord.LOST, vpnNetworkAgent);
 
         // 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
@@ -5310,7 +5066,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(CallbackState.NETWORK_CAPABILITIES, vpnNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED, vpnNetworkAgent);
 
         ranges.add(new UidRange(uid, uid));
         mMockVpn.setUids(ranges);
@@ -5322,23 +5078,23 @@
         vpnNetworkCallback.expectAvailableCallbacksValidated(vpnNetworkAgent);
         // 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(CallbackState.NETWORK_CAPABILITIES, vpnNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED, vpnNetworkAgent);
 
         mWiFiNetworkAgent.disconnect();
 
-        genericNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        genericNotVpnNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
-        wifiNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        genericNotVpnNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
+        wifiNetworkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         vpnNetworkCallback.assertNoCallback();
         defaultCallback.assertNoCallback();
 
         vpnNetworkAgent.disconnect();
 
-        genericNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        genericNetworkCallback.expectCallback(CallbackRecord.LOST, vpnNetworkAgent);
         genericNotVpnNetworkCallback.assertNoCallback();
         wifiNetworkCallback.assertNoCallback();
-        vpnNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
-        defaultCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        vpnNetworkCallback.expectCallback(CallbackRecord.LOST, vpnNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, vpnNetworkAgent);
         assertEquals(null, mCm.getActiveNetwork());
 
         mCm.unregisterNetworkCallback(genericNetworkCallback);
@@ -5402,7 +5158,7 @@
         assertEquals(defaultCallback.getLastAvailableNetwork(), mCm.getActiveNetwork());
 
         vpnNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, vpnNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidated(mWiFiNetworkAgent);
 
         mCm.unregisterNetworkCallback(defaultCallback);
@@ -5432,7 +5188,7 @@
         // Even though the VPN is unvalidated, it becomes the default network for our app.
         callback.expectAvailableCallbacksUnvalidated(vpnNetworkAgent);
         // TODO: this looks like a spurious callback.
-        callback.expectCallback(CallbackState.NETWORK_CAPABILITIES, vpnNetworkAgent);
+        callback.expectCallback(CallbackRecord.NETWORK_CAPS_UPDATED, vpnNetworkAgent);
         callback.assertNoCallback();
 
         assertTrue(vpnNetworkAgent.getScore() > mEthernetNetworkAgent.getScore());
@@ -5456,7 +5212,7 @@
         callback.assertNoCallback();
 
         vpnNetworkAgent.disconnect();
-        callback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        callback.expectCallback(CallbackRecord.LOST, vpnNetworkAgent);
         callback.expectAvailableCallbacksValidated(mEthernetNetworkAgent);
     }
 
@@ -5496,10 +5252,10 @@
         mService.setUnderlyingNetworksForVpn(
                 new Network[] { mCellNetworkAgent.getNetwork() });
 
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
-                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.addCapability(NET_CAPABILITY_NOT_METERED);
@@ -5508,52 +5264,52 @@
         mService.setUnderlyingNetworksForVpn(
                 new Network[] { mCellNetworkAgent.getNetwork(), mWiFiNetworkAgent.getNetwork() });
 
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
-                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // Don't disconnect, but note the VPN is not using wifi any more.
         mService.setUnderlyingNetworksForVpn(
                 new Network[] { mCellNetworkAgent.getNetwork() });
 
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
-                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // Use Wifi but not cell. Note the VPN is now unmetered.
         mService.setUnderlyingNetworksForVpn(
                 new Network[] { mWiFiNetworkAgent.getNetwork() });
 
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && !caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
-                && caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // Use both again.
         mService.setUnderlyingNetworksForVpn(
                 new Network[] { mCellNetworkAgent.getNetwork(), mWiFiNetworkAgent.getNetwork() });
 
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
-                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // Disconnect cell. Receive update without even removing the dead network from the
         // underlying networks – it's dead anyway. Not metered any more.
         mCellNetworkAgent.disconnect();
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && !caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
-                && caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // Disconnect wifi too. No underlying networks means this is now metered.
         mWiFiNetworkAgent.disconnect();
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && !caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
-                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         mMockVpn.disconnect();
     }
@@ -5592,20 +5348,20 @@
         mCellNetworkAgent = new MockNetworkAgent(TRANSPORT_CELLULAR);
         mCellNetworkAgent.connect(true);
 
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
-                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // Connect to WiFi; WiFi is the new default.
         mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
         mWiFiNetworkAgent.addCapability(NET_CAPABILITY_NOT_METERED);
         mWiFiNetworkAgent.connect(true);
 
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && !caps.hasTransport(TRANSPORT_CELLULAR) && caps.hasTransport(TRANSPORT_WIFI)
-                && caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // Disconnect Cell. The default network did not change, so there shouldn't be any changes in
         // the capabilities.
@@ -5614,10 +5370,10 @@
         // Disconnect wifi too. Now we have no default network.
         mWiFiNetworkAgent.disconnect();
 
-        vpnNetworkCallback.expectCapabilitiesLike((caps) -> caps.hasTransport(TRANSPORT_VPN)
+        vpnNetworkCallback.expectCapabilitiesThat(vpnNetworkAgent,
+                (caps) -> caps.hasTransport(TRANSPORT_VPN)
                 && !caps.hasTransport(TRANSPORT_CELLULAR) && !caps.hasTransport(TRANSPORT_WIFI)
-                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED),
-                vpnNetworkAgent);
+                && !caps.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         mMockVpn.disconnect();
     }
@@ -5893,7 +5649,7 @@
 
         // Switch to METERED network. Restrict the use of the network.
         mWiFiNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         defaultCallback.expectAvailableCallbacksValidatedAndBlocked(mCellNetworkAgent);
 
         // Network becomes NOT_METERED.
@@ -5907,7 +5663,7 @@
         defaultCallback.assertNoCallback();
 
         mCellNetworkAgent.disconnect();
-        defaultCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        defaultCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         defaultCallback.assertNoCallback();
 
         mCm.unregisterNetworkCallback(defaultCallback);
@@ -5983,7 +5739,7 @@
         // the NAT64 prefix was removed because one was never discovered.
         cellLp.addLinkAddress(myIpv4);
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         verify(mMockDnsResolver, times(1)).stopPrefix64Discovery(cellNetId);
         verify(mMockDnsResolver, atLeastOnce()).setResolverConfiguration(any());
 
@@ -5996,7 +5752,7 @@
         cellLp.removeLinkAddress(myIpv4);
         cellLp.removeRoute(new RouteInfo(myIpv4, null, MOBILE_IFNAME));
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         verify(mMockDnsResolver, times(1)).startPrefix64Discovery(cellNetId);
 
         // When NAT64 prefix discovery succeeds, LinkProperties are updated and clatd is started.
@@ -6004,15 +5760,15 @@
         assertNull(mCm.getLinkProperties(mCellNetworkAgent.getNetwork()).getNat64Prefix());
         mService.mNetdEventCallback.onNat64PrefixEvent(cellNetId, true /* added */,
                 kNat64PrefixString, 96);
-        LinkProperties lpBeforeClat = (LinkProperties) networkCallback.expectCallback(
-                CallbackState.LINK_PROPERTIES, mCellNetworkAgent).arg;
+        LinkProperties lpBeforeClat = networkCallback.expectCallback(
+                CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent).getLp();
         assertEquals(0, lpBeforeClat.getStackedLinks().size());
         assertEquals(kNat64Prefix, lpBeforeClat.getNat64Prefix());
         verify(mMockNetd, times(1)).clatdStart(MOBILE_IFNAME, kNat64Prefix.toString());
 
         // Clat iface comes up. Expect stacked link to be added.
         clat.interfaceLinkStateChanged(CLAT_PREFIX + MOBILE_IFNAME, true);
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         List<LinkProperties> stackedLps = mCm.getLinkProperties(mCellNetworkAgent.getNetwork())
                 .getStackedLinks();
         assertEquals(makeClatLinkProperties(myIpv4), stackedLps.get(0));
@@ -6020,7 +5776,7 @@
         // Change trivial linkproperties and see if stacked link is preserved.
         cellLp.addDnsServer(InetAddress.getByName("8.8.8.8"));
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
 
         List<LinkProperties> stackedLpsAfterChange =
                 mCm.getLinkProperties(mCellNetworkAgent.getNetwork()).getStackedLinks();
@@ -6038,12 +5794,12 @@
         cellLp.addLinkAddress(myIpv4);
         cellLp.addRoute(new RouteInfo(myIpv4, null, MOBILE_IFNAME));
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         verify(mMockNetd, times(1)).clatdStop(MOBILE_IFNAME);
         verify(mMockDnsResolver, times(1)).stopPrefix64Discovery(cellNetId);
 
         // As soon as stop is called, the linkproperties lose the stacked interface.
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         LinkProperties actualLpAfterIpv4 = mCm.getLinkProperties(mCellNetworkAgent.getNetwork());
         LinkProperties expected = new LinkProperties(cellLp);
         expected.setNat64Prefix(kNat64Prefix);
@@ -6062,41 +5818,39 @@
         // Stopping prefix discovery causes netd to tell us that the NAT64 prefix is gone.
         mService.mNetdEventCallback.onNat64PrefixEvent(cellNetId, false /* added */,
                 kNat64PrefixString, 96);
-        networkCallback.expectLinkPropertiesLike((lp) -> lp.getNat64Prefix() == null,
-                mCellNetworkAgent);
+        networkCallback.expectLinkPropertiesThat(mCellNetworkAgent,
+                (lp) -> lp.getNat64Prefix() == null);
 
         // Remove IPv4 address and expect prefix discovery and clatd to be started again.
         cellLp.removeLinkAddress(myIpv4);
         cellLp.removeRoute(new RouteInfo(myIpv4, null, MOBILE_IFNAME));
         cellLp.removeDnsServer(InetAddress.getByName("8.8.8.8"));
         mCellNetworkAgent.sendLinkProperties(cellLp);
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         verify(mMockDnsResolver, times(1)).startPrefix64Discovery(cellNetId);
         mService.mNetdEventCallback.onNat64PrefixEvent(cellNetId, true /* added */,
                 kNat64PrefixString, 96);
-        networkCallback.expectCallback(CallbackState.LINK_PROPERTIES, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         verify(mMockNetd, times(1)).clatdStart(MOBILE_IFNAME, kNat64Prefix.toString());
 
 
         // Clat iface comes up. Expect stacked link to be added.
         clat.interfaceLinkStateChanged(CLAT_PREFIX + MOBILE_IFNAME, true);
-        networkCallback.expectLinkPropertiesLike(
-                (lp) -> lp.getStackedLinks().size() == 1 && lp.getNat64Prefix() != null,
-                mCellNetworkAgent);
+        networkCallback.expectLinkPropertiesThat(mCellNetworkAgent,
+                (lp) -> lp.getStackedLinks().size() == 1 && lp.getNat64Prefix() != null);
 
         // NAT64 prefix is removed. Expect that clat is stopped.
         mService.mNetdEventCallback.onNat64PrefixEvent(cellNetId, false /* added */,
                 kNat64PrefixString, 96);
-        networkCallback.expectLinkPropertiesLike(
-                (lp) -> lp.getStackedLinks().size() == 0 && lp.getNat64Prefix() == null,
-                mCellNetworkAgent);
+        networkCallback.expectLinkPropertiesThat(mCellNetworkAgent,
+                (lp) -> lp.getStackedLinks().size() == 0 && lp.getNat64Prefix() == null);
         verify(mMockNetd, times(1)).clatdStop(MOBILE_IFNAME);
-        networkCallback.expectLinkPropertiesLike((lp) -> lp.getStackedLinks().size() == 0,
-                mCellNetworkAgent);
+        networkCallback.expectLinkPropertiesThat(mCellNetworkAgent,
+                (lp) -> lp.getStackedLinks().size() == 0);
 
         // Clean up.
         mCellNetworkAgent.disconnect();
-        networkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LOST, mCellNetworkAgent);
         networkCallback.assertNoCallback();
         mCm.unregisterNetworkCallback(networkCallback);
     }
@@ -6128,7 +5882,7 @@
         reset(mNetworkManagementService);
         mWiFiNetworkAgent.connect(true);
         networkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        networkCallback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         networkCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
         verify(mNetworkManagementService, times(1)).addIdleTimer(eq(WIFI_IFNAME), anyInt(),
                 eq(ConnectivityManager.TYPE_WIFI));
@@ -6137,7 +5891,7 @@
         // Disconnect wifi and switch back to cell
         reset(mNetworkManagementService);
         mWiFiNetworkAgent.disconnect();
-        networkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LOST, mWiFiNetworkAgent);
         assertNoCallbacks(networkCallback);
         verify(mNetworkManagementService, times(1)).removeIdleTimer(eq(WIFI_IFNAME));
         verify(mNetworkManagementService, times(1)).addIdleTimer(eq(MOBILE_IFNAME), anyInt(),
@@ -6149,14 +5903,14 @@
         mWiFiNetworkAgent.sendLinkProperties(wifiLp);
         mWiFiNetworkAgent.connect(true);
         networkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        networkCallback.expectCallback(CallbackState.LOSING, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.LOSING, mCellNetworkAgent);
         networkCallback.expectCapabilitiesWith(NET_CAPABILITY_VALIDATED, mWiFiNetworkAgent);
 
         // Disconnect cell
         reset(mNetworkManagementService);
         reset(mMockNetd);
         mCellNetworkAgent.disconnect();
-        networkCallback.expectCallback(CallbackState.LOST, mCellNetworkAgent);
+        networkCallback.expectCallback(CallbackRecord.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.