Only report tryTest errors when not caught

Instead of reporting tryTest errors when the tryTest block throws,
report them in the cleanup block if no catch block has caught it.

This fixes a bug where tests that use tryTest/catch would collect
diagnostics in expected cases.

Test: atest
Change-Id: Ic53ae54f1945a2784becd50e5890fad8e7461474
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 851d09a..bde55c3 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
@@ -17,14 +17,19 @@
 package com.android.net.module.util
 
 import android.util.Log
+import com.android.testutils.TryTestConfig
 import com.android.testutils.tryTest
+import java.util.function.Consumer
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+import kotlin.test.assertNull
+import kotlin.test.assertTrue
+import kotlin.test.fail
+import org.junit.After
+import org.junit.Before
 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
 
@@ -34,6 +39,18 @@
     class TestException2 : Exception()
     class TestException3 : Exception()
 
+    private var originalDiagnosticsCollector: Consumer<Throwable>? = null
+
+    @Before
+    fun setUp() {
+        originalDiagnosticsCollector = TryTestConfig.swapDiagnosticsCollector(null)
+    }
+
+    @After
+    fun tearDown() {
+        TryTestConfig.swapDiagnosticsCollector(originalDiagnosticsCollector)
+    }
+
     @Test
     fun testNotThrow() {
         var x = 1
@@ -220,4 +237,74 @@
         assertTrue(thrown.suppressedExceptions[1] is TestException3)
         assert(x == 7)
     }
+
+    @Test
+    fun testNoErrorReportingWhenCaught() {
+        var error: Throwable? = null
+        TryTestConfig.swapDiagnosticsCollector {
+            error = it
+        }
+        var x = 1
+        tryTest {
+            x = 2
+            throw TestException1()
+            x = 3
+        }.catch<TestException1> {
+            x = 4
+        } cleanup {
+            x = 5
+        }
+
+        assertEquals(5, x)
+        assertNull(error)
+    }
+
+    @Test
+    fun testErrorReportingInTry() {
+        var error: Throwable? = null
+        TryTestConfig.swapDiagnosticsCollector {
+            assertNull(error)
+            error = it
+        }
+        var x = 1
+        assertFailsWith<TestException1> {
+            tryTest {
+                x = 2
+                throw TestException1()
+                x = 3
+            } cleanupStep {
+                throw TestException2()
+                x = 4
+            } cleanup {
+                x = 5
+            }
+        }
+
+        assertEquals(5, x)
+        assertTrue(error is TestException1)
+    }
+
+    @Test
+    fun testErrorReportingInCatch() {
+        var error: Throwable? = null
+        TryTestConfig.swapDiagnosticsCollector {
+            assertNull(error)
+            error = it
+        }
+        var x = 1
+        assertFailsWith<TestException2> {
+            tryTest {
+                throw TestException1()
+                x = 2
+            }.catch<TestException1> {
+                throw TestException2()
+                x = 3
+            } cleanup {
+                x = 4
+            }
+        }
+
+        assertEquals(4, x)
+        assertTrue(error is TestException2)
+    }
 }
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt b/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt
index ea86281..d4d33d9 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt
@@ -112,7 +112,7 @@
     override fun onSetUp() {
         assertNull(instance, "ConnectivityDiagnosticsCollectors were set up multiple times")
         instance = this
-        TryTestConfig.setDiagnosticsCollector { throwable ->
+        TryTestConfig.swapDiagnosticsCollector { throwable ->
             if (runOnFailure(throwable)) {
                 collectTestFailureDiagnostics(throwable)
             }
diff --git a/staticlibs/testutils/hostdevice/com/android/testutils/Cleanup.kt b/staticlibs/testutils/hostdevice/com/android/testutils/Cleanup.kt
index dcd422c..45c69c9 100644
--- a/staticlibs/testutils/hostdevice/com/android/testutils/Cleanup.kt
+++ b/staticlibs/testutils/hostdevice/com/android/testutils/Cleanup.kt
@@ -75,13 +75,21 @@
  */
 
 object TryTestConfig {
-    internal var diagnosticsCollector: Consumer<Throwable>? = null
+    private var diagnosticsCollector: Consumer<Throwable>? = null
 
     /**
      * Set the diagnostics collector to be used in case of failure in [tryTest].
+     *
+     * @return The previous collector.
      */
-    fun setDiagnosticsCollector(collector: Consumer<Throwable>) {
+    fun swapDiagnosticsCollector(collector: Consumer<Throwable>?): Consumer<Throwable>? {
+        val oldCollector = diagnosticsCollector
         diagnosticsCollector = collector
+        return oldCollector
+    }
+
+    fun reportError(e: Throwable) {
+        diagnosticsCollector?.accept(e)
     }
 }
 
@@ -90,14 +98,10 @@
         try {
             Result.success(block())
         } catch (e: Throwable) {
-            TryTestConfig.diagnosticsCollector?.accept(e)
             Result.failure(e)
-        })
+        }, skipErrorReporting = false)
 
-// 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")
-inline class TryExpr<T>(val result: Result<T>) {
+class TryExpr<T>(val result: Result<T>, val skipErrorReporting: Boolean) {
     inline infix fun <reified E : Throwable> catch(block: (E) -> T): TryExpr<T> {
         val originalException = result.exceptionOrNull()
         if (originalException !is E) return this
@@ -105,23 +109,32 @@
             Result.success(block(originalException))
         } catch (e: Throwable) {
             Result.failure(e)
-        })
+        }, this.skipErrorReporting)
     }
 
     @CheckReturnValue
     inline infix fun cleanupStep(block: () -> Unit): TryExpr<T> {
+        // Report errors before the cleanup step, but after catch blocks that may suppress it
+        val originalException = result.exceptionOrNull()
+        var nextSkipErrorReporting = skipErrorReporting
+        if (!skipErrorReporting && originalException != null) {
+            TryTestConfig.reportError(originalException)
+            nextSkipErrorReporting = true
+        }
         try {
             block()
         } catch (e: Throwable) {
-            val originalException = result.exceptionOrNull()
-            return TryExpr(if (null == originalException) {
-                Result.failure(e)
+            return if (null == originalException) {
+                if (!skipErrorReporting) {
+                    TryTestConfig.reportError(e)
+                }
+                TryExpr(Result.failure(e), skipErrorReporting = true)
             } else {
                 originalException.addSuppressed(e)
-                Result.failure(originalException)
-            })
+                TryExpr(Result.failure(originalException), true)
+            }
         }
-        return this
+        return TryExpr(result, nextSkipErrorReporting)
     }
 
     inline infix fun cleanup(block: () -> Unit): T = cleanupStep(block).result.getOrThrow()