Support multiple cleanup blocks in tryTest{}

This is useful for some tests that need to make sure a few
things all happen, but some of them may throw an exception.

In vanilla Java you'd say

try {
  ...
} finally {
  try {
    ...
  } finally {
    try {
      ...
    } finally {
      ...
    }
  }
}

With this patch, you can pass a list of blocks to clean.
Kotlin :

tryTest {
  ...
} cleanupStep {
  ...
} cleanupStep {
  ...
} cleanup {
  ...
}

Java :
tryAndCleanup(() -> {
  ...
  }, () -> {
  ...
  }, () -> {
  ...
  })

This keeps the semantics of tryTest{} of throwing any
exception that was thrown in tryTest{} and adding the
exception in the cleanup steps as suppressed.

Test: new tests for this
Change-Id: Ie26b802cb4893e13a70646aa0b7887701dee1ec6
diff --git a/staticlibs/tests/unit/src/com/android/net/module/util/CleanupTest.kt b/staticlibs/tests/unit/src/com/android/net/module/util/CleanupTest.kt
index 0067931..649b30e 100644
--- a/staticlibs/tests/unit/src/com/android/net/module/util/CleanupTest.kt
+++ b/staticlibs/tests/unit/src/com/android/net/module/util/CleanupTest.kt
@@ -32,6 +32,7 @@
 class CleanupTest {
     class TestException1 : Exception()
     class TestException2 : Exception()
+    class TestException3 : Exception()
 
     @Test
     fun testNotThrow() {
@@ -172,4 +173,32 @@
         assertTrue(x == 4)
         assertTrue(thrown.suppressedExceptions.isEmpty())
     }
+
+    @Test
+    fun testMultipleCleanups() {
+        var x = 1
+        val thrown = assertFailsWith<TestException1> {
+            tryTest {
+                x = 2
+                throw TestException1()
+            } cleanupStep {
+                assertTrue(x == 2)
+                x = 3
+                throw TestException2()
+                x = 4
+            } cleanupStep {
+                assertTrue(x == 3)
+                x = 5
+                throw TestException3()
+                x = 6
+            } cleanup {
+                assertTrue(x == 5)
+                x = 7
+            }
+        }
+        assertEquals(2, thrown.suppressedExceptions.size)
+        assertTrue(thrown.suppressedExceptions[0] is TestException2)
+        assertTrue(thrown.suppressedExceptions[1] is TestException3)
+        assert(x == 7)
+    }
 }
diff --git a/staticlibs/tests/unit/src/com/android/net/module/util/CleanupTestJava.java b/staticlibs/tests/unit/src/com/android/net/module/util/CleanupTestJava.java
index 83abfa1..8a13397 100644
--- a/staticlibs/tests/unit/src/com/android/net/module/util/CleanupTestJava.java
+++ b/staticlibs/tests/unit/src/com/android/net/module/util/CleanupTestJava.java
@@ -20,6 +20,7 @@
 import static com.android.testutils.MiscAsserts.assertThrows;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 import android.util.Log;
 
@@ -31,6 +32,7 @@
     private static final String TAG = CleanupTestJava.class.getSimpleName();
     private static final class TestException1 extends Exception {}
     private static final class TestException2 extends Exception {}
+    private static final class TestException3 extends Exception {}
 
     @Test
     public void testNotThrow() {
@@ -93,4 +95,27 @@
         );
         assertEquals(3, x.get());
     }
+
+    @Test
+    public void testMultipleCleanups() {
+        final AtomicInteger x = new AtomicInteger(1);
+        final TestException1 exception = assertThrows(TestException1.class, () ->
+                testAndCleanup(() -> {
+                    x.compareAndSet(1, 2);
+                    throw new TestException1();
+                }, () -> {
+                        x.compareAndSet(2, 3);
+                        throw new TestException2();
+                    }, () -> {
+                        x.compareAndSet(3, 4);
+                        throw new TestException3();
+                    }, () -> {
+                        x.compareAndSet(4, 5);
+                    })
+        );
+        assertEquals(2, exception.getSuppressed().length);
+        assertTrue(exception.getSuppressed()[0] instanceof TestException2);
+        assertTrue(exception.getSuppressed()[1] instanceof TestException3);
+        assertEquals(5, x.get());
+    }
 }
