Merge "Check connectivity is restored after cleanup"
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index 5edf727..6ff2458 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -413,14 +413,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)
     }
 }