Snap for 7478986 from 994231a768545a76dc3afbfdf8dcbc28657796dc to sc-release

Change-Id: I5664293f116339fac944765707dcb236c826fbbd

2 files changed, 83 insertions, 2 deletions

diff --git a/common/testutils/devicetests/com/android/testutils/TestableNetworkOfferCallback.kt b/common/testutils/devicetests/com/android/testutils/TestableNetworkOfferCallback.kt
new file mode 100644
index 00000000..21bd60c9
--- /dev/null
+++ b/common/testutils/devicetests/com/android/testutils/TestableNetworkOfferCallback.kt
@@ -0,0 +1,75 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.testutils
+
+import android.net.NetworkCapabilities
+import android.net.NetworkProvider
+import android.net.NetworkRequest
+import android.util.Log
+import com.android.net.module.util.ArrayTrackRecord
+import kotlin.test.fail
+
+class TestableNetworkOfferCallback(val timeoutMs: Long, private val noCallbackTimeoutMs: Long)
+            : NetworkProvider.NetworkOfferCallback {
+    private val TAG = this::class.simpleName
+    val history = ArrayTrackRecord<CallbackEntry>().newReadHead()
+
+    sealed class CallbackEntry {
+        data class OnNetworkNeeded(val request: NetworkRequest) : CallbackEntry()
+        data class OnNetworkUnneeded(val request: NetworkRequest) : CallbackEntry()
+    }
+
+    /**
+     * Called by the system when a network for this offer is needed to satisfy some
+     * networking request.
+     */
+    override fun onNetworkNeeded(request: NetworkRequest) {
+        Log.d(TAG, "onNetworkNeeded $request")
+        history.add(CallbackEntry.OnNetworkNeeded(request))
+    }
+
+    /**
+     * Called by the system when this offer is no longer valuable for this request.
+     */
+    override fun onNetworkUnneeded(request: NetworkRequest) {
+        Log.d(TAG, "onNetworkUnneeded $request")
+        history.add(CallbackEntry.OnNetworkUnneeded(request))
+    }
+
+    inline fun <reified T : CallbackEntry> expectCallbackThat(
+        crossinline predicate: (T) -> Boolean
+    ) {
+        val event = history.poll(timeoutMs)
+                ?: fail("Did not receive callback after ${timeoutMs}ms")
+        if (event !is T || !predicate(event)) fail("Received unexpected callback $event")
+    }
+
+    fun expectOnNetworkNeeded(capabilities: NetworkCapabilities) =
+            expectCallbackThat<CallbackEntry.OnNetworkNeeded> {
+                it.request.canBeSatisfiedBy(capabilities)
+            }
+
+    fun expectOnNetworkUnneeded(capabilities: NetworkCapabilities) =
+            expectCallbackThat<CallbackEntry.OnNetworkUnneeded> {
+                it.request.canBeSatisfiedBy(capabilities)
+            }
+
+    fun assertNoCallback() {
+        val cb = history.poll(noCallbackTimeoutMs)
+        if (null != cb) fail("Expected no callback but got $cb")
+    }
+}
\ No newline at end of file
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) {