diff --git a/compos/benchmark/src/java/com/android/compos/benchmark/ComposBenchmark.java b/compos/benchmark/src/java/com/android/compos/benchmark/ComposBenchmark.java
index 996d32a..dd113a6 100644
--- a/compos/benchmark/src/java/com/android/compos/benchmark/ComposBenchmark.java
+++ b/compos/benchmark/src/java/com/android/compos/benchmark/ComposBenchmark.java
@@ -23,7 +23,6 @@
 
 import android.app.Instrumentation;
 import android.os.Bundle;
-import android.os.ParcelFileDescriptor;
 import android.util.Log;
 
 import com.android.microdroid.test.common.MetricsProcessor;
@@ -36,9 +35,7 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
-import java.io.InputStream;
 import java.sql.Timestamp;
 import java.text.DateFormat;
 import java.text.ParseException;
@@ -198,37 +195,8 @@
         processMemory.forEach((k, v) -> reportMetric(prefix + k, unit, v));
     }
 
-    private byte[] executeCommandBlocking(String command) {
-        try (InputStream is =
-                        new ParcelFileDescriptor.AutoCloseInputStream(
-                                mInstrumentation.getUiAutomation().executeShellCommand(command));
-                ByteArrayOutputStream out = new ByteArrayOutputStream()) {
-            byte[] buf = new byte[BUFFER_SIZE];
-            int length;
-            while ((length = is.read(buf)) >= 0) {
-                out.write(buf, 0, length);
-            }
-            return out.toByteArray();
-        } catch (IOException e) {
-            Log.e(TAG, "Error executing: " + command, e);
-            return null;
-        }
-    }
-
     private String executeCommand(String command) {
-        try {
-            byte[] output = executeCommandBlocking(command);
-
-            if (output == null) {
-                throw new RuntimeException("Failed to run the command.");
-            } else {
-                String stdout = new String(output, "UTF-8");
-                Log.i(TAG, "Get stdout : " + stdout);
-                return stdout;
-            }
-        } catch (Exception e) {
-            throw new RuntimeException("Error executing: " + command + " , Exception: " + e);
-        }
+        return runInShell(TAG, mInstrumentation.getUiAutomation(), command);
     }
 
     private class GetMetricsRunnable implements Runnable {
diff --git a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
index a07731e..fdc846e 100644
--- a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
+++ b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
@@ -19,6 +19,7 @@
 
 import static org.junit.Assume.assumeNoException;
 
+import android.app.UiAutomation;
 import android.content.Context;
 import android.os.ParcelFileDescriptor;
 import android.os.SystemProperties;
@@ -35,6 +36,8 @@
 import com.android.virt.VirtualizationTestHelper;
 
 import java.io.BufferedReader;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.util.OptionalLong;
@@ -334,4 +337,20 @@
                 listener.getInitStartedNanoTime(),
                 listener.getPayloadStartedNanoTime());
     }
+
+    /** Execute a command. Returns stdout. */
+    protected String runInShell(String tag, UiAutomation uiAutomation, String command) {
+        try (InputStream is =
+                        new ParcelFileDescriptor.AutoCloseInputStream(
+                                uiAutomation.executeShellCommand(command));
+                ByteArrayOutputStream out = new ByteArrayOutputStream()) {
+            is.transferTo(out);
+            String stdout = out.toString("UTF-8");
+            Log.i(tag, "Got stdout : " + stdout);
+            return stdout;
+        } catch (IOException e) {
+            Log.e(tag, "Error executing: " + command, e);
+            throw new RuntimeException("Failed to run the command.");
+        }
+    }
 }