Merge changes I677dff58,I5748fbbb,I7803be1e

* changes:
  Remove microdroid_payload.json
  microdroid_manager waits for task to exit
  Refactor MicrodroidTestCase
diff --git a/apex/Android.bp b/apex/Android.bp
index 53c77db..5b6f0a9 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -55,7 +55,6 @@
         "com.android.virt.init.rc",
         "microdroid_cdisk.json",
         "microdroid_cdisk_env.json",
-        "microdroid_payload.json",
         "microdroid_uboot_env",
         "microdroid_bootloader",
     ],
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index b454ea2..6f731bb 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -398,11 +398,6 @@
 }
 
 prebuilt_etc {
-    name: "microdroid_payload.json",
-    src: "microdroid_payload.json",
-}
-
-prebuilt_etc {
     name: "microdroid_vendor_manifest",
     src: "microdroid_vendor_manifest.xml",
     filename: "manifest.xml",
diff --git a/microdroid/README.md b/microdroid/README.md
index 6b9f4b1..c2a65d5 100644
--- a/microdroid/README.md
+++ b/microdroid/README.md
@@ -42,9 +42,11 @@
 
 ## Running
 
-Create a config file, `microdroid.json`:
+Create two config files, `microdroid.json` and `payload.json`:
 
 ```json
+microdroid.json:
+
 {
   "bootloader": "/apex/com.android.virt/etc/microdroid_bootloader",
   "disks": [
@@ -108,9 +110,20 @@
     }
   ]
 }
+
+payload.json:
+
+{
+  "system_apexes" : [
+    "com.android.adbd",
+    "com.android.i18n",
+    "com.android.os.statsd",
+    "com.android.sdkext"
+  ],
+}
 ```
 
-Copy the artifacts to the temp directory, create the composite image using
+Copy the these files to the temp directory, create the composite image using
 `mk_cdisk` and copy the VM config file. For now, some other files have to be
 manually created. In the future, you won't need these, and this shall be done
 via [`virtualizationservice`](../virtualizationservice/).
@@ -119,7 +132,8 @@
 $ adb root
 $ adb shell 'mkdir /data/local/tmp/microdroid'
 $ adb shell 'dd if=/dev/zero of=/data/local/tmp/microdroid/misc.img bs=4k count=256'
-$ adb shell 'cd /data/local/tmp/microdroid; /apex/com.android.virt/bin/mk_payload /apex/com.android.virt/etc/microdroid_payload.json payload.img'
+$ adb push payload.json /data/local/tmp/microdroid/payload.json
+$ adb shell 'cd /data/local/tmp/microdroid; /apex/com.android.virt/bin/mk_payload payload.json payload.img'
 $ adb shell 'chmod go+r /data/local/tmp/microdroid/payload*'
 $ adb push microdroid.json /data/local/tmp/microdroid/microdroid.json
 ```
diff --git a/microdroid/microdroid_cdisk.json b/microdroid/microdroid_cdisk.json
index 5721775..e1ca826 100644
--- a/microdroid/microdroid_cdisk.json
+++ b/microdroid/microdroid_cdisk.json
@@ -2,43 +2,43 @@
   "partitions": [
     {
       "label": "misc",
-      "path": "misc.img"
+      "path": "/data/local/tmp/virt/misc.img"
     },
     {
       "label": "boot_a",
-      "path": "microdroid_boot-5.10.img"
+      "path": "/apex/com.android.virt/etc/fs/microdroid_boot-5.10.img"
     },
     {
       "label": "boot_b",
-      "path": "microdroid_boot-5.10.img"
+      "path": "/apex/com.android.virt/etc/fs/microdroid_boot-5.10.img"
     },
     {
       "label": "vendor_boot_a",
-      "path": "microdroid_vendor_boot-5.10.img"
+      "path": "/apex/com.android.virt/etc/fs/microdroid_vendor_boot-5.10.img"
     },
     {
       "label": "vendor_boot_b",
-      "path": "microdroid_vendor_boot-5.10.img"
+      "path": "/apex/com.android.virt/etc/fs/microdroid_vendor_boot-5.10.img"
     },
     {
       "label": "vbmeta_a",
-      "path": "microdroid_vbmeta.img"
+      "path": "/apex/com.android.virt/etc/fs/microdroid_vbmeta.img"
     },
     {
       "label": "vbmeta_b",
-      "path": "microdroid_vbmeta.img"
+      "path": "/apex/com.android.virt/etc/fs/microdroid_vbmeta.img"
     },
     {
       "label": "vbmeta_system_a",
-      "path": "microdroid_vbmeta_system.img"
+      "path": "/apex/com.android.virt/etc/fs/microdroid_vbmeta_system.img"
     },
     {
       "label": "vbmeta_system_b",
-      "path": "microdroid_vbmeta_system.img"
+      "path": "/apex/com.android.virt/etc/fs/microdroid_vbmeta_system.img"
     },
     {
       "label": "super",
-      "path": "microdroid_super.img"
+      "path": "/apex/com.android.virt/etc/fs/microdroid_super.img"
     }
   ]
 }
diff --git a/microdroid/microdroid_cdisk_env.json b/microdroid/microdroid_cdisk_env.json
index b0fbe44..8c84ed5 100644
--- a/microdroid/microdroid_cdisk_env.json
+++ b/microdroid/microdroid_cdisk_env.json
@@ -2,7 +2,7 @@
   "partitions": [
     {
       "label": "uboot_env",
-      "path": "uboot_env.img"
+      "path": "/apex/com.android.virt/etc/uboot_env.img"
     }
   ]
 }
diff --git a/microdroid/microdroid_payload.json b/microdroid/microdroid_payload.json
deleted file mode 100644
index 83766de..0000000
--- a/microdroid/microdroid_payload.json
+++ /dev/null
@@ -1,13 +0,0 @@
-{
-  "system_apexes": [
-    "com.android.adbd",
-    "com.android.i18n",
-    "com.android.os.statsd",
-    "com.android.sdkext"
-  ],
-  "apk": {
-    "name": "com.android.microdroid.test",
-    "path": "APK_PATH"
-  },
-  "payload_config_path": "/mnt/apk/assets/vm_config.json"
-}
diff --git a/microdroid/sepolicy/system/private/microdroid_launcher.te b/microdroid/sepolicy/system/private/microdroid_launcher.te
index e68b687..5a313b6 100644
--- a/microdroid/sepolicy/system/private/microdroid_launcher.te
+++ b/microdroid/sepolicy/system/private/microdroid_launcher.te
@@ -21,3 +21,6 @@
 
 # Allow to use terminal
 allow microdroid_launcher devpts:chr_file rw_file_perms;
+
+# Allow to set debug prop
+set_prop(microdroid_launcher, debug_prop)
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index ae72a59..b71daa8 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -57,8 +57,15 @@
 
 fn exec_task(task: &Task) -> Result<()> {
     info!("executing main task {:?}...", task);
-    build_command(task)?.spawn()?;
-    Ok(())
+    let exit_status = build_command(task)?.spawn()?.wait()?;
+    if exit_status.success() {
+        Ok(())
+    } else {
+        match exit_status.code() {
+            Some(code) => bail!("task exited with exit code: {}", code),
+            None => bail!("task terminated by signal"),
+        }
+    }
 }
 
 fn build_command(task: &Task) -> Result<Command> {
diff --git a/tests/hostside/java/android/virt/test/MicrodroidTestCase.java b/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
index baa383b..87a9822 100644
--- a/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
+++ b/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
@@ -28,22 +28,28 @@
 import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
 import com.android.tradefed.util.CommandResult;
 import com.android.tradefed.util.CommandStatus;
-import com.android.tradefed.util.FileUtil;
 import com.android.tradefed.util.RunUtil;
 
+import org.json.JSONArray;
+import org.json.JSONObject;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.io.BufferedReader;
 import java.io.File;
 import java.io.FileWriter;
+import java.io.InputStream;
+import java.io.InputStreamReader;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+import java.util.zip.ZipFile;
 
 @RunWith(DeviceJUnit4ClassRunner.class)
 public class MicrodroidTestCase extends BaseHostJUnit4Test {
@@ -52,158 +58,210 @@
     private static final int TEST_VM_CID = 10;
     private static final int TEST_VM_ADB_PORT = 8000;
     private static final String MICRODROID_SERIAL = "localhost:" + TEST_VM_ADB_PORT;
-    private static final String TEST_APK = "MicrodroidTestApp.apk";
+
     // This is really slow on GCE (2m 40s) but fast on localhost or actual Android phones (< 10s)
     // Set the maximum timeout value big enough.
     private static final long MICRODROID_BOOT_TIMEOUT_MINUTES = 5;
 
-    private String executeCommand(String cmd) {
-        final long defaultCommandTimeoutMillis = 3000; // 3 sec. Can be slow on GCE
-        return executeCommand(defaultCommandTimeoutMillis, cmd);
-    }
-
-    private String executeCommand(long timeout, String cmd) {
-        CommandResult result = RunUtil.getDefault().runTimedCmd(timeout, cmd.split(" "));
-        return result.getStdout().trim(); // remove the trailing whitespace including newline
-    }
-
-    private String executeCommandOnMicrodroid(String cmd) {
-        cmd = "adb -s " + MICRODROID_SERIAL + " " + cmd;
-        return executeCommand(cmd);
-    }
-
     @Test
     public void testMicrodroidBoots() throws Exception {
-        // Prepare input files
-        String prepareImagesCmd =
-                String.format(
-                        "mkdir -p %s; cd %s; "
-                                + "cp %setc/microdroid_bootloader bootloader && "
-                                + "cp %setc/fs/*.img . && "
-                                + "cp %setc/uboot_env.img . && "
-                                + "dd if=/dev/zero of=misc.img bs=4k count=256",
-                        TEST_ROOT, TEST_ROOT, VIRT_APEX, VIRT_APEX, VIRT_APEX);
-        getDevice().executeShellCommand(prepareImagesCmd);
+        final String apkName = "MicrodroidTestApp.apk";
+        final String packageName = "com.android.microdroid.test";
+        final String configPath = "assets/vm_config.json"; // path inside the APK
+        startMicrodroid(apkName, packageName, configPath);
+        waitForMicrodroidBoot(MICRODROID_BOOT_TIMEOUT_MINUTES);
+        adbConnectToMicrodroid();
 
-        // Create os_composite.img, env_composite.img, and payload.img
-        String makeOsCompositeCmd =
-                String.format(
-                        "cd %s; %sbin/mk_cdisk %setc/microdroid_cdisk.json os_composite.img",
-                        TEST_ROOT, VIRT_APEX, VIRT_APEX);
-        getDevice().executeShellCommand(makeOsCompositeCmd);
-        String makeEnvCompositeCmd =
-                String.format(
-                        "cd %s; %sbin/mk_cdisk %setc/microdroid_cdisk_env.json env_composite.img",
-                        TEST_ROOT, VIRT_APEX, VIRT_APEX);
-        getDevice().executeShellCommand(makeEnvCompositeCmd);
+        // Check if it actually booted by reading a sysprop.
+        assertThat(runOnMicrodroid("getprop", "ro.hardware"), is("microdroid"));
 
-        String idsigPath = TEST_ROOT + TEST_APK + ".idsig";
-        File idsigFile = (new CompatibilityBuildHelper(getBuild())).getTestFile(TEST_APK + ".idsig");
-        getDevice().pushFile(idsigFile, idsigPath);
+        // Test writing to /data partition
+        runOnMicrodroid("echo MicrodroidTest > /data/local/tmp/test.txt");
+        assertThat(runOnMicrodroid("cat /data/local/tmp/test.txt"), is("MicrodroidTest"));
 
-        // TODO(b/190343842): pass this file to mk_payload
-        String idsigExists = getDevice().executeShellV2Command("[ -f " + idsigPath + " ] && echo ok")
-                .getStdout().trim();
-        assertThat(idsigExists, is("ok"));
+        // Check if the APK partition exists
+        final String apkPartition = "/dev/block/by-name/microdroid-apk";
+        assertThat(runOnMicrodroid("ls", apkPartition), is(apkPartition));
 
-        String payloadJsonOrig = VIRT_APEX + "etc/microdroid_payload.json";
-        String mkPayload = VIRT_APEX + "bin/mk_payload";
+        // Check if the APK is mounted using zipfuse
+        final String mountEntry = "zipfuse on /mnt/apk type fuse.zipfuse";
+        assertThat(runOnMicrodroid("mount"), containsString(mountEntry));
+
+        // Check if the native library in the APK is has correct filesystem info
+        final String[] abis = runOnMicrodroid("getprop", "ro.product.cpu.abilist").split(",");
+        assertThat(abis.length, is(1));
+        final String testLib = "/mnt/apk/lib/" + abis[0] + "/MicrodroidTestNativeLib.so";
+        final String label = "u:object_r:system_file:s0";
+        assertThat(runOnMicrodroid("ls", "-Z", testLib), is(label + " " + testLib));
+
+        // Check if the command in vm_config.json was executed by examining the side effect of the
+        // command
+        assertThat(runOnMicrodroid("getprop", "debug.microdroid.app.run"), is("true"));
+
+        // Manually execute the library and check the output
+        final String microdroidLauncher = "system/bin/microdroid_launcher";
+        assertThat(
+                runOnMicrodroid(microdroidLauncher, testLib, "arg1", "arg2"),
+                is("Hello Microdroid " + testLib + " arg1 arg2"));
+
+        // Shutdown microdroid
+        runOnMicrodroid("reboot");
+    }
+
+    // Run an arbitrary command in the host side and returns the result
+    private String runOnHost(String... cmd) {
+        final long timeout = 10000;
+        CommandResult result = RunUtil.getDefault().runTimedCmd(timeout, cmd);
+        assertThat(result.getStatus(), is(CommandStatus.SUCCESS));
+        return result.getStdout().trim();
+    }
+
+    // Same as runOnHost, but failure is not an error
+    private String tryRunOnHost(String... cmd) {
+        final long timeout = 10000;
+        CommandResult result = RunUtil.getDefault().runTimedCmd(timeout, cmd);
+        return result.getStdout().trim();
+    }
+
+    // Run a shell command on Android
+    private String runOnAndroid(String... cmd) throws Exception {
+        CommandResult result = getDevice().executeShellV2Command(join(cmd));
+        assertThat(result.getStatus(), is(CommandStatus.SUCCESS));
+        return result.getStdout().trim();
+    }
+
+    // Same as runOnAndroid, but failutre is not an error
+    private String tryRunOnAndroid(String... cmd) throws Exception {
+        CommandResult result = getDevice().executeShellV2Command(join(cmd));
+        return result.getStdout().trim();
+    }
+
+    // Run a shell command on Microdroid
+    private String runOnMicrodroid(String... cmd) {
+        final long timeout = 3000; // 3 sec. Microdroid is extremely slow on GCE-on-CF.
+        CommandResult result =
+                RunUtil.getDefault()
+                        .runTimedCmd(timeout, "adb", "-s", MICRODROID_SERIAL, "shell", join(cmd));
+        assertThat(result.getStatus(), is(CommandStatus.SUCCESS));
+        return result.getStdout().trim();
+    }
+
+    private String join(String... strs) {
+        return String.join(" ", Arrays.asList(strs));
+    }
+
+    // TODO(b/191131043) remove this step
+    private String createMiscImage() throws Exception {
+        final String output = TEST_ROOT + "misc.img";
+        runOnAndroid("dd", "if=/dev/zero", "of=" + output, "bs=4k", "count=256");
+        assertThat(runOnAndroid("du", "-b", output), is(not("")));
+        return output;
+    }
+
+    private String createPayloadImage(String apkName, String packageName, String configPath)
+            throws Exception {
+        File apkFile = findTestFile(apkName);
+        getDevice().installPackage(apkFile, /* reinstall */ true);
+
+        // Read the config file from the apk and parse it to know the list of APEXes needed
+        ZipFile apkAsZip = new ZipFile(apkFile);
+        InputStream is = apkAsZip.getInputStream(apkAsZip.getEntry(configPath));
+        String configString =
+                new BufferedReader(new InputStreamReader(is))
+                        .lines()
+                        .collect(Collectors.joining("\n"));
+        JSONObject configObject = new JSONObject(configString);
+        JSONArray apexes = configObject.getJSONArray("apexes");
+        List<String> apexNames = new ArrayList<>();
+        for (int i = 0; i < apexes.length(); i++) {
+            JSONObject anApex = apexes.getJSONObject(i);
+            apexNames.add(anApex.getString("name"));
+        }
 
         // Get the path to the installed apk. Note that
         // getDevice().getAppPackageInfo(...).getCodePath() doesn't work due to the incorrect
-        // parsing of the "=" character. (b/190975227)
-        String apkPath = getDevice().executeShellV2Command("pm path com.android.microdroid.test")
-                .getStdout().trim();
+        // parsing of the "=" character. (b/190975227). So we use the `pm path` command directly.
+        String apkPath = runOnAndroid("pm", "path", packageName);
         assertTrue(apkPath.startsWith("package:"));
         apkPath = apkPath.substring("package:".length());
-        apkPath = apkPath.replace("/", "\\\\/");
 
-        // Replace APK_PATH in the json file with the actual path of the test APK
-        String makePayloadCompositeCmd =
-                "cd " + TEST_ROOT + ";" +
-                String.format("sed s/APK_PATH/%s/ %s > %s;", apkPath, payloadJsonOrig, "payload.json") +
-                mkPayload + " payload.json payload.img";
-        getDevice().executeShellCommand(makePayloadCompositeCmd);
+        // Create payload.json from the gathered data
+        JSONObject payloadObject = new JSONObject();
+        payloadObject.put("system_apexes", new JSONArray(apexNames));
+        payloadObject.put("payload_config_path", "/mnt/apk/" + configPath);
+        JSONObject apkObject = new JSONObject();
+        apkObject.put("path", apkPath);
+        apkObject.put("name", packageName);
+        payloadObject.put("apk", apkObject);
 
-        // Make sure that the composite images are created
-        final List<String> compositeImages =
-                new ArrayList<>(
-                        Arrays.asList(
-                                TEST_ROOT + "/os_composite.img",
-                                TEST_ROOT + "/env_composite.img",
-                                TEST_ROOT + "/payload.img"));
-        CommandResult result =
-                getDevice().executeShellV2Command("du -b " + String.join(" ", compositeImages));
-        assertThat(result.getExitCode(), is(0));
-        assertThat(result.getStdout(), is(not("")));
+        // Push the idsig file to the device
+        // TODO(b/190343842): pass path to this file to payloadObject
+        // File idsigOnHost = findTestFile(apkFile + ".idsig");
+        // final String testApkIdsig = TEST_ROOT + apkFile + ".idsig";
+        // getDevice().pushFile(idsigOnHost, testApkIdsig);
+
+        // Copy the json file to Android
+        File payloadJsonOnHost = File.createTempFile("payload", "json");
+        FileWriter writer = new FileWriter(payloadJsonOnHost);
+        writer.write(payloadObject.toString());
+        writer.close();
+        final String payloadJson = TEST_ROOT + "payload.json";
+        getDevice().pushFile(payloadJsonOnHost, payloadJson);
+
+        // Finally run mk_payload to create payload.img
+        final String mkPayload = VIRT_APEX + "bin/mk_payload";
+        final String payloadImg = TEST_ROOT + "payload.img";
+        runOnAndroid(mkPayload, payloadJson, payloadImg);
+        assertThat(runOnAndroid("du", "-b", payloadImg), is(not("")));
+
+        return payloadImg;
+    }
+
+    private File findTestFile(String name) throws Exception {
+        return (new CompatibilityBuildHelper(getBuild())).getTestFile(name);
+    }
+
+    private void startMicrodroid(String apkName, String packageName, String configPath)
+            throws Exception {
+        // Create misc.img and payload.img
+        final String payloadImg = createPayloadImage(apkName, packageName, configPath);
+        final String miscImg = createMiscImage();
+
+        // Tools and executables
+        final String mkCdisk = VIRT_APEX + "bin/mk_cdisk";
+        final String crosvm = VIRT_APEX + "bin/crosvm";
+
+        // Create os_composisite.img and env_composite.img
+        // TODO(jiyong): remove this when running a VM is done by `vm`
+        final String cdiskJson = VIRT_APEX + "etc/microdroid_cdisk.json";
+        final String cdiskEnvJson = VIRT_APEX + "etc/microdroid_cdisk_env.json";
+        final String osImg = TEST_ROOT + "os_composite.img";
+        final String envImg = TEST_ROOT + "env_composite.img";
+        final String bootloader = VIRT_APEX + "etc/microdroid_bootloader";
+        runOnAndroid(mkCdisk, cdiskJson, osImg);
+        runOnAndroid(mkCdisk, cdiskEnvJson, envImg);
 
         // Start microdroid using crosvm
+        // TODO(jiyong): do this via the `vm` command
         ExecutorService executor = Executors.newFixedThreadPool(1);
-        String runMicrodroidCmd =
-                String.format(
-                        "cd %s; %sbin/crosvm run --cid=%d --disable-sandbox --bios=bootloader"
-                                + " --serial=type=syslog --disk=os_composite.img"
-                                + " --disk=env_composite.img --disk=payload.img &",
-                        TEST_ROOT, VIRT_APEX, TEST_VM_CID);
         executor.execute(
                 () -> {
                     try {
-                        getDevice().executeShellV2Command(runMicrodroidCmd);
+                        runOnAndroid(
+                                crosvm,
+                                "run",
+                                "--cid=" + TEST_VM_CID,
+                                "--disable-sandbox",
+                                "--bios=" + bootloader,
+                                "--serial=type=syslog",
+                                "--disk=" + osImg,
+                                "--disk=" + envImg,
+                                "--disk=" + payloadImg,
+                                "&");
                     } catch (Exception e) {
                         throw new RuntimeException(e);
                     }
                 });
-        waitForMicrodroidBoot(MICRODROID_BOOT_TIMEOUT_MINUTES);
-
-        // Connect to microdroid and read a system property from there
-        executeCommand(
-                "adb -s "
-                        + getDevice().getSerialNumber()
-                        + " forward tcp:"
-                        + TEST_VM_ADB_PORT
-                        + " vsock:"
-                        + TEST_VM_CID
-                        + ":5555");
-        executeCommand("adb connect " + MICRODROID_SERIAL);
-        String prop = executeCommandOnMicrodroid("shell getprop ro.hardware");
-        assertThat(prop, is("microdroid"));
-
-        // Test writing to /data partition
-        File tmpFile = FileUtil.createTempFile("test", ".txt");
-        tmpFile.deleteOnExit();
-        FileWriter writer = new FileWriter(tmpFile);
-        writer.write("MicrodroidTest");
-        writer.close();
-
-        executeCommandOnMicrodroid("push " + tmpFile.getPath() + " /data/local/tmp/test.txt");
-        assertThat(
-                executeCommandOnMicrodroid("shell cat /data/local/tmp/test.txt"),
-                is("MicrodroidTest"));
-
-        assertThat(
-                executeCommandOnMicrodroid("shell ls /dev/block/by-name/microdroid-apk"),
-                is("/dev/block/by-name/microdroid-apk"));
-
-        assertThat(
-                executeCommandOnMicrodroid("shell mount"),
-                containsString("zipfuse on /mnt/apk type fuse.zipfuse"));
-
-        final String[] abiList =
-                executeCommandOnMicrodroid("shell getprop ro.product.cpu.abilist").split(",");
-        assertThat(abiList.length, is(1));
-
-        final String libPath = "/mnt/apk/lib/" + abiList[0] + "/MicrodroidTestNativeLib.so";
-        assertThat(
-                executeCommandOnMicrodroid("shell ls -Z " + libPath),
-                is("u:object_r:system_file:s0 " + libPath));
-
-        assertThat(
-                executeCommandOnMicrodroid(
-                        "shell /system/bin/microdroid_launcher " + libPath + " arg1 arg2"),
-                is("Hello Microdroid " + libPath + " arg1 arg2"));
-
-        // Shutdown microdroid
-        executeCommand("adb -s localhost:" + TEST_VM_ADB_PORT + " shell reboot");
     }
 
     private void waitForMicrodroidBoot(long timeoutMinutes) throws Exception {
@@ -217,9 +275,19 @@
                         TimeUnit.MINUTES);
     }
 
+    // Establish an adb connection to microdroid by letting Android forward the connection to
+    // microdroid.
+    private void adbConnectToMicrodroid() {
+        final String serial = getDevice().getSerialNumber();
+        final String from = "tcp:" + TEST_VM_ADB_PORT;
+        final String to = "vsock:" + TEST_VM_CID + ":5555";
+        runOnHost("adb", "-s", serial, "forward", from, to);
+        runOnHost("adb", "connect", MICRODROID_SERIAL);
+    }
+
     private void skipIfFail(String command) throws Exception {
-        assumeThat(
-                getDevice().executeShellV2Command(command).getStatus(), is(CommandStatus.SUCCESS));
+        CommandResult result = getDevice().executeShellV2Command(command);
+        assumeThat(result.getStatus(), is(CommandStatus.SUCCESS));
     }
 
     @Before
@@ -234,24 +302,25 @@
     @Before
     public void setUp() throws Exception {
         // kill stale crosvm processes
-        getDevice().executeShellV2Command("killall crosvm");
+        tryRunOnAndroid("killall", "crosvm");
 
-        // delete the test root
-        getDevice().executeShellCommand("rm -rf " + TEST_ROOT);
+        // Prepare the test root
+        tryRunOnAndroid("rm", "-rf", TEST_ROOT);
+        tryRunOnAndroid("mkdir", "-p", TEST_ROOT);
 
         // disconnect from microdroid
-        executeCommand("adb disconnect " + MICRODROID_SERIAL);
+        tryRunOnHost("adb", "disconnect", MICRODROID_SERIAL);
 
         // clear the log
-        getDevice().executeShellV2Command("logcat -c");
+        tryRunOnAndroid("logcat", "-c");
     }
 
     @After
     public void shutdown() throws Exception {
         // disconnect from microdroid
-        executeCommand("adb disconnect " + MICRODROID_SERIAL);
+        tryRunOnHost("adb", "disconnect", MICRODROID_SERIAL);
 
         // kill stale crosvm processes
-        getDevice().executeShellV2Command("killall crosvm");
+        tryRunOnAndroid("killall", "crosvm");
     }
 }
diff --git a/tests/testapk/assets/vm_config.json b/tests/testapk/assets/vm_config.json
index 7a3df7a..8312f4d 100644
--- a/tests/testapk/assets/vm_config.json
+++ b/tests/testapk/assets/vm_config.json
@@ -9,5 +9,19 @@
       "hello",
       "microdroid"
     ]
-  }
-}
\ No newline at end of file
+  },
+  "apexes": [
+    {
+      "name": "com.android.adbd"
+    },
+    {
+      "name": "com.android.i18n"
+    },
+    {
+      "name": "com.android.os.statsd"
+    },
+    {
+      "name": "com.android.sdkext"
+    }
+  ]
+}
diff --git a/tests/testapk/src/native/testbinary.cpp b/tests/testapk/src/native/testbinary.cpp
index c317cd2..c3eefc4 100644
--- a/tests/testapk/src/native/testbinary.cpp
+++ b/tests/testapk/src/native/testbinary.cpp
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 #include <stdio.h>
+#include <sys/system_properties.h>
 
 extern "C" int android_native_main(int argc, char* argv[]) {
     printf("Hello Microdroid ");
@@ -25,5 +26,7 @@
         }
     }
     printf("\n");
+
+    __system_property_set("debug.microdroid.app.run", "true");
     return 0;
 }