Merge "Add options to disable diagnostics collector" into main
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt b/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt
index 0cb5c34..80bb99b 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt
@@ -18,6 +18,7 @@
 
 import android.device.collectors.BaseMetricListener
 import android.device.collectors.DataRecord
+import android.os.Build
 import android.os.ParcelFileDescriptor
 import android.util.Log
 import androidx.test.platform.app.InstrumentationRegistry
@@ -29,8 +30,19 @@
 import org.junit.runner.Description
 import org.junit.runner.notification.Failure
 
+/**
+ * A diagnostics collector that outputs diagnostics files as test artifacts.
+ *
+ * <p>Collects diagnostics automatically by default on non-local builds. Can be enabled/disabled
+ * manually with:
+ * ```
+ * atest MyModule -- \
+ *     --module-arg MyModule:instrumentation-arg:connectivity-diagnostics-on-failure:=false
+ * ```
+ */
 class ConnectivityDiagnosticsCollector : BaseMetricListener() {
     companion object {
+        private const val ARG_RUN_ON_FAILURE = "connectivity-diagnostics-on-failure"
         private const val COLLECTOR_DIR = "run_listeners/connectivity_diagnostics"
         private const val FILENAME_SUFFIX = "_conndiag.txt"
         private const val MAX_DUMPS = 20
@@ -60,7 +72,11 @@
     override fun onTestFail(testData: DataRecord, description: Description, failure: Failure) {
         // TODO: find a way to disable this behavior only on local runs, to avoid slowing them down
         // when iterating on failing tests.
+        if (!runOnFailure()) return
         if (outputFiles.size >= MAX_DUMPS) return
+        Log.i(TAG, "Collecting diagnostics for test failure. Disable by running tests with: " +
+                "atest MyModule -- " +
+                "--module-arg MyModule:instrumentation-arg:$ARG_RUN_ON_FAILURE:=false")
         collectTestFailureDiagnostics(failure.exception)
 
         val baseFilename = "${description.className}#${description.methodName}_failure"
@@ -80,6 +96,13 @@
         flushBufferToFileMetric(testData, baseFilename)
     }
 
+    private fun runOnFailure(): Boolean {
+        // Do not run on local builds (which have ro.build.version.incremental set to eng.username)
+        // to avoid slowing down local runs.
+        val enabledByDefault = !Build.VERSION.INCREMENTAL.startsWith("eng.")
+        return argsBundle.getString(ARG_RUN_ON_FAILURE)?.toBooleanStrictOrNull() ?: enabledByDefault
+    }
+
     private fun flushBufferToFileMetric(testData: DataRecord, baseFilename: String) {
         var filename = baseFilename
         // In case a method was run multiple times (typically retries), append a number