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;