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

This reverts commit 886c3f16352f372f8e384c2dafd74c304008af38.

Reason for revert: From result of forrest, it turns out the culprit is not the CL. Hence the revert is not necessary.
See https://buganizer.corp.google.com/issues/191742804#comment10 for more detail.

Change-Id: If8492233c69d7947300f35f764199a7db1bd0045
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) {