Allow collecting diagnostics from shell commands

Allow adding output from any shell command to the test diagnostics file,
instead of just dumpsys.

This introduces a new ShellUtil that can also help run more complex
commands, with pipes and redirects or as root for example.

Bug: 317602748
Test: atest with tests using this util
Change-Id: I0a0b5c7be5afd9e592dd820bfdbcee39ade418f8
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt b/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt
index c7d6850..4b9429b 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/ConnectivityDiagnosticsCollector.kt
@@ -430,19 +430,32 @@
      * @param dumpsysCmd The dumpsys command to run (for example "connectivity").
      * @param exceptionContext An exception to write a stacktrace to the dump for context.
      */
-    fun collectDumpsys(dumpsysCmd: String, exceptionContext: Throwable? = null) {
-        Log.i(TAG, "Collecting dumpsys $dumpsysCmd for test artifacts")
+    fun collectDumpsys(dumpsysCmd: String, exceptionContext: Throwable? = null) =
+        collectCommandOutput("dumpsys $dumpsysCmd", exceptionContext = exceptionContext)
+
+    /**
+     * Add the output of a command to the test data dump.
+     *
+     * <p>The output will be collected immediately, and exported to a test artifact file when the
+     * test ends.
+     * @param cmd The command to run. Stdout of the command will be collected.
+     * @param shell The shell to run the command in.
+     * @param exceptionContext An exception to write a stacktrace to the dump for context.
+     */
+    fun collectCommandOutput(
+        cmd: String,
+        shell: String = "sh",
+        exceptionContext: Throwable? = null
+    ) {
+        Log.i(TAG, "Collecting '$cmd' for test artifacts")
         PrintWriter(buffer).let {
-            it.println("--- Dumpsys $dumpsysCmd at ${ZonedDateTime.now()} ---")
+            it.println("--- $cmd at ${ZonedDateTime.now()} ---")
             maybeWriteExceptionContext(it, exceptionContext)
             it.flush()
         }
-        ParcelFileDescriptor.AutoCloseInputStream(
-            InstrumentationRegistry.getInstrumentation().uiAutomation.executeShellCommand(
-                "dumpsys $dumpsysCmd"
-            )
-        ).use {
-            it.copyTo(buffer)
+
+        runCommandInShell(cmd, shell) { stdout, _ ->
+            stdout.copyTo(buffer)
         }
     }
 
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/ShellUtil.kt b/staticlibs/testutils/devicetests/com/android/testutils/ShellUtil.kt
new file mode 100644
index 0000000..fadc2ab
--- /dev/null
+++ b/staticlibs/testutils/devicetests/com/android/testutils/ShellUtil.kt
@@ -0,0 +1,75 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@file:JvmName("ShellUtil")
+
+package com.android.testutils
+
+import android.app.UiAutomation
+import android.os.ParcelFileDescriptor.AutoCloseInputStream
+import android.os.ParcelFileDescriptor.AutoCloseOutputStream
+import androidx.test.platform.app.InstrumentationRegistry
+import java.io.InputStream
+
+/**
+ * Run a command in a shell.
+ *
+ * Compared to [UiAutomation.executeShellCommand], this allows running commands with pipes and
+ * redirections. [UiAutomation.executeShellCommand] splits the command on spaces regardless of
+ * quotes, so it is not able to run commands like `sh -c "echo 123 > some_file"`.
+ *
+ * @param cmd Shell command to run.
+ * @param shell Command used to run the shell.
+ * @param outputProcessor Function taking stdout, stderr as argument. The streams will be closed
+ *                        when this function returns.
+ * @return Result of [outputProcessor].
+ */
+fun <T> runCommandInShell(
+    cmd: String,
+    shell: String = "sh",
+    outputProcessor: (InputStream, InputStream) -> T,
+): T {
+    val (stdout, stdin, stderr) = InstrumentationRegistry.getInstrumentation().uiAutomation
+        .executeShellCommandRwe(shell)
+    AutoCloseOutputStream(stdin).bufferedWriter().use { it.write(cmd) }
+    AutoCloseInputStream(stdout).use { outStream ->
+        AutoCloseInputStream(stderr).use { errStream ->
+            return outputProcessor(outStream, errStream)
+        }
+    }
+}
+
+/**
+ * Run a command in a shell.
+ *
+ * Overload of [runCommandInShell] that reads and returns stdout as String.
+ */
+fun runCommandInShell(
+    cmd: String,
+    shell: String = "sh",
+) = runCommandInShell(cmd, shell) { stdout, _ ->
+    stdout.reader().use { it.readText() }
+}
+
+/**
+ * Run a command in a root shell.
+ *
+ * This is generally only usable on devices on which [DeviceInfoUtils.isDebuggable] is true.
+ * @see runCommandInShell
+ */
+fun runCommandInRootShell(
+    cmd: String
+) = runCommandInShell(cmd, shell = "su root sh")