Merge "Remove redundant zipfuse mount check"
diff --git a/authfs/TEST_MAPPING b/authfs/TEST_MAPPING
index d0c0b09..14f1824 100644
--- a/authfs/TEST_MAPPING
+++ b/authfs/TEST_MAPPING
@@ -2,6 +2,9 @@
   "presubmit": [
     {
       "name": "authfs_device_test_src_lib"
+    },
+    {
+      "name": "AuthFsHostTest"
     }
   ]
 }
diff --git a/authfs/tests/AndroidTest.xml b/authfs/tests/AndroidTest.xml
index 8f940f6..6100ab9 100644
--- a/authfs/tests/AndroidTest.xml
+++ b/authfs/tests/AndroidTest.xml
@@ -18,18 +18,11 @@
     <!-- Need root to start virtualizationservice -->
     <target_preparer class="com.android.tradefed.targetprep.RootTargetPreparer"/>
 
-    <!-- virtualizationservice doesn't have access to shell_data_file. Instead of giving it
-          a test-only permission, run it without selinux -->
+    <!-- Still need to define SELinux policy for authfs and fd_server properly. -->
     <target_preparer class="com.android.tradefed.targetprep.DisableSELinuxTargetPreparer"/>
 
-    <!-- Basic checks that the device has all the prerequisites. -->
     <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
         <option name="throw-if-cmd-fail" value="true" />
-        <!-- Make sure kernel has FUSE enabled. -->
-        <option name="run-command" value="ls /dev/fuse" />
-        <!-- Make sure necessary executables are installed. -->
-        <option name="run-command" value="ls /apex/com.android.virt/bin/fd_server" />
-        <option name="run-command" value="ls /apex/com.android.virt/bin/authfs" />
         <!-- Prepare test directory. -->
         <option name="run-command" value="mkdir -p /data/local/tmp/authfs/mnt" />
         <option name="teardown-command" value="rm -rf /data/local/tmp/authfs" />
diff --git a/authfs/tests/java/src/com/android/fs/AuthFsHostTest.java b/authfs/tests/java/src/com/android/fs/AuthFsHostTest.java
index 426b333..6e1c890 100644
--- a/authfs/tests/java/src/com/android/fs/AuthFsHostTest.java
+++ b/authfs/tests/java/src/com/android/fs/AuthFsHostTest.java
@@ -18,18 +18,26 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assume.assumeFalse;
 
 import android.platform.test.annotations.RootPermissionTest;
+import android.virt.test.CommandRunner;
 import android.virt.test.VirtualizationTestCaseBase;
 
 import com.android.compatibility.common.util.PollingCheck;
 import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.invoker.TestInformation;
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
+import com.android.tradefed.testtype.junit4.AfterClassWithInfo;
+import com.android.tradefed.testtype.junit4.BeforeClassWithInfo;
 import com.android.tradefed.util.CommandResult;
 
 import org.junit.After;
