Check that wifi, cell validates before tests

Tests are sometimes run on hardware devices with wrongly configured wifi
or cell data. Ensure that this is reported as an infra error, and not a
test error, so the root cause is easier to identify.

Bug: 264170054
Test: atest
Change-Id: I4f964fbd4ee497e8ac92f7729375b75b6c4594a3
diff --git a/staticlibs/testutils/app/connectivitychecker/src/com/android/testutils/connectivitypreparer/ConnectivityCheckTest.kt b/staticlibs/testutils/app/connectivitychecker/src/com/android/testutils/connectivitypreparer/ConnectivityCheckTest.kt
index f1f0975..d75d9ca 100644
--- a/staticlibs/testutils/app/connectivitychecker/src/com/android/testutils/connectivitypreparer/ConnectivityCheckTest.kt
+++ b/staticlibs/testutils/app/connectivitychecker/src/com/android/testutils/connectivitypreparer/ConnectivityCheckTest.kt
@@ -18,17 +18,10 @@
 
 import android.content.pm.PackageManager.FEATURE_TELEPHONY
 import android.content.pm.PackageManager.FEATURE_WIFI
-import android.net.ConnectivityManager
-import android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET
-import android.net.NetworkCapabilities.TRANSPORT_CELLULAR
-import android.net.NetworkRequest
 import android.telephony.TelephonyManager
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.testutils.ConnectUtil
-import com.android.testutils.RecorderCallback
-import com.android.testutils.TestableNetworkCallback
-import com.android.testutils.tryTest
 import kotlin.test.assertTrue
 import kotlin.test.fail
 import org.junit.Test
