Allow expect onRequestStatsUpdated with any token

Altough token field is not used for now. But by design,
the token will be the version number of networkIdentity mappings
that maintain by the service, which cannot be predicted by the
test. Thus, allow expecting any token is necessary.

Ignore-AOSP-First: Counter part CL is not available in aosp.
Test: m gts && atest \
      GtsNetworkStackHostTestCases:NetworkStatsHostTest#testNetworkStatsManager
Bug: 191327585
Change-Id: Ibf10040a5c10b7b1c7d381c01a071fb157409004
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) {