diff --git a/service/jarjar-rules.txt b/service/jarjar-rules.txt
index e5d1a88..cc0cbcc 100644
--- a/service/jarjar-rules.txt
+++ b/service/jarjar-rules.txt
@@ -90,5 +90,9 @@
 # From services-connectivity-shared-srcs
 rule android.net.util.NetworkConstants* com.android.connectivity.@0
 
+# From modules-utils-statemachine
+rule com.android.internal.util.IState* com.android.connectivity.@0
+rule com.android.internal.util.State* com.android.connectivity.@0
+
 # Remaining are connectivity sources in com.android.server and com.android.server.connectivity:
 # TODO: move to a subpackage of com.android.connectivity (such as com.android.connectivity.server)
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index 80c2db4..53e4ab7 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -168,6 +168,7 @@
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
 import com.android.testutils.DevSdkIgnoreRuleKt;
+import com.android.testutils.DumpTestUtils;
 import com.android.testutils.RecorderCallback.CallbackEntry;
 import com.android.testutils.TestHttpServer;
 import com.android.testutils.TestNetworkTracker;
@@ -3021,6 +3022,13 @@
         }
     }
 
+    @Test
+    public void testDump() throws Exception {
+        final String dumpOutput = DumpTestUtils.dumpServiceWithShellPermission(
+                Context.CONNECTIVITY_SERVICE, "--short");
+        assertTrue(dumpOutput, dumpOutput.contains("Active default network"));
+    }
+
     private void unregisterRegisteredCallbacks() {
         for (NetworkCallback callback: mRegisteredCallbacks) {
             mCm.unregisterNetworkCallback(callback);
diff --git a/tests/unit/Android.bp b/tests/unit/Android.bp
index 4c605f3..0bef0e2 100644
--- a/tests/unit/Android.bp
+++ b/tests/unit/Android.bp
@@ -90,6 +90,7 @@
         "java/com/android/server/net/ipmemorystore/*.java",
         "java/com/android/server/net/BpfInterfaceMapUpdaterTest.java",
         "java/com/android/server/net/NetworkStats*.java",
+        "java/com/android/server/net/TestableUsageCallback.kt",
     ]
 }
 
diff --git a/tests/unit/java/com/android/server/NetworkManagementServiceTest.java b/tests/unit/java/com/android/server/NetworkManagementServiceTest.java
index ea29da0..0c58582 100644
--- a/tests/unit/java/com/android/server/NetworkManagementServiceTest.java
+++ b/tests/unit/java/com/android/server/NetworkManagementServiceTest.java
@@ -22,6 +22,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
@@ -32,6 +33,7 @@
 
 import android.annotation.NonNull;
 import android.content.Context;
+import android.net.ConnectivityManager;
 import android.net.INetd;
 import android.net.INetdUnsolicitedEventListener;
 import android.net.LinkAddress;
@@ -71,6 +73,7 @@
 public class NetworkManagementServiceTest {
     private NetworkManagementService mNMService;
     @Mock private Context mContext;
+    @Mock private ConnectivityManager mCm;
     @Mock private IBatteryStats.Stub mBatteryStatsService;
     @Mock private INetd.Stub mNetdService;
 
@@ -113,6 +116,9 @@
         MockitoAnnotations.initMocks(this);
         doNothing().when(mNetdService)
                 .registerUnsolicitedEventListener(mUnsolListenerCaptor.capture());
+        doReturn(Context.CONNECTIVITY_SERVICE).when(mContext).getSystemServiceName(
+                eq(ConnectivityManager.class));
+        doReturn(mCm).when(mContext).getSystemService(eq(Context.CONNECTIVITY_SERVICE));
         // Start the service and wait until it connects to our socket.
         mNMService = NetworkManagementService.create(mContext, mDeps);
     }
@@ -239,6 +245,7 @@
         mNMService.setUidOnMeteredNetworkDenylist(TEST_UID, true);
         assertTrue("Should be true since mobile data usage is restricted",
                 mNMService.isNetworkRestricted(TEST_UID));
+        verify(mCm).updateMeteredNetworkDenyList(TEST_UID, true /* enabled */);
 
         mNMService.setDataSaverModeEnabled(true);
         verify(mNetdService).bandwidthEnableDataSaver(true);