@@ -36,8 +29,9 @@
 
 @RunWith(AndroidJUnit4::class)
 class ConnectivityCheckTest {
-    val context by lazy { InstrumentationRegistry.getInstrumentation().context }
-    val pm by lazy { context.packageManager }
+    private val context by lazy { InstrumentationRegistry.getInstrumentation().context }
+    private val pm by lazy { context.packageManager }
+    private val connectUtil by lazy { ConnectUtil(context) }
 
     @Test
     fun testCheckConnectivity() {
@@ -47,7 +41,7 @@
 
     private fun checkWifiSetup() {
         if (!pm.hasSystemFeature(FEATURE_WIFI)) return
-        ConnectUtil(context).ensureWifiConnected()
+        connectUtil.ensureWifiValidated()
     }
 
     private fun checkTelephonySetup() {
@@ -69,20 +63,6 @@
         assertTrue(tm.isDataConnectivityPossible,
             "The device is not setup with a SIM card that supports data connectivity. " +
                     commonError)
-        val cb = TestableNetworkCallback()
-        val cm = context.getSystemService(ConnectivityManager::class.java)
-                ?: fail("Could not get ConnectivityManager")
-        cm.requestNetwork(
-                NetworkRequest.Builder()
-                        .addTransportType(TRANSPORT_CELLULAR)
-                        .addCapability(NET_CAPABILITY_INTERNET).build(), cb)
-        tryTest {
-            cb.poll { it is RecorderCallback.CallbackEntry.Available }
-                    ?: fail("The device does not have mobile data available. Check that it is " +
-                            "setup with a SIM card that has a working data plan, and that the " +
-                            "APN configuration is valid.")
-        } cleanup {
-            cm.unregisterNetworkCallback(cb)
-        }
+        connectUtil.ensureCellularValidated()
     }
 }
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/ConnectUtil.kt b/staticlibs/testutils/devicetests/com/android/testutils/ConnectUtil.kt
index 71f7877..b1d64f8 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/ConnectUtil.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/ConnectUtil.kt
@@ -23,6 +23,9 @@
 import android.content.IntentFilter
 import android.net.ConnectivityManager
 import android.net.Network
+import android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET
+import android.net.NetworkCapabilities.NET_CAPABILITY_VALIDATED
+import android.net.NetworkCapabilities.TRANSPORT_CELLULAR
 import android.net.NetworkCapabilities.TRANSPORT_WIFI
 import android.net.NetworkRequest
 import android.net.wifi.ScanResult
@@ -33,6 +36,7 @@
 import android.util.Log
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.testutils.RecorderCallback.CallbackEntry
+import com.android.testutils.RecorderCallback.CallbackEntry.CapabilitiesChanged
 import java.util.concurrent.CompletableFuture
 import java.util.concurrent.TimeUnit
 import kotlin.test.assertNotNull
@@ -56,13 +60,35 @@
     private val wifiManager = context.getSystemService(WifiManager::class.java)
             ?: fail("Could not find WifiManager")
 
-    fun ensureWifiConnected(): Network {
-        val callback = TestableNetworkCallback()
+    fun ensureWifiConnected(): Network = ensureWifiConnected(requireValidated = false)
+    fun ensureWifiValidated(): Network = ensureWifiConnected(requireValidated = true)
+
+    fun ensureCellularValidated(): Network {
+        val cb = TestableNetworkCallback()
+        cm.requestNetwork(
+            NetworkRequest.Builder()
+                .addTransportType(TRANSPORT_CELLULAR)
+                .addCapability(NET_CAPABILITY_INTERNET).build(), cb)
+        return tryTest {
+            val errorMsg = "The device does not have mobile data available. Check that it is " +
+                    "setup with a SIM card that has a working data plan, that the APN " +
+                    "configuration is valid, and that the device can access the internet through " +
+                    "mobile data."
+            cb.eventuallyExpect<CapabilitiesChanged>(errorMsg) {
+                it.caps.hasCapability(NET_CAPABILITY_VALIDATED)
+            }.network
+        } cleanup {
+            cm.unregisterNetworkCallback(cb)
+        }
+    }
+
+    private fun ensureWifiConnected(requireValidated: Boolean): Network {
+        val callback = TestableNetworkCallback(timeoutMs = WIFI_CONNECT_TIMEOUT_MS)
         cm.registerNetworkCallback(NetworkRequest.Builder()
                 .addTransportType(TRANSPORT_WIFI)
                 .build(), callback)
 
-        try {
+        return tryTest {
             val connInfo = wifiManager.connectionInfo
             Log.d(TAG, "connInfo=" + connInfo)
             if (connInfo == null || connInfo.networkId == -1) {
@@ -73,12 +99,19 @@
                 val config = getOrCreateWifiConfiguration()
                 connectToWifiConfig(config)
             }
-            val cb = callback.poll(WIFI_CONNECT_TIMEOUT_MS) { it is CallbackEntry.Available }
-            assertNotNull(cb, "Could not connect to a wifi access point within " +
-                    "$WIFI_CONNECT_TIMEOUT_MS ms. Check that the test device has a wifi network " +
-                    "configured, and that the test access point is functioning properly.")
-            return cb.network
-        } finally {
+            val errorMsg = if (requireValidated) {
+                "The wifi access point did not have access to the internet after " +
+                        "$WIFI_CONNECT_TIMEOUT_MS ms. Check that it has a working connection."
+            } else {
+                "Could not connect to a wifi access point within $WIFI_CONNECT_TIMEOUT_MS ms. " +
+                        "Check that the test device has a wifi network configured, and that the " +
+                        "test access point is functioning properly."
+            }
+            val cb = callback.eventuallyExpect<CapabilitiesChanged>(errorMsg) {
+                (!requireValidated || it.caps.hasCapability(NET_CAPABILITY_VALIDATED))
+            }
+            cb.network
+        } cleanup {
             cm.unregisterNetworkCallback(callback)
         }
     }
@@ -201,3 +234,10 @@
         }
     }
 }
+
+private inline fun <reified T : CallbackEntry> TestableNetworkCallback.eventuallyExpect(
+    errorMsg: String,
+    crossinline predicate: (T) -> Boolean = { true }
+): T = history.poll(defaultTimeoutMs, mark) { it is T && predicate(it) }.also {
+    assertNotNull(it, errorMsg)
+} as T