commit d1573ed5e5313bb5c1b467d531d035d3cd3eba40
author Roy Want <roywant@google.com> Tue Jun 22 22:25:27 2021 +0000
committer Gerrit Code Review <noreply-gerritcodereview@google.com> Tue Jun 22 22:25:27 2021 +0000
tree 93a1fdc96403bf656132f82df0f9f0fa6d2e0959
parent c0786a5ec8be6f58559d63ca0df86c9d066f24e5
parent 75b9e6773bfa46eac305f03e75f6b397cb24b878

Merge "Revert "Allow expect onRequestStatsUpdated with any token""

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..d5c3a2a 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt
@@ -23,7 +23,6 @@
 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
@@ -50,13 +49,8 @@
         history.add(CallbackType.OnSetAlert(quotaBytes))
     }
 
-    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 expectOnRequestStatsUpdate(token: Int, timeout: Long = defaultTimeoutMs) {
+        assertEquals(CallbackType.OnRequestStatsUpdate(token), history.poll(timeout))
     }
 
     fun expectOnSetLimit(iface: String?, quotaBytes: Long, timeout: Long = defaultTimeoutMs) {