@@ -246,13 +253,16 @@
         mNMService.setUidOnMeteredNetworkDenylist(TEST_UID, false);
         assertTrue("Should be true since data saver is on and the uid is not allowlisted",
                 mNMService.isNetworkRestricted(TEST_UID));
+        verify(mCm).updateMeteredNetworkDenyList(TEST_UID, true /* false */);
 
         mNMService.setUidOnMeteredNetworkAllowlist(TEST_UID, true);
         assertFalse("Should be false since data saver is on and the uid is allowlisted",
                 mNMService.isNetworkRestricted(TEST_UID));
+        verify(mCm).updateMeteredNetworkAllowList(TEST_UID, true /* enabled */);
 
         // remove uid from allowlist and turn datasaver off again
         mNMService.setUidOnMeteredNetworkAllowlist(TEST_UID, false);
+        verify(mCm).updateMeteredNetworkAllowList(TEST_UID, false /* enabled */);
         mNMService.setDataSaverModeEnabled(false);
         verify(mNetdService).bandwidthEnableDataSaver(false);
         assertFalse("Network should not be restricted when data saver is off",
@@ -306,12 +316,14 @@
         for (int chain : chains) {
             final ArrayMap<Integer, Boolean> expectedValues = expected.get(chain);
             mNMService.setFirewallChainEnabled(chain, true);
+            verify(mCm).setFirewallChainEnabled(chain, true /* enabled */);
             for (int state : states) {
                 mNMService.setFirewallUidRule(chain, TEST_UID, state);
                 assertEquals(errorMsg.apply(chain, state),
                         expectedValues.get(state), mNMService.isNetworkRestricted(TEST_UID));
             }
             mNMService.setFirewallChainEnabled(chain, false);
+            verify(mCm).setFirewallChainEnabled(chain, false /* enabled */);
         }
     }
 }
diff --git a/tests/unit/java/com/android/server/connectivity/MultipathPolicyTrackerTest.java b/tests/unit/java/com/android/server/connectivity/MultipathPolicyTrackerTest.java
index c86e699..ec51537 100644
--- a/tests/unit/java/com/android/server/connectivity/MultipathPolicyTrackerTest.java
+++ b/tests/unit/java/com/android/server/connectivity/MultipathPolicyTrackerTest.java
@@ -296,7 +296,7 @@
                 false /* roaming */);
 
         verify(mStatsManager, times(1)).registerUsageCallback(
-                any(), anyInt(), eq(DataUnit.MEGABYTES.toBytes(12)), any(), any());
+                any(), eq(DataUnit.MEGABYTES.toBytes(12)), any(), any());
     }
 
     @Test
@@ -315,7 +315,7 @@
 
         // Daily budget should be 15MB (5% of daily quota), 7MB used today: callback set for 8MB
         verify(mStatsManager, times(1)).registerUsageCallback(
-                any(), anyInt(), eq(DataUnit.MEGABYTES.toBytes(8)), any(), any());
+                any(), eq(DataUnit.MEGABYTES.toBytes(8)), any(), any());
     }
 
     @Test
@@ -334,7 +334,7 @@
 
         // Daily budget should be 15MB (5% of daily quota), 7MB used today: callback set for 8MB
         verify(mStatsManager, times(1)).registerUsageCallback(
-                any(), anyInt(), eq(DataUnit.MEGABYTES.toBytes(8)), any(), any());
+                any(), eq(DataUnit.MEGABYTES.toBytes(8)), any(), any());
     }
 
     @Test
@@ -351,7 +351,7 @@
 
         // Default global setting should be used: 12 - 7 = 5
         verify(mStatsManager, times(1)).registerUsageCallback(
-                any(), anyInt(), eq(DataUnit.MEGABYTES.toBytes(5)), any(), any());
+                any(), eq(DataUnit.MEGABYTES.toBytes(5)), any(), any());
     }
 
     @Test
