diff --git a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
index f5656e2..b176cfc 100644
--- a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
+++ b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
@@ -45,7 +45,6 @@
 import org.junit.runner.RunWith;
 
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.regex.Matcher;
@@ -413,11 +412,16 @@
 
         for (int round = 0; round < ROUND_COUNT; ++round) {
             reInstallApex(REINSTALL_APEX_TIMEOUT_SEC);
-            if (isWithCompos) {
-                compileStagedApex(COMPILE_STAGED_APEX_TIMEOUT_SEC);
+            try {
+                if (isWithCompos) {
+                    compileStagedApex(COMPILE_STAGED_APEX_TIMEOUT_SEC);
+                }
+            } finally {
+                // If compilation fails, we still have a staged APEX, and we need to reboot to
+                // clean that up for further tests.
+                getDevice().nonBlockingReboot();
+                waitForBootCompleted();
             }
-            getDevice().nonBlockingReboot();
-            waitForBootCompleted();
 
             double elapsedSec = getDmesgBootTime();
             bootDmesgTime.add(elapsedSec);
@@ -458,8 +462,9 @@
             try {
                 CommandRunner android = new CommandRunner(getDevice());
 
-                String result = android.run(
-                        COMPOSD_CMD_BIN + " staged-apex-compile");
+                String result =
+                        android.runWithTimeout(
+                                3 * 60 * 1000, COMPOSD_CMD_BIN + " staged-apex-compile");
                 assertWithMessage("Failed to compile staged APEX. Reason: " + result)
                     .that(result).ignoringCase().contains("all ok");
 
diff --git a/tests/hostside/helper/java/com/android/microdroid/test/host/CommandRunner.java b/tests/hostside/helper/java/com/android/microdroid/test/host/CommandRunner.java
index 846531d..242dbde 100644
--- a/tests/hostside/helper/java/com/android/microdroid/test/host/CommandRunner.java
+++ b/tests/hostside/helper/java/com/android/microdroid/test/host/CommandRunner.java
@@ -66,9 +66,7 @@
 
     public String runWithTimeout(long timeoutMillis, String... cmd)
             throws DeviceNotAvailableException {
-        CommandResult result =
-                mDevice.executeShellV2Command(
-                        join(cmd), timeoutMillis, java.util.concurrent.TimeUnit.MILLISECONDS);
+        CommandResult result = runForResultWithTimeout(timeoutMillis, cmd);
         if (result.getStatus() != CommandStatus.SUCCESS) {
             fail(join(cmd) + " has failed: " + result);
         }