commit a7ace49440b4a0168fa9bc019a231198d9f35fb1
author Junyu Lai <junyulai@google.com> Tue Jun 22 01:21:24 2021 +0000
committer Junyu Lai <junyulai@google.com> Thu Jun 24 01:19:25 2021 +0000
tree b2bc2af7b12a71542f60daa7d76fbf0f03742727
parent 8dc6e175cc02b8ca6455c4eff7a35f613367471e

Add OnSetWarningAndLimit into TestableNetworkStatsProvider

This change also address left comments at ag/15010113.

Test: m gts && atest \
      GtsNetworkStackHostTestCases:NetworkStatsHostTest#testNetworkStatsProvider
Bug: 191327585
Merged-In: I100668cc6d4cecbaff29d027d37d313dd48854f3
Change-Id: I100668cc6d4cecbaff29d027d37d313dd48854f3
  (cherry-picked from ag/15061296)

staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt

diff --git a/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
index d034a7d..8477151 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
@@ -30,7 +30,12 @@
 ) : NetworkStatsProvider() {
     sealed class CallbackType {
         data class OnRequestStatsUpdate(val token: Int) : CallbackType()
-        data class OnSetLimit(val iface: String?, val quotaBytes: Long) : CallbackType()
+        data class OnSetWarningAndLimit(
+            val iface: String?,
+            val warningBytes: Long,
+            val limitBytes: Long
+        ) : CallbackType()
+        data class OnSetLimit(val iface: String?, val limitBytes: Long) : CallbackType()
         data class OnSetAlert(val quotaBytes: Long) : CallbackType()
     }
 
@@ -42,6 +47,10 @@
         history.add(CallbackType.OnRequestStatsUpdate(token))
     }
 
+    override fun onSetWarningAndLimit(iface: String, warningBytes: Long, limitBytes: Long) {
+        history.add(CallbackType.OnSetWarningAndLimit(iface, warningBytes, limitBytes))
+    }
+
     override fun onSetLimit(iface: String, quotaBytes: Long) {
         history.add(CallbackType.OnSetLimit(iface, quotaBytes))
     }
@@ -56,7 +65,7 @@
         if (token != TOKEN_ANY) {
             assertEquals(token, event.token)
         }
-        return token
+        return event.token
     }
 
     fun expectOnSetLimit(iface: String?, quotaBytes: Long, timeout: Long = defaultTimeoutMs) {