@@ -366,7 +366,7 @@
                 false /* roaming */);
 
         verify(mStatsManager, times(1)).registerUsageCallback(
-                any(), anyInt(), eq(DataUnit.MEGABYTES.toBytes(8)), any(), any());
+                any(), eq(DataUnit.MEGABYTES.toBytes(8)), any(), any());
 
         // Update setting
         setDefaultQuotaGlobalSetting(DataUnit.MEGABYTES.toBytes(14));
@@ -376,7 +376,7 @@
         // Callback must have been re-registered with new setting
         verify(mStatsManager, times(1)).unregisterUsageCallback(any());
         verify(mStatsManager, times(1)).registerUsageCallback(
-                any(), anyInt(), eq(DataUnit.MEGABYTES.toBytes(12)), any(), any());
+                any(), eq(DataUnit.MEGABYTES.toBytes(12)), any(), any());
     }
 
     @Test
@@ -391,7 +391,7 @@
                 false /* roaming */);
 
         verify(mStatsManager, times(1)).registerUsageCallback(
-                any(), anyInt(), eq(DataUnit.MEGABYTES.toBytes(12)), any(), any());
+                any(), eq(DataUnit.MEGABYTES.toBytes(12)), any(), any());
 
         when(mResources.getInteger(R.integer.config_networkDefaultDailyMultipathQuotaBytes))
                 .thenReturn((int) DataUnit.MEGABYTES.toBytes(16));
@@ -402,6 +402,6 @@
 
         // Uses the new setting (16 - 2 = 14MB)
         verify(mStatsManager, times(1)).registerUsageCallback(
-                any(), anyInt(), eq(DataUnit.MEGABYTES.toBytes(14)), any(), any());
+                any(), eq(DataUnit.MEGABYTES.toBytes(14)), any(), any());
     }
 }
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsFactoryTest.java b/tests/unit/java/com/android/server/net/NetworkStatsFactoryTest.java
index 8340a13..0e2c293 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsFactoryTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsFactoryTest.java
@@ -32,9 +32,12 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.doReturn;
 
+import android.content.Context;
 import android.content.res.Resources;
-import android.net.INetd;
+import android.net.ConnectivityManager;
 import android.net.NetworkStats;
 import android.net.TrafficStats;
 import android.net.UnderlyingNetworkInfo;
@@ -73,8 +76,8 @@
 
     private File mTestProc;
     private NetworkStatsFactory mFactory;
-    @Mock
-    private INetd mNetd;
+    @Mock private Context mContext;
+    @Mock private ConnectivityManager mCm;
 
     @Before
     public void setUp() throws Exception {
@@ -85,7 +88,10 @@
         // applications. So in order to have a test support native library, the native code
         // related to networkStatsFactory is compiled to a minimal native library and loaded here.
         System.loadLibrary("networkstatsfactorytestjni");
-        mFactory = new NetworkStatsFactory(mTestProc, false, mNetd);
+        doReturn(Context.CONNECTIVITY_SERVICE).when(mContext).getSystemServiceName(
+                eq(ConnectivityManager.class));
+        doReturn(mCm).when(mContext).getSystemService(eq(Context.CONNECTIVITY_SERVICE));
+        mFactory = new NetworkStatsFactory(mContext, mTestProc, false);
         mFactory.updateUnderlyingNetworkInfos(new UnderlyingNetworkInfo[0]);
     }
 
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsObserversTest.java b/tests/unit/java/com/android/server/net/NetworkStatsObserversTest.java
index d993d1f..66dcf6d 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsObserversTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsObserversTest.java
@@ -36,19 +36,15 @@
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
 
-import android.app.usage.NetworkStatsManager;
 import android.net.DataUsageRequest;
 import android.net.NetworkIdentity;
 import android.net.NetworkIdentitySet;
 import android.net.NetworkStats;
 import android.net.NetworkStatsAccess;
 import android.net.NetworkTemplate;
-import android.os.ConditionVariable;
-import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.IBinder;
 import android.os.Looper;
-import android.os.Messenger;
 import android.os.Process;
 import android.os.UserHandle;
 import android.telephony.TelephonyManager;
@@ -56,7 +52,6 @@
 
 import androidx.test.filters.SmallTest;
 
-import com.android.server.net.NetworkStatsServiceTest.LatchedHandler;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRunner;
 import com.android.testutils.HandlerUtils;