+import org.junit.AssumptionViolatedException;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -37,7 +45,6 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 
-// TODO move to Virtualization/tests/hostside/
 @RootPermissionTest
 @RunWith(DeviceJUnit4ClassRunner.class)
 public final class AuthFsHostTest extends VirtualizationTestCaseBase {
@@ -55,50 +62,87 @@
     private static final String AUTHFS_BIN = "/system/bin/authfs";
 
     /** Plenty of time for authfs to get ready */
-    private static final int AUTHFS_INIT_TIMEOUT_MS = 1500;
+    private static final int AUTHFS_INIT_TIMEOUT_MS = 3000;
 
     /** FUSE's magic from statfs(2) */
     private static final String FUSE_SUPER_MAGIC_HEX = "65735546";
 
-    private ExecutorService mThreadPool;
-    private String mCid;
+    private static CommandRunner sAndroid;
+    private static String sCid;
+    private static boolean sAssumptionFailed;
 
-    @Before
-    public void setUp() throws DeviceNotAvailableException {
-        testIfDeviceIsCapable();
+    private ExecutorService mThreadPool = Executors.newCachedThreadPool();
 
-        cleanUpTestFiles();
+    @BeforeClassWithInfo
+    public static void beforeClassWithDevice(TestInformation testInfo)
+            throws DeviceNotAvailableException {
+        assertNotNull(testInfo.getDevice());
+        ITestDevice androidDevice = testInfo.getDevice();
+        sAndroid = new CommandRunner(androidDevice);
 
-        prepareVirtualizationTestSetup();
+        try {
+            testIfDeviceIsCapable(androidDevice);
+        } catch (AssumptionViolatedException e) {
+            // NB: The assumption exception is NOT handled by the test infra when it is thrown from
+            // a class method (see b/37502066). This has not only caused the loss of log, but also
+            // prevented the test cases to be reported at all and thus confused the test infra.
+            //
+            // Since we want to avoid the big overhead to start the VM repeatedly on CF, let's catch
+            // AssumptionViolatedException and emulate it artifitially.
+            CLog.e("Assumption failed: " + e);
+            sAssumptionFailed = true;
+            return;
+        }
 
-        mThreadPool = Executors.newCachedThreadPool();
+        prepareVirtualizationTestSetup(androidDevice);
 
         // For each test case, boot and adb connect to a new Microdroid
+        CLog.i("Starting the shared VM");
         final String apkName = "MicrodroidTestApp.apk";
         final String packageName = "com.android.microdroid.test";
         final String configPath = "assets/vm_config.json"; // path inside the APK
-        mCid = startMicrodroid(apkName, packageName, configPath, /* debug */ false);
-        adbConnectToMicrodroid(mCid);
+        sCid =
+                startMicrodroid(
+                        androidDevice,
+                        testInfo.getBuildInfo(),
+                        apkName,
+                        packageName,
+                        configPath,
+                        /* debug */ false);
+        adbConnectToMicrodroid(androidDevice, sCid);
 
         // Root because authfs (started from shell in this test) currently require root to open
         // /dev/fuse and mount the FUSE.
         rootMicrodroid();
     }
 
-    @After
-    public void tearDown() throws DeviceNotAvailableException {
-        if (mCid != null) {
-            shutdownMicrodroid(mCid);
-            mCid = null;
+    @AfterClassWithInfo
+    public static void afterClassWithDevice(TestInformation testInfo)
+            throws DeviceNotAvailableException {
+        assertNotNull(sAndroid);
+
+        if (sCid != null) {
+            CLog.i("Shutting down shared VM");
+            shutdownMicrodroid(sAndroid.getDevice(), sCid);
+            sCid = null;
         }
 
-        tryRunOnAndroid("killall fd_server");
-        cleanUpTestFiles();
-        cleanUpVirtualizationTestSetup();
+        cleanUpVirtualizationTestSetup(sAndroid.getDevice());
+        sAndroid = null;
     }
 
-    private void cleanUpTestFiles() throws DeviceNotAvailableException {
-        tryRunOnAndroid("rm -f " + TEST_DIR + "/output");
+    @Before
+    public void setUp() {
+        assumeFalse(sAssumptionFailed);
+    }
+
+    @After
+    public void tearDown() throws DeviceNotAvailableException {
+        sAndroid.tryRun("killall fd_server");
+        sAndroid.tryRun("rm -f " + TEST_DIR + "/output");
+
+        tryRunOnMicrodroid("killall authfs");
+        tryRunOnMicrodroid("umount " + MOUNT_DIR);
     }
 
     @Test
@@ -194,7 +238,7 @@
 
         // Action
         // Tampering with the first 2 4K block of the backing file.
-        runOnAndroid("dd if=/dev/zero of=" + backendPath + " bs=1 count=8192");
+        sAndroid.run("dd if=/dev/zero of=" + backendPath + " bs=1 count=8192");
 
         // Verify
         // Write to a block partially requires a read back to calculate the new hash. It should fail
@@ -274,7 +318,7 @@
     }
 
     private String computeFileHashOnAndroid(String path) throws DeviceNotAvailableException {
-        String result = runOnAndroid("sha256sum " + path);
+        String result = sAndroid.run("sha256sum " + path);
         String[] tokens = result.split("\\s");
         if (tokens.length > 0) {
             return tokens[0];
@@ -320,7 +364,7 @@
                 () -> {
                     try {
                         CLog.i("Starting fd_server");
-                        CommandResult result = getDevice().executeShellV2Command(cmd);
+                        CommandResult result = sAndroid.runForResult(cmd);
                         CLog.w("fd_server has stopped: " + result);
                     } catch (DeviceNotAvailableException e) {
                         CLog.e("Error running fd_server", e);
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 522651b..c46bb2b 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -18,6 +18,8 @@
 
 import static android.os.ParcelFileDescriptor.MODE_READ_WRITE;
 
+import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.content.Context;
 import android.os.IBinder;
 import android.os.ParcelFileDescriptor;
@@ -71,35 +73,36 @@
     }
 
     /** The package which owns this VM. */
-    private final String mPackageName;
+    private final @NonNull String mPackageName;
 
     /** Name of this VM within the package. The name should be unique in the package. */
-    private final String mName;
+    private final @NonNull String mName;
 
     /**
      * Path to the config file for this VM. The config file is where the configuration is persisted.
      */
-    private final File mConfigFilePath;
+    private final @NonNull File mConfigFilePath;
 
     /** Path to the instance image file for this VM. */
-    private final File mInstanceFilePath;
+    private final @NonNull File mInstanceFilePath;
 
     /** Size of the instance image. 10 MB. */
     private static final long INSTANCE_FILE_SIZE = 10 * 1024 * 1024;
 
     /** The configuration that is currently associated with this VM. */
-    private VirtualMachineConfig mConfig;
+    private @NonNull VirtualMachineConfig mConfig;
 
     /** Handle to the "running" VM. */
-    private IVirtualMachine mVirtualMachine;
+    private @Nullable IVirtualMachine mVirtualMachine;
 
     /** The registered callback */
-    private VirtualMachineCallback mCallback;
+    private @Nullable VirtualMachineCallback mCallback;
 
-    private ParcelFileDescriptor mConsoleReader;
-    private ParcelFileDescriptor mConsoleWriter;
+    private @Nullable ParcelFileDescriptor mConsoleReader;
+    private @Nullable ParcelFileDescriptor mConsoleWriter;
 
-    private VirtualMachine(Context context, String name, VirtualMachineConfig config) {
+    private VirtualMachine(
+            @NonNull Context context, @NonNull String name, @NonNull VirtualMachineConfig config) {
         mPackageName = context.getPackageName();
         mName = name;
         mConfig = config;
@@ -115,8 +118,8 @@
      * it is persisted until it is deleted by calling {@link #delete()}. The created virtual machine
      * is in {@link #STOPPED} state. To run the VM, call {@link #run()}.
      */
-    /* package */ static VirtualMachine create(
-            Context context, String name, VirtualMachineConfig config)
+    /* package */ static @NonNull VirtualMachine create(
+            @NonNull Context context, @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
         if (config == null) {
             throw new VirtualMachineException("null config");
@@ -165,8 +168,8 @@
     }
 
     /** Loads a virtual machine that is already created before. */
-    /* package */ static VirtualMachine load(Context context, String name)
-            throws VirtualMachineException {
+    /* package */ static @NonNull VirtualMachine load(
+            @NonNull Context context, @NonNull String name) throws VirtualMachineException {
         VirtualMachine vm = new VirtualMachine(context, name, /* config */ null);
 
         try (FileInputStream input = new FileInputStream(vm.mConfigFilePath)) {
@@ -193,7 +196,7 @@
      * Returns the name of this virtual machine. The name is unique in the package and can't be
      * changed.
      */
-    public String getName() {
+    public @NonNull String getName() {
         return mName;
     }
 
@@ -204,12 +207,12 @@
      * share the same config. It is also possible that a virtual machine can switch its config,
      * which can be done by calling {@link #setConfig(VirtualMachineCOnfig)}.
      */
-    public VirtualMachineConfig getConfig() {
+    public @NonNull VirtualMachineConfig getConfig() {
         return mConfig;
     }
 
     /** Returns the current status of this virtual machine. */
-    public Status getStatus() throws VirtualMachineException {
+    public @NonNull Status getStatus() throws VirtualMachineException {
         try {
             if (mVirtualMachine != null && mVirtualMachine.isRunning()) {
                 return Status.RUNNING;
@@ -227,12 +230,12 @@
      * Registers the callback object to get events from the virtual machine. If a callback was
      * already registered, it is replaced with the new one.
      */
-    public void setCallback(VirtualMachineCallback callback) {
+    public void setCallback(@Nullable VirtualMachineCallback callback) {
         mCallback = callback;
     }
 
     /** Returns the currently registered callback. */
-    public VirtualMachineCallback getCallback() {
+    public @Nullable VirtualMachineCallback getCallback() {
         return mCallback;
     }
 
@@ -304,7 +307,7 @@
     }
 
     /** Returns the stream object representing the console output from the virtual machine. */
-    public InputStream getConsoleOutputStream() throws VirtualMachineException {
+    public @NonNull InputStream getConsoleOutputStream() throws VirtualMachineException {
         if (mConsoleReader == null) {
             throw new VirtualMachineException("Console output not available");
         }
@@ -339,7 +342,7 @@
     }
 
     /** Returns the CID of this virtual machine, if it is running. */
-    public Optional<Integer> getCid() throws VirtualMachineException {
+    public @NonNull Optional<Integer> getCid() throws VirtualMachineException {
         if (getStatus() != Status.RUNNING) {
             return Optional.empty();
         }
@@ -361,7 +364,7 @@
      *
      * @return the old config
      */
-    public VirtualMachineConfig setConfig(VirtualMachineConfig newConfig)
+    public @NonNull VirtualMachineConfig setConfig(@NonNull VirtualMachineConfig newConfig)
             throws VirtualMachineException {
         final VirtualMachineConfig oldConfig = getConfig();
         if (!oldConfig.isCompatibleWith(newConfig)) {
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java b/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java
index 0267de8..07af4a1 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java
@@ -16,6 +16,7 @@
 
 package android.system.virtualmachine;
 
+import android.annotation.NonNull;
 import android.os.ParcelFileDescriptor;
 
 /**
@@ -27,8 +28,8 @@
 public interface VirtualMachineCallback {
 
     /** Called when the payload starts in the VM. */
-    void onPayloadStarted(VirtualMachine vm, ParcelFileDescriptor stdout);
+    void onPayloadStarted(@NonNull VirtualMachine vm, @NonNull ParcelFileDescriptor stdout);
 
     /** Called when the VM died. */
-    void onDied(VirtualMachine vm);
+    void onDied(@NonNull VirtualMachine vm);
 }
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index f0e1ce6..21e1a46 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -18,6 +18,7 @@
 
 import static android.os.ParcelFileDescriptor.MODE_READ_ONLY;
 
+import android.annotation.NonNull;
 import android.content.Context;
 import android.content.pm.PackageManager;
 import android.content.pm.Signature; // This actually is certificate!
@@ -52,23 +53,23 @@
     private static final String KEY_DEBUGMODE = "debugMode";
 
     // Paths to the APK and its idsig file of this application.
-    private final String mApkPath;
-    private final Signature[] mCerts;
-    private final String mIdsigPath;
+    private final @NonNull String mApkPath;
+    private final @NonNull Signature[] mCerts;
+    private final @NonNull String mIdsigPath;
     private final boolean mDebugMode;
 
     /**
      * Path within the APK to the payload config file that defines software aspects of this config.
      */
-    private final String mPayloadConfigPath;
+    private final @NonNull String mPayloadConfigPath;
 
     // TODO(jiyong): add more items like # of cpu, size of ram, debuggability, etc.
 
     private VirtualMachineConfig(
-            String apkPath,
-            Signature[] certs,
-            String idsigPath,
-            String payloadConfigPath,
+            @NonNull String apkPath,
+            @NonNull Signature[] certs,
+            @NonNull String idsigPath,
+            @NonNull String payloadConfigPath,
             boolean debugMode) {
         mApkPath = apkPath;
         mCerts = certs;
@@ -78,7 +79,7 @@
     }
 
     /** Loads a config from a stream, for example a file. */
-    /* package */ static VirtualMachineConfig from(InputStream input)
+    /* package */ static @NonNull VirtualMachineConfig from(@NonNull InputStream input)
             throws IOException, VirtualMachineException {
         PersistableBundle b = PersistableBundle.readFromStream(input);
         final int version = b.getInt(KEY_VERSION);
@@ -111,7 +112,7 @@
     }
 
     /** Persists this config to a stream, for example a file. */
-    /* package */ void serialize(OutputStream output) throws IOException {
+    /* package */ void serialize(@NonNull OutputStream output) throws IOException {
         PersistableBundle b = new PersistableBundle();
         b.putInt(KEY_VERSION, VERSION);
         b.putString(KEY_APKPATH, mApkPath);
@@ -128,7 +129,7 @@
     }
 
     /** Returns the path to the payload config within the owning application. */
-    public String getPayloadConfigPath() {
+    public @NonNull String getPayloadConfigPath() {
         return mPayloadConfigPath;
     }
 
@@ -139,7 +140,7 @@
      * signed by the same signer. All other changes (e.g. using a payload from a different signer,
      * change of the debug mode, etc.) are considered as incompatible.
      */
-    public boolean isCompatibleWith(VirtualMachineConfig other) {
+    public boolean isCompatibleWith(@NonNull VirtualMachineConfig other) {
         if (!Arrays.equals(this.mCerts, other.mCerts)) {
             return false;
         }
@@ -173,7 +174,7 @@
         // TODO(jiyong): add more items like # of cpu, size of ram, debuggability, etc.
 
         /** Creates a builder for the given context (APK), and the payload config file in APK. */
-        public Builder(Context context, String payloadConfigPath) {
+        public Builder(@NonNull Context context, @NonNull String payloadConfigPath) {
             mContext = context;
             mPayloadConfigPath = payloadConfigPath;
             mDebugMode = false;
@@ -188,13 +189,13 @@
         // TODO(jiyong): remove this. Apps shouldn't need to set the path to the idsig file. It
         // should be automatically found or created on demand.
         /** Set the path to the idsig file for the current application. */
-        public Builder idsigPath(String idsigPath) {
+        public Builder idsigPath(@NonNull String idsigPath) {
             mIdsigPath = idsigPath;
             return this;
         }
 
         /** Builds an immutable {@link VirtualMachineConfig} */
-        public VirtualMachineConfig build() {
+        public @NonNull VirtualMachineConfig build() {
             final String apkPath = mContext.getPackageCodePath();
             final String packageName = mContext.getPackageName();
             Signature[] certs;
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
index 317caee..3654886 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
@@ -16,6 +16,7 @@
 
 package android.system.virtualmachine;
 
+import android.annotation.NonNull;
 import android.content.Context;
 
 import java.lang.ref.WeakReference;
@@ -28,16 +29,16 @@
  * @hide
  */
 public class VirtualMachineManager {
-    private final Context mContext;
+    private final @NonNull Context mContext;
 
-    private VirtualMachineManager(Context context) {
+    private VirtualMachineManager(@NonNull Context context) {
         mContext = context;
     }
 
     static Map<Context, WeakReference<VirtualMachineManager>> sInstances = new WeakHashMap<>();
 
     /** Returns the per-context instance. */
-    public static VirtualMachineManager getInstance(Context context) {
+    public static @NonNull VirtualMachineManager getInstance(@NonNull Context context) {
         synchronized (sInstances) {
             VirtualMachineManager vmm =
                     sInstances.containsKey(context) ? sInstances.get(context).get() : null;
@@ -59,7 +60,8 @@
      * new (and different) virtual machine even if the name and the config are the same as the
      * deleted one.
      */
-    public VirtualMachine create(String name, VirtualMachineConfig config)
+    public @NonNull VirtualMachine create(
+            @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
         synchronized (sCreateLock) {
             return VirtualMachine.create(mContext, name, config);
@@ -70,31 +72,24 @@
      * Returns an existing {@link VirtualMachine} with the given name. Returns null if there is no
      * such virtual machine.
      */
-    public VirtualMachine get(String name) throws VirtualMachineException {
+    public @NonNull VirtualMachine get(@NonNull String name) throws VirtualMachineException {
         return VirtualMachine.load(mContext, name);
     }
 
     /**
-     * Returns an existing {@link VirtualMachine} if it exists, or create a new one. If the virtual
-     * machine exists, and config is not null, the virtual machine is re-configured with the new
-     * config. However, if the config is not compatible with the original config of the virtual
-     * machine, exception is thrown.
+     * Returns an existing {@link VirtualMachine} if it exists, or create a new one. The config
+     * parameter is used only when a new virtual machine is created.
      */
-    public VirtualMachine getOrCreate(String name, VirtualMachineConfig config)
+    public @NonNull VirtualMachine getOrCreate(
+            @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
         VirtualMachine vm;
         synchronized (sCreateLock) {
             vm = get(name);
             if (vm == null) {
-                return create(name, config);
+                vm = create(name, config);
             }
         }
-
-        if (config != null) {
-            // Can throw VirtualMachineException is the new config is not compatible with the
-            // old config.
-            vm.setConfig(config);
-        }
         return vm;
     }
 }
diff --git a/tests/hostside/helper/java/android/virt/test/CommandRunner.java b/tests/hostside/helper/java/android/virt/test/CommandRunner.java
new file mode 100644
index 0000000..696c89a
--- /dev/null
+++ b/tests/hostside/helper/java/android/virt/test/CommandRunner.java
@@ -0,0 +1,89 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.virt.test;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.junit.Assert.fail;
+import static org.junit.Assume.assumeThat;
+
+import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.log.LogUtil.CLog;
+import com.android.tradefed.util.CommandResult;
+import com.android.tradefed.util.CommandStatus;
+
+import java.util.Arrays;
+
+import javax.annotation.Nonnull;
+
+/** A helper class to provide easy way to run commands on a test device. */
+public class CommandRunner {
+
+    /** Default timeout. 30 sec because Microdroid is extremely slow on GCE-on-CF. */
+    private static final long DEFAULT_TIMEOUT = 30000;
+
+    private ITestDevice mDevice;
+
+    public CommandRunner(@Nonnull ITestDevice device) {
+        mDevice = device;
+    }
+
+    public ITestDevice getDevice() {
+        return mDevice;
+    }
+
+    public String run(String... cmd) throws DeviceNotAvailableException {
+        CommandResult result = runForResult(cmd);
+        if (result.getStatus() != CommandStatus.SUCCESS) {
+            fail(join(cmd) + " has failed: " + result);
+        }
+        return result.getStdout().trim();
+    }
+
+    public String tryRun(String... cmd) throws DeviceNotAvailableException {
+        CommandResult result = runForResult(cmd);
+        if (result.getStatus() == CommandStatus.SUCCESS) {
+            return result.getStdout().trim();
+        } else {
+            CLog.d(join(cmd) + " has failed (but ok): " + result);
+            return null;
+        }
+    }
+
+    public String runWithTimeout(long timeoutMillis, String... cmd)
+            throws DeviceNotAvailableException {
+        CommandResult result =
+                mDevice.executeShellV2Command(
+                        join(cmd), timeoutMillis, java.util.concurrent.TimeUnit.MILLISECONDS);
+        if (result.getStatus() != CommandStatus.SUCCESS) {
+            fail(join(cmd) + " has failed: " + result);
+        }
+        return result.getStdout().trim();
+    }
+
+    public CommandResult runForResult(String... cmd) throws DeviceNotAvailableException {
+        return mDevice.executeShellV2Command(join(cmd));
+    }
+
+    public void assumeSuccess(String... cmd) throws DeviceNotAvailableException {
+        assumeThat(runForResult(cmd).getStatus(), is(CommandStatus.SUCCESS));
+    }
+
+    private static String join(String... strs) {
+        return String.join(" ", Arrays.asList(strs));
+    }
+}
diff --git a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
index 2d55a9c..fef8864 100644
--- a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
+++ b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
@@ -20,10 +20,11 @@
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
-import static org.junit.Assume.assumeThat;
 
 import com.android.compatibility.common.tradefed.build.CompatibilityBuildHelper;
+import com.android.tradefed.build.IBuildInfo;
 import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
 import com.android.tradefed.util.CommandResult;
@@ -51,92 +52,63 @@
 
     private static final long MICRODROID_ADB_CONNECT_TIMEOUT_MINUTES = 5;
 
-    public void prepareVirtualizationTestSetup() throws DeviceNotAvailableException {
-        // kill stale crosvm processes
-        tryRunOnAndroid("killall", "crosvm");
+    public static void prepareVirtualizationTestSetup(ITestDevice androidDevice)
+            throws DeviceNotAvailableException {
+        CommandRunner android = new CommandRunner(androidDevice);
 
-        // Prepare the test root
-        tryRunOnAndroid("rm", "-rf", TEST_ROOT);
-        tryRunOnAndroid("mkdir", "-p", TEST_ROOT);
+        // kill stale crosvm processes
+        android.tryRun("killall", "crosvm");
 
         // disconnect from microdroid
         tryRunOnHost("adb", "disconnect", MICRODROID_SERIAL);
     }
 
-    public void cleanUpVirtualizationTestSetup() throws DeviceNotAvailableException {
+    public static void cleanUpVirtualizationTestSetup(ITestDevice androidDevice)
+            throws DeviceNotAvailableException {
+        CommandRunner android = new CommandRunner(androidDevice);
+
         // disconnect from microdroid
         tryRunOnHost("adb", "disconnect", MICRODROID_SERIAL);
 
         // kill stale VMs and directories
-        tryRunOnAndroid("killall", "crosvm");
-        tryRunOnAndroid("rm", "-rf", "/data/misc/virtualizationservice/*");
-        tryRunOnAndroid("stop", "virtualizationservice");
+        android.tryRun("killall", "crosvm");
+        android.tryRun("rm", "-rf", "/data/misc/virtualizationservice/*");
+        android.tryRun("stop", "virtualizationservice");
     }
 
-    public void testIfDeviceIsCapable() throws DeviceNotAvailableException {
+    public static void testIfDeviceIsCapable(ITestDevice androidDevice)
+            throws DeviceNotAvailableException {
+        CommandRunner android = new CommandRunner(androidDevice);
+
         // Checks the preconditions to run microdroid. If the condition is not satisfied
         // don't run the test (instead of failing)
-        skipIfFail("ls /dev/kvm");
-        skipIfFail("ls /dev/vhost-vsock");
-        skipIfFail("ls /apex/com.android.virt/bin/crosvm");
+        android.assumeSuccess("ls /dev/kvm");
+        android.assumeSuccess("ls /dev/vhost-vsock");
+        android.assumeSuccess("ls /apex/com.android.virt/bin/crosvm");
     }
 
     // Run an arbitrary command in the host side and returns the result
-    private String runOnHost(String... cmd) {
+    private static String runOnHost(String... cmd) {
         return runOnHostWithTimeout(10000, cmd);
     }
 
     // Same as runOnHost, but failure is not an error
-    private String tryRunOnHost(String... cmd) {
+    private static String tryRunOnHost(String... cmd) {
         final long timeout = 10000;
         CommandResult result = RunUtil.getDefault().runTimedCmd(timeout, cmd);
         return result.getStdout().trim();
     }
 
     // Same as runOnHost, but with custom timeout
-    private String runOnHostWithTimeout(long timeoutMillis, String... cmd) {
+    private static String runOnHostWithTimeout(long timeoutMillis, String... cmd) {
         assertTrue(timeoutMillis >= 0);
         CommandResult result = RunUtil.getDefault().runTimedCmd(timeoutMillis, cmd);
         assertThat(result.getStatus(), is(CommandStatus.SUCCESS));
         return result.getStdout().trim();
     }
 
-    // Run a shell command on Android. the default timeout is 2 min by tradefed
-    public String runOnAndroid(String... cmd) throws DeviceNotAvailableException {
-        CommandResult result = getDevice().executeShellV2Command(join(cmd));
-        if (result.getStatus() != CommandStatus.SUCCESS) {
-            fail(join(cmd) + " has failed: " + result);
-        }
-        return result.getStdout().trim();
-    }
-
-    // Same as runOnAndroid, but returns null on error.
-    public String tryRunOnAndroid(String... cmd) throws DeviceNotAvailableException {
-        CommandResult result = getDevice().executeShellV2Command(join(cmd));
-        if (result.getStatus() == CommandStatus.SUCCESS) {
-            return result.getStdout().trim();
-        } else {
-            CLog.d(join(cmd) + " has failed (but ok): " + result);
-            return null;
-        }
-    }
-
-    private String runOnAndroidWithTimeout(long timeoutMillis, String... cmd)
-            throws DeviceNotAvailableException {
-        CommandResult result =
-                getDevice()
-                        .executeShellV2Command(
-                                join(cmd),
-                                timeoutMillis,
-                                java.util.concurrent.TimeUnit.MILLISECONDS);
-        if (result.getStatus() != CommandStatus.SUCCESS) {
-            fail(join(cmd) + " has failed: " + result);
-        }
-        return result.getStdout().trim();
-    }
-
     // Run a shell command on Microdroid
-    public String runOnMicrodroid(String... cmd) {
+    public static String runOnMicrodroid(String... cmd) {
         CommandResult result = runOnMicrodroidForResult(cmd);
         if (result.getStatus() != CommandStatus.SUCCESS) {
             fail(join(cmd) + " has failed: " + result);
@@ -145,7 +117,7 @@
     }
 
     // Same as runOnMicrodroid, but returns null on error.
-    public String tryRunOnMicrodroid(String... cmd) {
+    public static String tryRunOnMicrodroid(String... cmd) {
         CommandResult result = runOnMicrodroidForResult(cmd);
         if (result.getStatus() == CommandStatus.SUCCESS) {
             return result.getStdout().trim();
@@ -155,52 +127,63 @@
         }
     }
 
-    public CommandResult runOnMicrodroidForResult(String... cmd) {
+    public static CommandResult runOnMicrodroidForResult(String... cmd) {
         final long timeout = 30000; // 30 sec. Microdroid is extremely slow on GCE-on-CF.
         return RunUtil.getDefault()
                 .runTimedCmd(timeout, "adb", "-s", MICRODROID_SERIAL, "shell", join(cmd));
     }
 
-    private String join(String... strs) {
+    private static String join(String... strs) {
         return String.join(" ", Arrays.asList(strs));
     }
 
     public File findTestFile(String name) {
+        return findTestFile(getBuild(), name);
+    }
+
+    private static File findTestFile(IBuildInfo buildInfo, String name) {
         try {
-            return (new CompatibilityBuildHelper(getBuild())).getTestFile(name);
+            return (new CompatibilityBuildHelper(buildInfo)).getTestFile(name);
         } catch (FileNotFoundException e) {
             fail("Missing test file: " + name);
             return null;
         }
     }
 
-    public String startMicrodroid(
-            String apkName, String packageName, String configPath, boolean debug)
+    public static String startMicrodroid(
+            ITestDevice androidDevice,
+            IBuildInfo buildInfo,
+            String apkName,
+            String packageName,
+            String configPath,
+            boolean debug)
             throws DeviceNotAvailableException {
+        CommandRunner android = new CommandRunner(androidDevice);
+
         // Install APK
-        File apkFile = findTestFile(apkName);
-        getDevice().installPackage(apkFile, /* reinstall */ true);
+        File apkFile = findTestFile(buildInfo, apkName);
+        androidDevice.installPackage(apkFile, /* reinstall */ true);
 
         // 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). So we use the `pm path` command directly.
-        String apkPath = runOnAndroid("pm", "path", packageName);
+        String apkPath = android.run("pm", "path", packageName);
         assertTrue(apkPath.startsWith("package:"));
         apkPath = apkPath.substring("package:".length());
 
         // Push the idsig file to the device
-        File idsigOnHost = findTestFile(apkName + ".idsig");
+        File idsigOnHost = findTestFile(buildInfo, apkName + ".idsig");
         final String apkIdsigPath = TEST_ROOT + apkName + ".idsig";
-        getDevice().pushFile(idsigOnHost, apkIdsigPath);
+        androidDevice.pushFile(idsigOnHost, apkIdsigPath);
 
         final String instanceImg = TEST_ROOT + INSTANCE_IMG;
         final String logPath = TEST_ROOT + "log.txt";
         final String debugFlag = debug ? "--debug " : "";
 
         // Run the VM
-        runOnAndroid("start", "virtualizationservice");
+        android.run("start", "virtualizationservice");
         String ret =
-                runOnAndroid(
+                android.run(
                         VIRT_APEX + "bin/vm",
                         "run-app",
                         "--daemonize",
@@ -217,7 +200,7 @@
                 () -> {
                     try {
                         // Keep redirecting sufficiently long enough
-                        runOnAndroidWithTimeout(
+                        android.runWithTimeout(
                                 MICRODROID_BOOT_TIMEOUT_MINUTES * 60 * 1000,
                                 "logwrapper",
                                 "tail",
@@ -236,17 +219,20 @@
         return matcher.group(1);
     }
 
-    public void shutdownMicrodroid(String cid) throws DeviceNotAvailableException {
+    public static void shutdownMicrodroid(ITestDevice androidDevice, String cid)
+            throws DeviceNotAvailableException {
+        CommandRunner android = new CommandRunner(androidDevice);
+
         // Close the connection before shutting the VM down. Otherwise, b/192660485.
         tryRunOnHost("adb", "disconnect", MICRODROID_SERIAL);
-        final String serial = getDevice().getSerialNumber();
+        final String serial = androidDevice.getSerialNumber();
         tryRunOnHost("adb", "-s", serial, "forward", "--remove", "tcp:" + TEST_VM_ADB_PORT);
 
         // Shutdown the VM
-        runOnAndroid(VIRT_APEX + "bin/vm", "stop", cid);
+        android.run(VIRT_APEX + "bin/vm", "stop", cid);
     }
 
-    public void rootMicrodroid() throws DeviceNotAvailableException {
+    public static void rootMicrodroid() throws DeviceNotAvailableException {
         runOnHost("adb", "-s", MICRODROID_SERIAL, "root");
 
         // TODO(192660959): Figure out the root cause and remove the sleep. For unknown reason,
@@ -254,6 +240,12 @@
         // `adb -s $MICRODROID_SERIAL shell ...` often fails with "adb: device offline".
         try {
             Thread.sleep(1000);
+            runOnHostWithTimeout(
+                    MICRODROID_ADB_CONNECT_TIMEOUT_MINUTES * 60 * 1000,
+                    "adb",
+                    "-s",
+                    MICRODROID_SERIAL,
+                    "wait-for-device");
         } catch (InterruptedException e) {
             Thread.currentThread().interrupt();
         }
@@ -261,12 +253,13 @@
 
     // Establish an adb connection to microdroid by letting Android forward the connection to
     // microdroid. Wait until the connection is established and microdroid is booted.
-    public void adbConnectToMicrodroid(String cid) throws DeviceNotAvailableException {
+    public static void adbConnectToMicrodroid(ITestDevice androidDevice, String cid)
+            throws DeviceNotAvailableException {
         long start = System.currentTimeMillis();
         long timeoutMillis = MICRODROID_ADB_CONNECT_TIMEOUT_MINUTES * 60 * 1000;
         long elapsed = 0;
 
-        final String serial = getDevice().getSerialNumber();
+        final String serial = androidDevice.getSerialNumber();
         final String from = "tcp:" + TEST_VM_ADB_PORT;
         final String to = "vsock:" + cid + ":5555";
         runOnHost("adb", "-s", serial, "forward", from, to);
@@ -300,9 +293,4 @@
         // Check if it actually booted by reading a sysprop.
         assertThat(runOnMicrodroid("getprop", "ro.hardware"), is("microdroid"));
     }
-
-    private void skipIfFail(String command) throws DeviceNotAvailableException {
-        CommandResult result = getDevice().executeShellV2Command(command);
-        assumeThat(result.getStatus(), is(CommandStatus.SUCCESS));
-    }
 }
diff --git a/tests/hostside/java/android/virt/test/MicrodroidTestCase.java b/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
index 78a16c4..8afc287 100644
--- a/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
+++ b/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
@@ -35,8 +35,15 @@
     @Test
     public void testMicrodroidBoots() throws Exception {
         final String configPath = "assets/vm_config.json"; // path inside the APK
-        final String cid = startMicrodroid(APK_NAME, PACKAGE_NAME, configPath, /* debug */ false);
-        adbConnectToMicrodroid(cid);
+        final String cid =
+                startMicrodroid(
+                        getDevice(),
+                        getBuild(),
+                        APK_NAME,
+                        PACKAGE_NAME,
+                        configPath,
+                        /* debug */ false);
+        adbConnectToMicrodroid(getDevice(), cid);
 
         // Test writing to /data partition
         runOnMicrodroid("echo MicrodroidTest > /data/local/tmp/test.txt");
@@ -69,26 +76,27 @@
         // Check that keystore was found by the payload
         assertThat(runOnMicrodroid("getprop", "debug.microdroid.test.keystore"), is("PASS"));
 
-        shutdownMicrodroid(cid);
+        shutdownMicrodroid(getDevice(), cid);
     }
 
     @Test
     public void testDebugMode() throws Exception {
         final String configPath = "assets/vm_config.json"; // path inside the APK
         final boolean debug = true;
-        final String cid = startMicrodroid(APK_NAME, PACKAGE_NAME, configPath, debug);
-        adbConnectToMicrodroid(cid);
+        final String cid =
+                startMicrodroid(getDevice(), getBuild(), APK_NAME, PACKAGE_NAME, configPath, debug);
+        adbConnectToMicrodroid(getDevice(), cid);
 
         assertThat(runOnMicrodroid("getenforce"), is("Permissive"));
 
-        shutdownMicrodroid(cid);
+        shutdownMicrodroid(getDevice(), cid);
     }
 
     @Before
     public void setUp() throws Exception {
-        testIfDeviceIsCapable();
+        testIfDeviceIsCapable(getDevice());
 
-        prepareVirtualizationTestSetup();
+        prepareVirtualizationTestSetup(getDevice());
 
         getDevice().installPackage(findTestFile(APK_NAME), /* reinstall */ false);
 
@@ -98,7 +106,7 @@
 
     @After
     public void shutdown() throws Exception {
-        cleanUpVirtualizationTestSetup();
+        cleanUpVirtualizationTestSetup(getDevice());
 
         getDevice().uninstallPackage(PACKAGE_NAME);
     }
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index a1dba43..239d729 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -33,8 +33,8 @@
         "libonce_cell",
         "libprotobuf",
         "libprotos",
-        "libregex",
         "libserde_json",
+        "libserde_xml_rs",
         "libserde",
         "libshared_child",
         "libuuid",
diff --git a/virtualizationservice/src/payload.rs b/virtualizationservice/src/payload.rs
index 76c55de..1df537c 100644
--- a/virtualizationservice/src/payload.rs
+++ b/virtualizationservice/src/payload.rs
@@ -16,27 +16,32 @@
 
 use crate::composite::align_to_partition_size;
 
-use anyhow::{anyhow, bail, Result};
+use anyhow::{anyhow, Context, Result};
 use microdroid_metadata::{ApexPayload, ApkPayload, Metadata};
 use microdroid_payload_config::ApexConfig;
 use once_cell::sync::OnceCell;
-use regex::Regex;
+use serde::Deserialize;
+use serde_xml_rs::from_reader;
 use std::fs;
-use std::fs::OpenOptions;
-use std::io::{BufRead, BufReader, Seek, SeekFrom, Write};
+use std::fs::{File, OpenOptions};
+use std::io::{Seek, SeekFrom, Write};
 use std::path::{Path, PathBuf};
-use std::process::Command;
 use vmconfig::{DiskImage, Partition};
 
+const APEX_INFO_LIST_PATH: &str = "/apex/apex-info-list.xml";
+
 /// Represents the list of APEXes
-#[derive(Debug)]
+#[derive(Debug, Deserialize)]
 struct ApexInfoList {
+    #[serde(rename = "apex-info")]
     list: Vec<ApexInfo>,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Deserialize)]
 struct ApexInfo {
+    #[serde(rename = "moduleName")]
     name: String,
+    #[serde(rename = "modulePath")]
     path: PathBuf,
 }
 
@@ -45,34 +50,11 @@
     fn load() -> Result<&'static ApexInfoList> {
         static INSTANCE: OnceCell<ApexInfoList> = OnceCell::new();
         INSTANCE.get_or_try_init(|| {
-            // TODO(b/191601801): look up /apex/apex-info-list.xml instead of apexservice
-            // Each APEX prints the line:
-            //   Module: <...> Version: <...> VersionName: <...> Path: <...> IsActive: <...> IsFactory: <...>
-            // We only care about "Module:" and "Path:" tagged values for now.
-            let info_pattern =
-                Regex::new(r"^Module: (?P<name>[^ ]*) .* Path: (?P<path>[^ ]*) .*$")?;
-            let output = Command::new("cmd")
-                .arg("-w")
-                .arg("apexservice")
-                .arg("getActivePackages")
-                .output()
-                .expect("failed to execute apexservice cmd");
-            let list = BufReader::new(output.stdout.as_slice())
-                .lines()
-                .map(|line| -> Result<ApexInfo> {
-                    let line = line?;
-                    let captures = info_pattern
-                        .captures(&line)
-                        .ok_or_else(|| anyhow!("can't parse: {}", line))?;
-                    let name = captures.name("name").unwrap();
-                    let path = captures.name("path").unwrap();
-                    Ok(ApexInfo { name: name.as_str().to_owned(), path: path.as_str().into() })
-                })
-                .collect::<Result<Vec<ApexInfo>>>()?;
-            if list.is_empty() {
-                bail!("failed to load apex info: empty");
-            }
-            Ok(ApexInfoList { list })
+            let apex_info_list = File::open(APEX_INFO_LIST_PATH)
+                .context(format!("Failed to open {}", APEX_INFO_LIST_PATH))?;
+            let apex_info_list: ApexInfoList = from_reader(apex_info_list)
+                .context(format!("Failed to parse {}", APEX_INFO_LIST_PATH))?;
+            Ok(apex_info_list)
         })
     }