diff --git a/compos/benchmark/Android.bp b/compos/benchmark/Android.bp
index 3d46a39..12b3ae8 100644
--- a/compos/benchmark/Android.bp
+++ b/compos/benchmark/Android.bp
@@ -15,7 +15,7 @@
         "MicrodroidTestHelper",
         "truth-prebuilt",
     ],
-    platform_apis: true,
+    sdk_version: "test_current",
     use_embedded_native_libs: true,
     compile_multilib: "64",
 }
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 b884b9e..3f5ff20 100644
--- a/compos/benchmark/src/java/com/android/compos/benchmark/ComposBenchmark.java
+++ b/compos/benchmark/src/java/com/android/compos/benchmark/ComposBenchmark.java
@@ -20,6 +20,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.TruthJUnit.assume;
 
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 
 import android.app.Instrumentation;
@@ -69,7 +70,7 @@
     }
 
     @After
-    public void tearDown() throws Exception {
+    public void tearDown() {
         mInstrumentation.getUiAutomation().dropShellPermissionIdentity();
     }
 
@@ -96,7 +97,7 @@
             Long compileEndTime = System.nanoTime();
             Timestamp afterCompileLatestTime = getLatestDex2oatSuccessTime();
 
-            assertTrue(afterCompileLatestTime != null);
+            assertNotNull(afterCompileLatestTime);
             assertTrue(
                     beforeCompileLatestTime == null
                             || beforeCompileLatestTime.before(afterCompileLatestTime));
@@ -135,7 +136,7 @@
             threadGetMetrics.start();
 
             Long compileStartTime = System.nanoTime();
-            String output = executeCommand(command);
+            String output = runInShellWithStderr(TAG, mInstrumentation.getUiAutomation(), command);
             Long compileEndTime = System.nanoTime();
             assertThat(output).containsMatch("All Ok");
             double elapsedSec = (compileEndTime - compileStartTime) / NANOS_IN_SEC;
diff --git a/tests/helper/Android.bp b/tests/helper/Android.bp
index c9eafad..f36a15d 100644
--- a/tests/helper/Android.bp
+++ b/tests/helper/Android.bp
@@ -19,5 +19,5 @@
         "MicrodroidTestHelper",
         "truth-prebuilt",
     ],
-    sdk_version: "system_current",
+    sdk_version: "test_current",
 }
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 744f94c..cf17e5b 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
@@ -444,7 +444,25 @@
             return stdout;
         } catch (IOException e) {
             Log.e(tag, "Error executing: " + command, e);
-            throw new RuntimeException("Failed to run the command.");
+            throw new RuntimeException("Failed to run the command.", e);
+        }
+    }
+
+    /** Execute a command. Returns the concatenation of stdout and stderr. */
+    protected String runInShellWithStderr(String tag, UiAutomation uiAutomation, String command) {
+        ParcelFileDescriptor[] files = uiAutomation.executeShellCommandRwe(command);
+        try (InputStream stdout = new ParcelFileDescriptor.AutoCloseInputStream(files[0]);
+                InputStream stderr = new ParcelFileDescriptor.AutoCloseInputStream(files[2]);
+                ByteArrayOutputStream out = new ByteArrayOutputStream()) {
+            files[1].close(); // The command's stdin
+            stdout.transferTo(out);
+            stderr.transferTo(out);
+            String output = out.toString("UTF-8");
+            Log.i(tag, "Got output : " + stdout);
+            return output;
+        } catch (IOException e) {
+            Log.e(tag, "Error executing: " + command, e);
+            throw new RuntimeException("Failed to run the command.", e);
         }
     }