@@ -97,21 +92,15 @@
     private static final long WAIT_TIMEOUT_MS = 500;
     private static final long THRESHOLD_BYTES = 2 * MB_IN_BYTES;
     private static final long BASE_BYTES = 7 * MB_IN_BYTES;
-    private static final int INVALID_TYPE = -1;
-
-    private long mElapsedRealtime;
 
     private HandlerThread mObserverHandlerThread;
-    private Handler mObserverNoopHandler;
-
-    private LatchedHandler mHandler;
 
     private NetworkStatsObservers mStatsObservers;
-    private Messenger mMessenger;
     private ArrayMap<String, NetworkIdentitySet> mActiveIfaces;
     private ArrayMap<String, NetworkIdentitySet> mActiveUidIfaces;
 
-    @Mock private IBinder mockBinder;
+    @Mock private IBinder mUsageCallbackBinder;
+    private TestableUsageCallback mUsageCallback;
 
     @Before
     public void setUp() throws Exception {
@@ -127,24 +116,29 @@
             }
         };
 
-        mHandler = new LatchedHandler(Looper.getMainLooper(), new ConditionVariable());
-        mMessenger = new Messenger(mHandler);
-
         mActiveIfaces = new ArrayMap<>();
         mActiveUidIfaces = new ArrayMap<>();
+        mUsageCallback = new TestableUsageCallback(mUsageCallbackBinder);
     }
 
     @Test
     public void testRegister_thresholdTooLow_setsDefaultThreshold() throws Exception {
-        long thresholdTooLowBytes = 1L;
-        DataUsageRequest inputRequest = new DataUsageRequest(
+        final long thresholdTooLowBytes = 1L;
+        final DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateWifi, thresholdTooLowBytes);
 
-        DataUsageRequest request = mStatsObservers.register(inputRequest, mMessenger, mockBinder,
-                Process.SYSTEM_UID, NetworkStatsAccess.Level.DEVICE);
-        assertTrue(request.requestId > 0);
-        assertTrue(Objects.equals(sTemplateWifi, request.template));
-        assertEquals(THRESHOLD_BYTES, request.thresholdInBytes);
+        final DataUsageRequest requestByApp = mStatsObservers.register(inputRequest, mUsageCallback,
+                UID_RED, NetworkStatsAccess.Level.DEVICE);
+        assertTrue(requestByApp.requestId > 0);
+        assertTrue(Objects.equals(sTemplateWifi, requestByApp.template));
+        assertEquals(THRESHOLD_BYTES, requestByApp.thresholdInBytes);
+
+        // Verify the threshold requested by system uid won't be overridden.
+        final DataUsageRequest requestBySystem = mStatsObservers.register(inputRequest,
+                mUsageCallback, Process.SYSTEM_UID, NetworkStatsAccess.Level.DEVICE);
+        assertTrue(requestBySystem.requestId > 0);
+        assertTrue(Objects.equals(sTemplateWifi, requestBySystem.template));
+        assertEquals(1, requestBySystem.thresholdInBytes);
     }
 
     @Test
@@ -153,7 +147,7 @@
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateWifi, highThresholdBytes);
 
-        DataUsageRequest request = mStatsObservers.register(inputRequest, mMessenger, mockBinder,
+        DataUsageRequest request = mStatsObservers.register(inputRequest, mUsageCallback,
                 Process.SYSTEM_UID, NetworkStatsAccess.Level.DEVICE);
         assertTrue(request.requestId > 0);
         assertTrue(Objects.equals(sTemplateWifi, request.template));
@@ -165,13 +159,13 @@
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateWifi, THRESHOLD_BYTES);
 
