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 cc0f9be..0067931 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
@@ -21,8 +21,10 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
+import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
 import kotlin.test.assertTrue
+import kotlin.test.fail
 
 private val TAG = CleanupTest::class.simpleName
 
@@ -104,6 +106,70 @@
 
     @Test
     fun testReturn() {
+        val resultIfSuccess = 11
+        val resultIfException = 12
+        fun doTestReturn(crash: Boolean) = tryTest {
+            if (crash) throw RuntimeException() else resultIfSuccess
+        }.catch<RuntimeException> {
+            resultIfException
+        } cleanup {}
+
         assertTrue(6 == tryTest { 6 } cleanup { Log.e(TAG, "tested") })
+        assertEquals(resultIfSuccess, doTestReturn(crash = false))
+        assertEquals(resultIfException, doTestReturn(crash = true))
+    }
+
+    @Test
+    fun testCatch() {
+        var x = 1
+        tryTest {
+            x = 2
+            throw TestException1()
+            x = 3
+        }.catch<TestException1> {
+            x = 4
+        }.catch<TestException2> {
+            x = 5
+        } cleanup {
+            assertTrue(x == 4)
+            x = 6
+        }
+        assertTrue(x == 6)
+    }
+
+    @Test
+    fun testNotCatch() {
+        var x = 1
+        assertFailsWith<TestException1> {
+            tryTest {
+                x = 2
+                throw TestException1()
+            }.catch<TestException2> {
+                fail("Caught TestException2 instead of TestException1")
+            } cleanup {
+                assertTrue(x == 2)
+                x = 3
+            }
+        }
+        assertTrue(x == 3)
+    }
+
+    @Test
+    fun testThrowInCatch() {
+        var x = 1
+        val thrown = assertFailsWith<TestException2> {
+            tryTest {
+                x = 2
+                throw TestException1()
+            }.catch<TestException1> {
+                x = 3
+                throw TestException2()
+            } cleanup {
+                assertTrue(x == 3)
+                x = 4
+            }
+        }
+        assertTrue(x == 4)
+        assertTrue(thrown.suppressedExceptions.isEmpty())
     }
 }
diff --git a/staticlibs/testutils/hostdevice/com/android/testutils/Cleanup.kt b/staticlibs/testutils/hostdevice/com/android/testutils/Cleanup.kt
index d93c7d0..0f7ec70 100644
--- a/staticlibs/testutils/hostdevice/com/android/testutils/Cleanup.kt
+++ b/staticlibs/testutils/hostdevice/com/android/testutils/Cleanup.kt
@@ -52,6 +52,12 @@
  * } cleanup {
  *   cleanup code
  * }
+ * Catch blocks can be added with the following syntax :
+ * tryTest {
+ *   testing code
+ * }.catch<ExceptionType> { it ->
+ *   do something to it
+ * }
  *
  * Java doesn't allow this kind of syntax, so instead a function taking 2 lambdas is provided.
  * testAndCleanup(() -> {
@@ -60,10 +66,21 @@
  *   cleanup code
  * });
  */
+
 // sc-mainline-prod has an older kotlin that doesn't know about value classes. TODO : Change this
 // to "value class" when aosp no longer merges into sc-mainline-prod.
 @Suppress("INLINE_CLASS_DEPRECATED")
 inline class ExceptionCleanupBlock<T>(val result: Result<T>) {
+    inline infix fun <reified E : Throwable> catch(block: (E) -> T): ExceptionCleanupBlock<T> {
+        val originalException = result.exceptionOrNull()
+        if (originalException !is E) return this
+        return ExceptionCleanupBlock(try {
+            Result.success(block(originalException))
+        } catch (e: Exception) {
+            Result.failure(e)
+        })
+    }
+
     inline infix fun cleanup(block: () -> Unit): T {
         try {
             block()