Check connectivity is restored after cleanup

Only check that internet connectivity is restored after test validation
URLs have been cleaned up, and other runAfterNextCleanup actions have
been run (such as reconnecting wifi).

The test currently fails on devices that only have one upstream network
(no mobile data for example); on devices that have mobile+wifi, mobile
data is still available so the test passes, but wifi is not actually
usable at that point.

Bug: 238844861
Test: atest ConnectivityManagerTest
Change-Id: Ib29e44245aa4e8c8da8994f9b8f1cd31ab48b939
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index 64238b3..834f11b 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -412,14 +412,17 @@
 
         // All tests in this class require a working Internet connection as they start. Make
         // sure there is still one as they end that's ready to use for the next test to use.
-        final TestNetworkCallback callback = new TestNetworkCallback();
-        registerDefaultNetworkCallback(callback);
-        try {
-            assertNotNull("Couldn't restore Internet connectivity", callback.waitForAvailable());
-        } finally {
-            // Unregister all registered callbacks.
-            unregisterRegisteredCallbacks();
-        }
+        mTestValidationConfigRule.runAfterNextCleanup(() -> {
+            final TestNetworkCallback callback = new TestNetworkCallback();
+            registerDefaultNetworkCallback(callback);
+            try {
+                assertNotNull("Couldn't restore Internet connectivity",
+                        callback.waitForAvailable());
+            } finally {
+                // Unregister all registered callbacks.
+                unregisterRegisteredCallbacks();
+            }
+        });
     }
 
     @Test
diff --git a/tests/cts/net/src/android/net/cts/DeviceConfigRule.kt b/tests/cts/net/src/android/net/cts/DeviceConfigRule.kt
index d31a4e0..3a739f2 100644
--- a/tests/cts/net/src/android/net/cts/DeviceConfigRule.kt
+++ b/tests/cts/net/src/android/net/cts/DeviceConfigRule.kt
@@ -21,6 +21,7 @@
 import android.provider.DeviceConfig
 import android.util.Log
 import com.android.modules.utils.build.SdkLevel
+import com.android.testutils.ExceptionUtils.ThrowingRunnable
 import com.android.testutils.runAsShell
 import com.android.testutils.tryTest
 import org.junit.rules.TestRule
@@ -51,7 +52,7 @@
     /**
      * Actions to be run after cleanup of the config, for the current test only.
      */
-    private val currentTestCleanupActions = mutableListOf<Runnable>()
+    private val currentTestCleanupActions = mutableListOf<ThrowingRunnable>()
 
     override fun apply(base: Statement, description: Description): Statement {
         return TestValidationUrlStatement(base, description)
@@ -93,8 +94,13 @@
                     originalConfig.clear()
                     usedConfig.clear()
                 } cleanup {
-                    currentTestCleanupActions.forEach { it.run() }
-                    currentTestCleanupActions.clear()
+                    // Fold all cleanup actions into cleanup steps of an empty tryTest, so they are
+                    // all run even if exceptions are thrown, and exceptions are reported properly.
+                    currentTestCleanupActions.fold(tryTest { }) {
+                        tryBlock, action -> tryBlock.cleanupStep { action.run() }
+                    }.cleanup {
+                        currentTestCleanupActions.clear()
+                    }
                 }
             }
         }
@@ -118,7 +124,7 @@
     /**
      * Add an action to be run after config cleanup when the current test case ends.
      */
-    fun runAfterNextCleanup(action: Runnable) {
+    fun runAfterNextCleanup(action: ThrowingRunnable) {
         currentTestCleanupActions.add(action)
     }
 }