-        DataUsageRequest request1 = mStatsObservers.register(inputRequest, mMessenger, mockBinder,
+        DataUsageRequest request1 = mStatsObservers.register(inputRequest, mUsageCallback,
                 Process.SYSTEM_UID, NetworkStatsAccess.Level.DEVICE);
         assertTrue(request1.requestId > 0);
         assertTrue(Objects.equals(sTemplateWifi, request1.template));
         assertEquals(THRESHOLD_BYTES, request1.thresholdInBytes);
 
-        DataUsageRequest request2 = mStatsObservers.register(inputRequest, mMessenger, mockBinder,
+        DataUsageRequest request2 = mStatsObservers.register(inputRequest, mUsageCallback,
                 Process.SYSTEM_UID, NetworkStatsAccess.Level.DEVICE);
         assertTrue(request2.requestId > request1.requestId);
         assertTrue(Objects.equals(sTemplateWifi, request2.template));
@@ -191,17 +185,19 @@
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
 
-        DataUsageRequest request = mStatsObservers.register(inputRequest, mMessenger, mockBinder,
+        DataUsageRequest request = mStatsObservers.register(inputRequest, mUsageCallback,
                 Process.SYSTEM_UID, NetworkStatsAccess.Level.DEVICE);
         assertTrue(request.requestId > 0);
         assertTrue(Objects.equals(sTemplateImsi1, request.template));
         assertEquals(THRESHOLD_BYTES, request.thresholdInBytes);
-        Mockito.verify(mockBinder).linkToDeath(any(IBinder.DeathRecipient.class), anyInt());
+        Mockito.verify(mUsageCallbackBinder).linkToDeath(any(IBinder.DeathRecipient.class),
+                anyInt());
 
         mStatsObservers.unregister(request, Process.SYSTEM_UID);
         waitForObserverToIdle();
 
-        Mockito.verify(mockBinder).unlinkToDeath(any(IBinder.DeathRecipient.class), anyInt());
+        Mockito.verify(mUsageCallbackBinder).unlinkToDeath(any(IBinder.DeathRecipient.class),
+                anyInt());
     }
 
     @Test
@@ -209,17 +205,18 @@
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
 
-        DataUsageRequest request = mStatsObservers.register(inputRequest, mMessenger, mockBinder,
+        DataUsageRequest request = mStatsObservers.register(inputRequest, mUsageCallback,
                 UID_RED, NetworkStatsAccess.Level.DEVICE);
         assertTrue(request.requestId > 0);
         assertTrue(Objects.equals(sTemplateImsi1, request.template));
         assertEquals(THRESHOLD_BYTES, request.thresholdInBytes);
-        Mockito.verify(mockBinder).linkToDeath(any(IBinder.DeathRecipient.class), anyInt());
+        Mockito.verify(mUsageCallbackBinder)
+                .linkToDeath(any(IBinder.DeathRecipient.class), anyInt());
 
         mStatsObservers.unregister(request, UID_BLUE);
         waitForObserverToIdle();
 
-        Mockito.verifyZeroInteractions(mockBinder);
+        Mockito.verifyZeroInteractions(mUsageCallbackBinder);
     }
 
     private NetworkIdentitySet makeTestIdentSet() {
@@ -236,7 +233,7 @@
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
 
-        DataUsageRequest request = mStatsObservers.register(inputRequest, mMessenger, mockBinder,
+        DataUsageRequest request = mStatsObservers.register(inputRequest, mUsageCallback,
                 Process.SYSTEM_UID, NetworkStatsAccess.Level.DEVICE);
         assertTrue(request.requestId > 0);
         assertTrue(Objects.equals(sTemplateImsi1, request.template));
@@ -260,7 +257,7 @@
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
 
-        DataUsageRequest request = mStatsObservers.register(inputRequest, mMessenger, mockBinder,
+        DataUsageRequest request = mStatsObservers.register(inputRequest, mUsageCallback,
                 Process.SYSTEM_UID, NetworkStatsAccess.Level.DEVICE);
         assertTrue(request.requestId > 0);
         assertTrue(Objects.equals(sTemplateImsi1, request.template));
@@ -290,7 +287,7 @@
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
 
-        DataUsageRequest request = mStatsObservers.register(inputRequest, mMessenger, mockBinder,
+        DataUsageRequest request = mStatsObservers.register(inputRequest, mUsageCallback,
                 Process.SYSTEM_UID, NetworkStatsAccess.Level.DEVICE);
         assertTrue(request.requestId > 0);
         assertTrue(Objects.equals(sTemplateImsi1, request.template));
@@ -313,7 +310,7 @@
         mStatsObservers.updateStats(
                 xtSnapshot, uidSnapshot, mActiveIfaces, mActiveUidIfaces, TEST_START);
         waitForObserverToIdle();
-        assertEquals(NetworkStatsManager.CALLBACK_LIMIT_REACHED, mHandler.lastMessageType);
+        mUsageCallback.expectOnThresholdReached(request);
     }
 
     @Test
@@ -321,7 +318,7 @@
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
 
-        DataUsageRequest request = mStatsObservers.register(inputRequest, mMessenger, mockBinder,
+        DataUsageRequest request = mStatsObservers.register(inputRequest, mUsageCallback,
                 UID_RED, NetworkStatsAccess.Level.DEFAULT);
         assertTrue(request.requestId > 0);
         assertTrue(Objects.equals(sTemplateImsi1, request.template));
@@ -346,7 +343,7 @@
         mStatsObservers.updateStats(
                 xtSnapshot, uidSnapshot, mActiveIfaces, mActiveUidIfaces, TEST_START);
         waitForObserverToIdle();
-        assertEquals(NetworkStatsManager.CALLBACK_LIMIT_REACHED, mHandler.lastMessageType);
+        mUsageCallback.expectOnThresholdReached(request);
     }
 
     @Test
@@ -354,7 +351,7 @@
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
 
-        DataUsageRequest request = mStatsObservers.register(inputRequest, mMessenger, mockBinder,
+        DataUsageRequest request = mStatsObservers.register(inputRequest, mUsageCallback,
                 UID_BLUE, NetworkStatsAccess.Level.DEFAULT);
         assertTrue(request.requestId > 0);
         assertTrue(Objects.equals(sTemplateImsi1, request.template));
@@ -386,7 +383,7 @@
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
 
-        DataUsageRequest request = mStatsObservers.register(inputRequest, mMessenger, mockBinder,
+        DataUsageRequest request = mStatsObservers.register(inputRequest, mUsageCallback,
                 UID_BLUE, NetworkStatsAccess.Level.USER);
         assertTrue(request.requestId > 0);
         assertTrue(Objects.equals(sTemplateImsi1, request.template));
@@ -411,7 +408,7 @@
         mStatsObservers.updateStats(
                 xtSnapshot, uidSnapshot, mActiveIfaces, mActiveUidIfaces, TEST_START);
         waitForObserverToIdle();
-        assertEquals(NetworkStatsManager.CALLBACK_LIMIT_REACHED, mHandler.lastMessageType);
+        mUsageCallback.expectOnThresholdReached(request);
     }
 
     @Test
@@ -419,7 +416,7 @@
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateImsi1, THRESHOLD_BYTES);
 
-        DataUsageRequest request = mStatsObservers.register(inputRequest, mMessenger, mockBinder,
+        DataUsageRequest request = mStatsObservers.register(inputRequest, mUsageCallback,
                 UID_RED, NetworkStatsAccess.Level.USER);
         assertTrue(request.requestId > 0);
         assertTrue(Objects.equals(sTemplateImsi1, request.template));
@@ -448,6 +445,5 @@
 
     private void waitForObserverToIdle() {
         HandlerUtils.waitForIdle(mObserverHandlerThread, WAIT_TIMEOUT_MS);
-        HandlerUtils.waitForIdle(mHandler, WAIT_TIMEOUT_MS);
     }
 }
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index 99b1194..ea35c31 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -80,7 +80,6 @@
 
 import android.annotation.NonNull;
 import android.app.AlarmManager;
-import android.app.usage.NetworkStatsManager;
 import android.content.Context;
 import android.content.Intent;
 import android.database.ContentObserver;
@@ -101,13 +100,9 @@
 import android.net.netstats.provider.INetworkStatsProviderCallback;
 import android.net.wifi.WifiInfo;
 import android.os.Build;
-import android.os.ConditionVariable;
 import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.IBinder;
-import android.os.Looper;
-import android.os.Message;
-import android.os.Messenger;
 import android.os.PowerManager;
 import android.os.SimpleClock;
 import android.provider.Settings;
@@ -188,7 +183,8 @@
     private @Mock TetheringManager mTetheringManager;
     private @Mock NetworkStatsFactory mStatsFactory;
     private @Mock NetworkStatsSettings mSettings;
-    private @Mock IBinder mBinder;
+    private @Mock IBinder mUsageCallbackBinder;
+    private TestableUsageCallback mUsageCallback;
     private @Mock AlarmManager mAlarmManager;
     @Mock
     private NetworkStatsSubscriptionsMonitor mNetworkStatsSubscriptionsMonitor;
@@ -314,6 +310,8 @@
         verify(mTetheringManager).registerTetheringEventCallback(
                 any(), tetheringEventCbCaptor.capture());
         mTetheringEventCallback = tetheringEventCbCaptor.getValue();
+
+        mUsageCallback = new TestableUsageCallback(mUsageCallbackBinder);
     }
 
     @NonNull
@@ -1247,20 +1245,14 @@
         DataUsageRequest inputRequest = new DataUsageRequest(
                 DataUsageRequest.REQUEST_ID_UNSET, sTemplateWifi, thresholdInBytes);
 
-        // Create a messenger that waits for callback activity
-        ConditionVariable cv = new ConditionVariable(false);
-        LatchedHandler latchedHandler = new LatchedHandler(Looper.getMainLooper(), cv);
-        Messenger messenger = new Messenger(latchedHandler);
-
         // Force poll
         expectDefaultSettings();
         expectNetworkStatsSummary(buildEmptyStats());
         expectNetworkStatsUidDetail(buildEmptyStats());
 
         // Register and verify request and that binder was called
-        DataUsageRequest request =
-                mService.registerUsageCallback(mServiceContext.getOpPackageName(), inputRequest,
-                        messenger, mBinder);
+        DataUsageRequest request = mService.registerUsageCallback(
+                mServiceContext.getOpPackageName(), inputRequest, mUsageCallback);
         assertTrue(request.requestId > 0);
         assertTrue(Objects.equals(sTemplateWifi, request.template));
         long minThresholdInBytes = 2 * 1024 * 1024; // 2 MB
@@ -1269,7 +1261,7 @@
         HandlerUtils.waitForIdle(mHandlerThread, WAIT_TIMEOUT);
 
         // Make sure that the caller binder gets connected
-        verify(mBinder).linkToDeath(any(IBinder.DeathRecipient.class), anyInt());
+        verify(mUsageCallbackBinder).linkToDeath(any(IBinder.DeathRecipient.class), anyInt());
 
         // modify some number on wifi, and trigger poll event
         // not enough traffic to call data usage callback
@@ -1284,7 +1276,7 @@
         assertNetworkTotal(sTemplateWifi, 1024L, 1L, 2048L, 2L, 0);
 
         // make sure callback has not being called
-        assertEquals(INVALID_TYPE, latchedHandler.lastMessageType);
+        mUsageCallback.assertNoCallback();
 
         // and bump forward again, with counters going higher. this is
         // important, since it will trigger the data usage callback
@@ -1299,23 +1291,21 @@
         assertNetworkTotal(sTemplateWifi, 4096000L, 4L, 8192000L, 8L, 0);
 
 
-        // Wait for the caller to ack receipt of CALLBACK_LIMIT_REACHED
-        assertTrue(cv.block(WAIT_TIMEOUT));
-        assertEquals(NetworkStatsManager.CALLBACK_LIMIT_REACHED, latchedHandler.lastMessageType);
-        cv.close();
+        // Wait for the caller to invoke expectOnThresholdReached.
+        mUsageCallback.expectOnThresholdReached(request);
 
         // Allow binder to disconnect
-        when(mBinder.unlinkToDeath(any(IBinder.DeathRecipient.class), anyInt())).thenReturn(true);
+        when(mUsageCallbackBinder.unlinkToDeath(any(IBinder.DeathRecipient.class), anyInt()))
+                .thenReturn(true);
 
         // Unregister request
         mService.unregisterUsageRequest(request);
 
-        // Wait for the caller to ack receipt of CALLBACK_RELEASED
-        assertTrue(cv.block(WAIT_TIMEOUT));
-        assertEquals(NetworkStatsManager.CALLBACK_RELEASED, latchedHandler.lastMessageType);
+        // Wait for the caller to invoke expectOnCallbackReleased.
+        mUsageCallback.expectOnCallbackReleased(request);
 
         // Make sure that the caller binder gets disconnected
-        verify(mBinder).unlinkToDeath(any(IBinder.DeathRecipient.class), anyInt());
+        verify(mUsageCallbackBinder).unlinkToDeath(any(IBinder.DeathRecipient.class), anyInt());
     }
 
     @Test
@@ -1891,21 +1881,4 @@
     private void waitForIdle() {
         HandlerUtils.waitForIdle(mHandlerThread, WAIT_TIMEOUT);
     }
-
-    static class LatchedHandler extends Handler {
-        private final ConditionVariable mCv;
-        int lastMessageType = INVALID_TYPE;
-
-        LatchedHandler(Looper looper, ConditionVariable cv) {
-            super(looper);
-            mCv = cv;
-        }
-
-        @Override
-        public void handleMessage(Message msg) {
-            lastMessageType = msg.what;
-            mCv.open();
-            super.handleMessage(msg);
-        }
-    }
 }
diff --git a/tests/unit/java/com/android/server/net/TestableUsageCallback.kt b/tests/unit/java/com/android/server/net/TestableUsageCallback.kt
new file mode 100644
index 0000000..1917ec3
--- /dev/null
+++ b/tests/unit/java/com/android/server/net/TestableUsageCallback.kt
@@ -0,0 +1,79 @@
+/*
+ * 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.
+ */
+
+package com.android.server.net
+
+import android.net.DataUsageRequest
+import android.net.netstats.IUsageCallback
+import android.os.IBinder
+import java.util.concurrent.LinkedBlockingQueue
+import java.util.concurrent.TimeUnit
+import kotlin.test.fail
+
+private const val DEFAULT_TIMEOUT_MS = 200L
+
+// TODO: Move the class to static libs once all downstream have IUsageCallback definition.
+class TestableUsageCallback(private val binder: IBinder) : IUsageCallback.Stub() {
+    sealed class CallbackType(val request: DataUsageRequest) {
+        class OnThresholdReached(request: DataUsageRequest) : CallbackType(request)
+        class OnCallbackReleased(request: DataUsageRequest) : CallbackType(request)
+    }
+
+    // TODO: Change to use ArrayTrackRecord once moved into to the module.
+    private val history = LinkedBlockingQueue<CallbackType>()
+
+    override fun onThresholdReached(request: DataUsageRequest) {
+        history.add(CallbackType.OnThresholdReached(request))
+    }
+
+    override fun onCallbackReleased(request: DataUsageRequest) {
+        history.add(CallbackType.OnCallbackReleased(request))
+    }
+
+    fun expectOnThresholdReached(request: DataUsageRequest) {
+        expectCallback<CallbackType.OnThresholdReached>(request, DEFAULT_TIMEOUT_MS)
+    }
+
+    fun expectOnCallbackReleased(request: DataUsageRequest) {
+        expectCallback<CallbackType.OnCallbackReleased>(request, DEFAULT_TIMEOUT_MS)
+    }
+
+    @JvmOverloads
+    fun assertNoCallback(timeout: Long = DEFAULT_TIMEOUT_MS) {
+        val cb = history.poll(timeout, TimeUnit.MILLISECONDS)
+        cb?.let { fail("Expected no callback but got $cb") }
+    }
+
+    // Expects a callback of the specified request on the specified network within the timeout.
+    // If no callback arrives, or a different callback arrives, fail.
+    private inline fun <reified T : CallbackType> expectCallback(
+        expectedRequest: DataUsageRequest,
+        timeoutMs: Long
+    ) {
+        history.poll(timeoutMs, TimeUnit.MILLISECONDS).let {
+            if (it !is T || it.request != expectedRequest) {
+                fail("Unexpected callback : $it," +
+                        " expected ${T::class} with Request[$expectedRequest]")
+            } else {
+                it
+            }
+        }
+    }
+
+    override fun asBinder(): IBinder {
+        return binder
+    }
+}
\ No newline at end of file