diff --git a/staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProviderBinder.kt b/staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProviderBinder.kt
index 02922d8..c19c9f9 100644
--- a/staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProviderBinder.kt
+++ b/staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProviderBinder.kt
@@ -26,8 +26,12 @@
 open class TestableNetworkStatsProviderBinder : INetworkStatsProvider.Stub() {
     sealed class CallbackType {
         data class OnRequestStatsUpdate(val token: Int) : CallbackType()
-        data class OnSetLimit(val iface: String?, val quotaBytes: Long) : CallbackType()
         data class OnSetAlert(val quotaBytes: Long) : CallbackType()
+        data class OnSetWarningAndLimit(
+            val iface: String,
+            val warningBytes: Long,
+            val limitBytes: Long
+        ) : CallbackType()
     }
 
     private val history = ArrayTrackRecord<CallbackType>().ReadHead()
@@ -36,20 +40,21 @@
         history.add(CallbackType.OnRequestStatsUpdate(token))
     }
 
-    override fun onSetLimit(iface: String?, quotaBytes: Long) {
-        history.add(CallbackType.OnSetLimit(iface, quotaBytes))
-    }
-
     override fun onSetAlert(quotaBytes: Long) {
         history.add(CallbackType.OnSetAlert(quotaBytes))
     }
 
+    override fun onSetWarningAndLimit(iface: String, warningBytes: Long, limitBytes: Long) {
+        history.add(CallbackType.OnSetWarningAndLimit(iface, warningBytes, limitBytes))
+    }
+
     fun expectOnRequestStatsUpdate(token: Int) {
         assertEquals(CallbackType.OnRequestStatsUpdate(token), history.poll(DEFAULT_TIMEOUT_MS))
     }
 
-    fun expectOnSetLimit(iface: String?, quotaBytes: Long) {
-        assertEquals(CallbackType.OnSetLimit(iface, quotaBytes), history.poll(DEFAULT_TIMEOUT_MS))
+    fun expectOnSetWarningAndLimit(iface: String, warningBytes: Long, limitBytes: Long) {
+        assertEquals(CallbackType.OnSetWarningAndLimit(iface, warningBytes, limitBytes),
+                history.poll(DEFAULT_TIMEOUT_MS))
     }
 
     fun expectOnSetAlert(quotaBytes: Long) {
diff --git a/staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProviderCbBinder.kt b/staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProviderCbBinder.kt
index 163473a..f15fd3b 100644
--- a/staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProviderCbBinder.kt
+++ b/staticlibs/devicetests/com/android/testutils/TestableNetworkStatsProviderCbBinder.kt
@@ -32,7 +32,7 @@
             val ifaceStats: NetworkStats,
             val uidStats: NetworkStats
         ) : CallbackType()
-        object NotifyLimitReached : CallbackType()
+        object NotifyWarningOrLimitReached : CallbackType()
         object NotifyAlertReached : CallbackType()
         object Unregister : CallbackType()
     }
@@ -43,8 +43,8 @@
         history.add(CallbackType.NotifyStatsUpdated(token, ifaceStats, uidStats))
     }
 
-    override fun notifyLimitReached() {
-        history.add(CallbackType.NotifyLimitReached)
+    override fun notifyWarningOrLimitReached() {
+        history.add(CallbackType.NotifyWarningOrLimitReached)
     }
 
     override fun notifyAlertReached() {
@@ -70,8 +70,8 @@
         assertNetworkStatsEquals(uidStats, event.uidStats)
     }
 
-    fun expectNotifyLimitReached() =
-            assertEquals(CallbackType.NotifyLimitReached, history.poll(DEFAULT_TIMEOUT_MS))
+    fun expectNotifyWarningOrLimitReached() =
+            assertEquals(CallbackType.NotifyWarningOrLimitReached, history.poll(DEFAULT_TIMEOUT_MS))
 
     fun expectNotifyAlertReached() =
             assertEquals(CallbackType.NotifyAlertReached, history.poll(DEFAULT_TIMEOUT_MS))