Merge "Revert "Revert "Allow expect onRequestStatsUpdated with any token"""
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
index d5c3a2a..d034a7d 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
@@ -23,6 +23,7 @@
import kotlin.test.fail
private const val DEFAULT_TIMEOUT_MS = 200L
+const val TOKEN_ANY = -1
open class TestableNetworkStatsProvider(
val defaultTimeoutMs: Long = DEFAULT_TIMEOUT_MS
@@ -49,8 +50,13 @@
history.add(CallbackType.OnSetAlert(quotaBytes))
}
- fun expectOnRequestStatsUpdate(token: Int, timeout: Long = defaultTimeoutMs) {
- assertEquals(CallbackType.OnRequestStatsUpdate(token), history.poll(timeout))
+ fun expectOnRequestStatsUpdate(token: Int, timeout: Long = defaultTimeoutMs): Int {
+ val event = history.poll(timeout)
+ assertTrue(event is CallbackType.OnRequestStatsUpdate)
+ if (token != TOKEN_ANY) {
+ assertEquals(token, event.token)
+ }
+ return token
}
fun expectOnSetLimit(iface: String?, quotaBytes: Long, timeout: Long = defaultTimeoutMs) {