diff options
| -rw-r--r-- | common/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt | 10 |
1 files changed, 8 insertions, 2 deletions
diff --git a/common/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt b/common/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt index d5c3a2a0..d034a7d5 100644 --- a/common/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt +++ b/common/testutils/devicetests/com/android/testutils/TestableNetworkStatsProvider.kt @@ -23,6 +23,7 @@ import kotlin.test.assertTrue 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 @@ open class TestableNetworkStatsProvider( 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) { |