diff --git a/staticlibs/testutils/hostdevice/com/android/testutils/Cleanup.kt b/staticlibs/testutils/hostdevice/com/android/testutils/Cleanup.kt
index 1b67f68..45783d8 100644
--- a/staticlibs/testutils/hostdevice/com/android/testutils/Cleanup.kt
+++ b/staticlibs/testutils/hostdevice/com/android/testutils/Cleanup.kt
@@ -22,14 +22,6 @@
 import com.android.testutils.ExceptionUtils.ThrowingSupplier
 import javax.annotation.CheckReturnValue
 
-@CheckReturnValue
-fun <T> tryTest(block: () -> T) = TryExpr(
-        try {
-            Result.success(block())
-        } catch (e: Throwable) {
-            Result.failure(e)
-        })
-
 /**
  * Utility to do cleanup in tests without replacing exceptions with those from a finally block.
  *
@@ -54,11 +46,15 @@
  * to the standard try{}finally{}, if both throws, the construct throws the exception that happened
  * in tryTest{} rather than the one that happened in cleanup{}.
  *
- * Kotlin usage is as try{}finally{} :
+ * Kotlin usage is as try{}finally{}, but with multiple finally{} blocks :
  * tryTest {
  *   testing code
+ * } cleanupStep {
+ *   cleanup code 1
+ * } cleanupStep {
+ *   cleanup code 2
  * } cleanup {
- *   cleanup code
+ *   cleanup code 3
  * }
  * Catch blocks can be added with the following syntax :
  * tryTest {
@@ -67,14 +63,24 @@
  *   do something to it
  * }
  *
- * Java doesn't allow this kind of syntax, so instead a function taking 2 lambdas is provided.
+ * Java doesn't allow this kind of syntax, so instead a function taking lambdas is provided.
  * testAndCleanup(() -> {
  *   testing code
  * }, () -> {
- *   cleanup code
+ *   cleanup code 1
+ * }, () -> {
+ *   cleanup code 2
  * });
  */
 
+@CheckReturnValue
+fun <T> tryTest(block: () -> T) = TryExpr(
+        try {
+            Result.success(block())
+        } catch (e: Throwable) {
+            Result.failure(e)
+        })
+
 // Some downstream branches have an older kotlin that doesn't know about value classes.
 // TODO : Change this to "value class" when aosp no longer merges into such branches.
 @Suppress("INLINE_CLASS_DEPRECATED")
@@ -89,30 +95,31 @@
         })
     }
 
-    inline infix fun cleanup(block: () -> Unit): T {
+    @CheckReturnValue
+    inline infix fun cleanupStep(block: () -> Unit): TryExpr<T> {
         try {
             block()
         } catch (e: Throwable) {
             val originalException = result.exceptionOrNull()
-            if (null == originalException) {
-                throw e
+            return TryExpr(if (null == originalException) {
+                Result.failure(e)
             } else {
                 originalException.addSuppressed(e)
-                throw originalException
-            }
+                Result.failure(originalException)
+            })
         }
-        return result.getOrThrow()
+        return this
     }
+
+    inline infix fun cleanup(block: () -> Unit): T = cleanupStep(block).result.getOrThrow()
 }
 
 // Java support
-fun <T> testAndCleanup(tryBlock: ThrowingSupplier<T>, cleanupBlock: ThrowingRunnable): T {
-    return tryTest {
-        tryBlock.get()
-    } cleanup {
-        cleanupBlock.run()
-    }
+fun <T> testAndCleanup(tryBlock: ThrowingSupplier<T>, vararg cleanupBlock: ThrowingRunnable): T {
+    return cleanupBlock.fold(tryTest { tryBlock.get() }) { previousExpr, nextCleanup ->
+        previousExpr.cleanupStep { nextCleanup.run() }
+    }.cleanup {}
 }
-fun testAndCleanup(tryBlock: ThrowingRunnable, cleanupBlock: ThrowingRunnable) {
-    return testAndCleanup(ThrowingSupplier { tryBlock.run() }, cleanupBlock)
+fun testAndCleanup(tryBlock: ThrowingRunnable, vararg cleanupBlock: ThrowingRunnable) {
+    return testAndCleanup(ThrowingSupplier { tryBlock.run() }, *cleanupBlock)
 }