Merge "Add MicrodroidTestApp to VTS" into main
diff --git a/java/framework/Android.bp b/java/framework/Android.bp
index 32b2aee..26ea214 100644
--- a/java/framework/Android.bp
+++ b/java/framework/Android.bp
@@ -43,4 +43,7 @@
     impl_library_visibility: [
         "//packages/modules/Virtualization:__subpackages__",
     ],
+    aconfig_declarations: [
+        "avf_aconfig_flags",
+    ],
 }
diff --git a/tests/pvmfw/java/com/android/pvmfw/test/CustomPvmfwHostTestCaseBase.java b/tests/pvmfw/java/com/android/pvmfw/test/CustomPvmfwHostTestCaseBase.java
new file mode 100644
index 0000000..d9d425a
--- /dev/null
+++ b/tests/pvmfw/java/com/android/pvmfw/test/CustomPvmfwHostTestCaseBase.java
@@ -0,0 +1,165 @@
+/*
+ * Copyright 2024 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 com.android.pvmfw.test;
+
+import static com.android.tradefed.device.TestDevice.MicrodroidBuilder;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.junit.Assume.assumeTrue;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import com.android.microdroid.test.host.MicrodroidHostTestCaseBase;
+import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.device.TestDevice;
+import com.android.tradefed.util.FileUtil;
+
+import org.junit.After;
+import org.junit.Before;
+
+import java.io.File;
+import java.util.Objects;
+import java.util.Map;
+
+/** Base class for testing custom pvmfw */
+public class CustomPvmfwHostTestCaseBase extends MicrodroidHostTestCaseBase {
+    @NonNull public static final String PVMFW_FILE_NAME = "pvmfw_test.bin";
+    @NonNull public static final String BCC_FILE_NAME = "bcc.dat";
+    @NonNull public static final String PACKAGE_FILE_NAME = "MicrodroidTestApp.apk";
+    @NonNull public static final String PACKAGE_NAME = "com.android.microdroid.test";
+    @NonNull public static final String MICRODROID_DEBUG_FULL = "full";
+    @NonNull public static final String MICRODROID_DEBUG_NONE = "none";
+
+    @NonNull
+    public static final String MICRODROID_CONFIG_PATH = "assets/microdroid/vm_config_apex.json";
+
+    @NonNull public static final String MICRODROID_LOG_PATH = TEST_ROOT + "log.txt";
+    public static final int BOOT_COMPLETE_TIMEOUT_MS = 30000; // 30 seconds
+    public static final int BOOT_FAILURE_WAIT_TIME_MS = 10000; // 10 seconds
+    public static final int CONSOLE_OUTPUT_WAIT_MS = 5000; // 5 seconds
+
+    @NonNull public static final String CUSTOM_PVMFW_FILE_PREFIX = "pvmfw";
+    @NonNull public static final String CUSTOM_PVMFW_FILE_SUFFIX = ".bin";
+    @NonNull public static final String CUSTOM_PVMFW_IMG_PATH = TEST_ROOT + PVMFW_FILE_NAME;
+    @NonNull public static final String CUSTOM_PVMFW_IMG_PATH_PROP = "hypervisor.pvmfw.path";
+
+    @Nullable private static File mPvmfwBinFileOnHost;
+    @Nullable private static File mBccFileOnHost;
+
+    @Nullable private TestDevice mAndroidDevice;
+    @Nullable private ITestDevice mMicrodroidDevice;
+
+    @Nullable private File mCustomPvmfwFileOnHost;
+
+    @Before
+    public void setUp() throws Exception {
+        mAndroidDevice = (TestDevice) Objects.requireNonNull(getDevice());
+
+        // Check device capabilities
+        assumeDeviceIsCapable(mAndroidDevice);
+        assumeTrue(
+                "Skip if protected VMs are not supported",
+                mAndroidDevice.supportsMicrodroid(/* protectedVm= */ true));
+
+        // tradefed copies the test artifacts under /tmp when running tests,
+        // so we should *find* the artifacts with the file name.
+        mPvmfwBinFileOnHost =
+                getTestInformation().getDependencyFile(PVMFW_FILE_NAME, /* targetFirst= */ false);
+        mBccFileOnHost =
+                getTestInformation().getDependencyFile(BCC_FILE_NAME, /* targetFirst= */ false);
+
+        // Prepare for system properties for custom pvmfw.img.
+        // File will be prepared later in individual test and then pushed to device
+        // when launching with launchProtectedVmAndWaitForBootCompleted().
+        mCustomPvmfwFileOnHost =
+                FileUtil.createTempFile(CUSTOM_PVMFW_FILE_PREFIX, CUSTOM_PVMFW_FILE_SUFFIX);
+        setPropertyOrThrow(mAndroidDevice, CUSTOM_PVMFW_IMG_PATH_PROP, CUSTOM_PVMFW_IMG_PATH);
+
+        // Prepare for launching microdroid
+        mAndroidDevice.installPackage(findTestFile(PACKAGE_FILE_NAME), /* reinstall */ false);
+        prepareVirtualizationTestSetup(mAndroidDevice);
+        mMicrodroidDevice = null;
+    }
+
+    @After
+    public void shutdown() throws Exception {
+        if (!mAndroidDevice.supportsMicrodroid(/* protectedVm= */ true)) {
+            return;
+        }
+        if (mMicrodroidDevice != null) {
+            mAndroidDevice.shutdownMicrodroid(mMicrodroidDevice);
+            mMicrodroidDevice = null;
+        }
+        mAndroidDevice.uninstallPackage(PACKAGE_NAME);
+
+        // Cleanup for custom pvmfw.img
+        setPropertyOrThrow(mAndroidDevice, CUSTOM_PVMFW_IMG_PATH_PROP, "");
+        FileUtil.deleteFile(mCustomPvmfwFileOnHost);
+
+        cleanUpVirtualizationTestSetup(mAndroidDevice);
+    }
+
+    /** Returns pvmfw.bin file on host for building custom pvmfw with */
+    public File getPvmfwBinFile() {
+        return mPvmfwBinFileOnHost;
+    }
+
+    /** Returns BCC file on host for building custom pvmfw with */
+    public File getBccFile() {
+        return mBccFileOnHost;
+    }
+
+    /**
+     * Returns a custom pvmfw file.
+     *
+     * <p>This is a temporary file on host. The file should been prepared as a custom pvmfw because
+     * calling {@link #launchProtectedVmAndWaitForBootCompleted}, so virtualization manager can read
+     * the file path from sysprop and boot pVM with it.
+     */
+    public File getCustomPvmfwFile() {
+        return mCustomPvmfwFileOnHost;
+    }
+
+    /**
+     * Launches protected VM with custom pvmfw ({@link #getCustomPvmfwFile}) and wait for boot
+     * completed. Throws exception when boot failed.
+     */
+    public ITestDevice launchProtectedVmAndWaitForBootCompleted(
+            String debugLevel, long adbTimeoutMs, @NonNull Map<String, File> bootFiles)
+            throws DeviceNotAvailableException {
+        MicrodroidBuilder builder =
+                MicrodroidBuilder.fromDevicePath(
+                                getPathForPackage(PACKAGE_NAME), MICRODROID_CONFIG_PATH)
+                        .debugLevel(debugLevel)
+                        .protectedVm(/* protectedVm= */ true)
+                        .addBootFile(mCustomPvmfwFileOnHost, PVMFW_FILE_NAME)
+                        .setAdbConnectTimeoutMs(adbTimeoutMs);
+        for (String name : bootFiles.keySet()) {
+            File file = bootFiles.get(name);
+            builder.addBootFile(file, name);
+        }
+
+        mMicrodroidDevice = builder.build(mAndroidDevice);
+
+        assertThat(mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT_MS)).isTrue();
+        assertThat(mMicrodroidDevice.enableAdbRoot()).isTrue();
+        return mMicrodroidDevice;
+    }
+}
diff --git a/tests/pvmfw/java/com/android/pvmfw/test/DebugPolicyHostTests.java b/tests/pvmfw/java/com/android/pvmfw/test/DebugPolicyHostTests.java
index 26f5993..803405d 100644
--- a/tests/pvmfw/java/com/android/pvmfw/test/DebugPolicyHostTests.java
+++ b/tests/pvmfw/java/com/android/pvmfw/test/DebugPolicyHostTests.java
@@ -16,28 +16,22 @@
 
 package com.android.pvmfw.test;
 
-import static com.android.tradefed.device.TestDevice.MicrodroidBuilder;
-
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
-import static org.junit.Assume.assumeTrue;
 import static org.junit.Assert.assertThrows;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import com.android.microdroid.test.host.CommandRunner;
-import com.android.microdroid.test.host.MicrodroidHostTestCaseBase;
 import com.android.pvmfw.test.host.Pvmfw;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.DeviceRuntimeException;
 import com.android.tradefed.device.ITestDevice;
-import com.android.tradefed.device.TestDevice;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.util.CommandStatus;
 import com.android.tradefed.util.CommandResult;
-import com.android.tradefed.util.FileUtil;
 
 import org.junit.After;
 import org.junit.Before;
@@ -45,32 +39,13 @@
 import org.junit.runner.RunWith;
 
 import java.io.File;
-import java.util.Objects;
+import java.util.Collections;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 /** Tests debug policy */
 @RunWith(DeviceJUnit4ClassRunner.class)
-public class DebugPolicyHostTests extends MicrodroidHostTestCaseBase {
-    @NonNull private static final String PVMFW_FILE_NAME = "pvmfw_test.bin";
-    @NonNull private static final String BCC_FILE_NAME = "bcc.dat";
-    @NonNull private static final String PACKAGE_FILE_NAME = "MicrodroidTestApp.apk";
-    @NonNull private static final String PACKAGE_NAME = "com.android.microdroid.test";
-    @NonNull private static final String MICRODROID_DEBUG_FULL = "full";
-    @NonNull private static final String MICRODROID_DEBUG_NONE = "none";
-
-    @NonNull
-    private static final String MICRODROID_CONFIG_PATH = "assets/microdroid/vm_config_apex.json";
-
-    @NonNull private static final String MICRODROID_LOG_PATH = TEST_ROOT + "log.txt";
-    private static final int BOOT_COMPLETE_TIMEOUT_MS = 30000; // 30 seconds
-    private static final int BOOT_FAILURE_WAIT_TIME_MS = 10000; // 10 seconds
-    private static final int CONSOLE_OUTPUT_WAIT_MS = 5000; // 5 seconds
-
-    @NonNull private static final String CUSTOM_PVMFW_FILE_PREFIX = "pvmfw";
-    @NonNull private static final String CUSTOM_PVMFW_FILE_SUFFIX = ".bin";
-    @NonNull private static final String CUSTOM_PVMFW_IMG_PATH = TEST_ROOT + PVMFW_FILE_NAME;
-    @NonNull private static final String CUSTOM_PVMFW_IMG_PATH_PROP = "hypervisor.pvmfw.path";
-
+public class DebugPolicyHostTests extends CustomPvmfwHostTestCaseBase {
     @NonNull private static final String CUSTOM_DEBUG_POLICY_FILE_NAME = "debug_policy.dtb";
 
     @NonNull
@@ -98,63 +73,22 @@
     @NonNull private static final String HEX_STRING_ZERO = "00000000";
     @NonNull private static final String HEX_STRING_ONE = "00000001";
 
-    @Nullable private static File mPvmfwBinFileOnHost;
-    @Nullable private static File mBccFileOnHost;
-
-    @Nullable private TestDevice mAndroidDevice;
-    @Nullable private ITestDevice mMicrodroidDevice;
-    @Nullable private File mCustomPvmfwBinFileOnHost;
     @Nullable private File mCustomDebugPolicyFileOnHost;
 
     @Before
     public void setUp() throws Exception {
-        mAndroidDevice = (TestDevice) Objects.requireNonNull(getDevice());
+        super.setUp();
 
-        // Check device capabilities
-        assumeDeviceIsCapable(mAndroidDevice);
-        assumeTrue(
-                "Skip if protected VMs are not supported",
-                mAndroidDevice.supportsMicrodroid(/* protectedVm= */ true));
-
-        // tradefed copies the test artfacts under /tmp when running tests,
-        // so we should *find* the artifacts with the file name.
-        mPvmfwBinFileOnHost =
-                getTestInformation().getDependencyFile(PVMFW_FILE_NAME, /* targetFirst= */ false);
-        mBccFileOnHost =
-                getTestInformation().getDependencyFile(BCC_FILE_NAME, /* targetFirst= */ false);
-
-        // Prepare for system properties for custom debug policy.
-        // File will be prepared later in individual test by setupCustomDebugPolicy()
-        // and then pushed to device when launching with launchProtectedVmAndWaitForBootCompleted()
-        // or tryLaunchProtectedNonDebuggableVm().
-        mCustomPvmfwBinFileOnHost =
-                FileUtil.createTempFile(CUSTOM_PVMFW_FILE_PREFIX, CUSTOM_PVMFW_FILE_SUFFIX);
-        setPropertyOrThrow(mAndroidDevice, CUSTOM_PVMFW_IMG_PATH_PROP, CUSTOM_PVMFW_IMG_PATH);
-        setPropertyOrThrow(mAndroidDevice, CUSTOM_DEBUG_POLICY_PATH_PROP, CUSTOM_DEBUG_POLICY_PATH);
-
-        // Prepare for launching microdroid
-        mAndroidDevice.installPackage(findTestFile(PACKAGE_FILE_NAME), /* reinstall */ false);
-        prepareVirtualizationTestSetup(mAndroidDevice);
-        mMicrodroidDevice = null;
+        // Prepare system properties for custom debug policy.
+        setPropertyOrThrow(getDevice(), CUSTOM_DEBUG_POLICY_PATH_PROP, CUSTOM_DEBUG_POLICY_PATH);
     }
 
     @After
     public void shutdown() throws Exception {
-        if (!mAndroidDevice.supportsMicrodroid(/* protectedVm= */ true)) {
-            return;
-        }
-        if (mMicrodroidDevice != null) {
-            mAndroidDevice.shutdownMicrodroid(mMicrodroidDevice);
-            mMicrodroidDevice = null;
-        }
-        mAndroidDevice.uninstallPackage(PACKAGE_NAME);
+        super.shutdown();
 
         // Cleanup for custom debug policies
-        setPropertyOrThrow(mAndroidDevice, CUSTOM_DEBUG_POLICY_PATH_PROP, "");
-        setPropertyOrThrow(mAndroidDevice, CUSTOM_PVMFW_IMG_PATH_PROP, "");
-        FileUtil.deleteFile(mCustomPvmfwBinFileOnHost);
-
-        cleanUpVirtualizationTestSetup(mAndroidDevice);
+        setPropertyOrThrow(getDevice(), CUSTOM_DEBUG_POLICY_PATH_PROP, "");
     }
 
     @Test
@@ -198,43 +132,41 @@
     @Test
     public void testRamdumpInDebugPolicy_withDebugLevelNone_hasRamdumpArgs() throws Exception {
         prepareCustomDebugPolicy("avf_debug_policy_with_ramdump.dtbo");
-        mMicrodroidDevice = launchProtectedVmAndWaitForBootCompleted(MICRODROID_DEBUG_NONE);
+        ITestDevice device = launchProtectedVmAndWaitForBootCompleted(MICRODROID_DEBUG_NONE);
 
-        assertThat(readMicrodroidFileAsString(MICRODROID_CMDLINE_PATH)).contains("crashkernel=");
-        assertThat(readMicrodroidFileAsString(MICRODROID_DT_BOOTARGS_PATH))
-                .contains("crashkernel=");
-        assertThat(readMicrodroidFileAsHexString(MICRODROID_DT_RAMDUMP_PATH))
+        assertThat(readFileAsString(device, MICRODROID_CMDLINE_PATH)).contains("crashkernel=");
+        assertThat(readFileAsString(device, MICRODROID_DT_BOOTARGS_PATH)).contains("crashkernel=");
+        assertThat(readFileAsHexString(device, MICRODROID_DT_RAMDUMP_PATH))
                 .isEqualTo(HEX_STRING_ONE);
     }
 
     @Test
     public void testNoRamdumpInDebugPolicy_withDebugLevelNone_noRamdumpArgs() throws Exception {
         prepareCustomDebugPolicy("avf_debug_policy_without_ramdump.dtbo");
-        mMicrodroidDevice = launchProtectedVmAndWaitForBootCompleted(MICRODROID_DEBUG_NONE);
+        ITestDevice device = launchProtectedVmAndWaitForBootCompleted(MICRODROID_DEBUG_NONE);
 
-        assertThat(readMicrodroidFileAsString(MICRODROID_CMDLINE_PATH))
+        assertThat(readFileAsString(device, MICRODROID_CMDLINE_PATH))
                 .doesNotContain("crashkernel=");
-        assertThat(readMicrodroidFileAsString(MICRODROID_DT_BOOTARGS_PATH))
+        assertThat(readFileAsString(device, MICRODROID_DT_BOOTARGS_PATH))
                 .doesNotContain("crashkernel=");
-        assertThat(readMicrodroidFileAsHexString(MICRODROID_DT_RAMDUMP_PATH))
+        assertThat(readFileAsHexString(device, MICRODROID_DT_RAMDUMP_PATH))
                 .isEqualTo(HEX_STRING_ZERO);
     }
 
     @Test
     public void testNoRamdumpInDebugPolicy_withDebugLevelFull_hasRamdumpArgs() throws Exception {
         prepareCustomDebugPolicy("avf_debug_policy_without_ramdump.dtbo");
-        mMicrodroidDevice = launchProtectedVmAndWaitForBootCompleted(MICRODROID_DEBUG_FULL);
+        ITestDevice device = launchProtectedVmAndWaitForBootCompleted(MICRODROID_DEBUG_FULL);
 
-        assertThat(readMicrodroidFileAsString(MICRODROID_CMDLINE_PATH)).contains("crashkernel=");
-        assertThat(readMicrodroidFileAsString(MICRODROID_DT_BOOTARGS_PATH))
-                .contains("crashkernel=");
-        assertThat(readMicrodroidFileAsHexString(MICRODROID_DT_RAMDUMP_PATH))
+        assertThat(readFileAsString(device, MICRODROID_CMDLINE_PATH)).contains("crashkernel=");
+        assertThat(readFileAsString(device, MICRODROID_DT_BOOTARGS_PATH)).contains("crashkernel=");
+        assertThat(readFileAsHexString(device, MICRODROID_DT_RAMDUMP_PATH))
                 .isEqualTo(HEX_STRING_ZERO);
     }
 
     private boolean isDebugPolicyEnabled(@NonNull String dtPropertyPath)
             throws DeviceNotAvailableException {
-        CommandRunner runner = new CommandRunner(mAndroidDevice);
+        CommandRunner runner = new CommandRunner(getDevice());
         CommandResult result =
                 runner.runForResult("xxd", "-p", "/proc/device-tree" + dtPropertyPath);
         if (result.getStatus() == CommandStatus.SUCCESS) {
@@ -244,15 +176,15 @@
     }
 
     @NonNull
-    private String readMicrodroidFileAsString(@NonNull String path)
+    private String readFileAsString(@NonNull ITestDevice device, @NonNull String path)
             throws DeviceNotAvailableException {
-        return new CommandRunner(mMicrodroidDevice).run("cat", path);
+        return new CommandRunner(device).run("cat", path);
     }
 
     @NonNull
-    private String readMicrodroidFileAsHexString(@NonNull String path)
+    private String readFileAsHexString(@NonNull ITestDevice device, @NonNull String path)
             throws DeviceNotAvailableException {
-        return new CommandRunner(mMicrodroidDevice).run("xxd", "-p", path);
+        return new CommandRunner(device).run("xxd", "-p", path);
     }
 
     private void prepareCustomDebugPolicy(@NonNull String debugPolicyFileName) throws Exception {
@@ -261,10 +193,10 @@
                         .getDependencyFile(debugPolicyFileName, /* targetFirst= */ false);
 
         Pvmfw pvmfw =
-                new Pvmfw.Builder(mPvmfwBinFileOnHost, mBccFileOnHost)
+                new Pvmfw.Builder(getPvmfwBinFile(), getBccFile())
                         .setDebugPolicyOverlay(mCustomDebugPolicyFileOnHost)
                         .build();
-        pvmfw.serialize(mCustomPvmfwBinFileOnHost);
+        pvmfw.serialize(getCustomPvmfwFile());
     }
 
     private boolean hasConsoleOutput(@NonNull CommandResult result)
@@ -274,29 +206,22 @@
 
     private boolean hasMicrodroidLogcatOutput() throws DeviceNotAvailableException {
         CommandResult result =
-                new CommandRunner(mAndroidDevice).runForResult("test", "-s", MICRODROID_LOG_PATH);
+                new CommandRunner(getDevice()).runForResult("test", "-s", MICRODROID_LOG_PATH);
         return result.getExitCode() == 0;
     }
 
-    private ITestDevice launchProtectedVmAndWaitForBootCompleted(String debugLevel)
+    public ITestDevice launchProtectedVmAndWaitForBootCompleted(String debugLevel)
             throws DeviceNotAvailableException {
         return launchProtectedVmAndWaitForBootCompleted(debugLevel, BOOT_COMPLETE_TIMEOUT_MS);
     }
 
-    private ITestDevice launchProtectedVmAndWaitForBootCompleted(
+    public ITestDevice launchProtectedVmAndWaitForBootCompleted(
             String debugLevel, long adbTimeoutMs) throws DeviceNotAvailableException {
-        mMicrodroidDevice =
-                MicrodroidBuilder.fromDevicePath(
-                                getPathForPackage(PACKAGE_NAME), MICRODROID_CONFIG_PATH)
-                        .debugLevel(debugLevel)
-                        .protectedVm(/* protectedVm= */ true)
-                        .addBootFile(mCustomPvmfwBinFileOnHost, PVMFW_FILE_NAME)
-                        .addBootFile(mCustomDebugPolicyFileOnHost, CUSTOM_DEBUG_POLICY_FILE_NAME)
-                        .setAdbConnectTimeoutMs(adbTimeoutMs)
-                        .build(mAndroidDevice);
-        assertThat(mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT_MS)).isTrue();
-        assertThat(mMicrodroidDevice.enableAdbRoot()).isTrue();
-        return mMicrodroidDevice;
+        Map<String, File> bootFiles =
+                Collections.singletonMap(
+                        CUSTOM_DEBUG_POLICY_FILE_NAME, mCustomDebugPolicyFileOnHost);
+
+        return launchProtectedVmAndWaitForBootCompleted(debugLevel, adbTimeoutMs, bootFiles);
     }
 
     // Try to launch protected non-debuggable VM for a while and quit.
@@ -304,10 +229,10 @@
     private CommandResult tryLaunchProtectedNonDebuggableVm() throws Exception {
         // Can't use MicrodroidBuilder because it expects adb connection
         // but non-debuggable VM may not enable adb.
-        CommandRunner runner = new CommandRunner(mAndroidDevice);
+        CommandRunner runner = new CommandRunner(getDevice());
         runner.run("mkdir", "-p", TEST_ROOT);
-        mAndroidDevice.pushFile(mCustomPvmfwBinFileOnHost, CUSTOM_PVMFW_IMG_PATH);
-        mAndroidDevice.pushFile(mCustomDebugPolicyFileOnHost, CUSTOM_DEBUG_POLICY_PATH);
+        getDevice().pushFile(getCustomPvmfwFile(), CUSTOM_PVMFW_IMG_PATH);
+        getDevice().pushFile(mCustomDebugPolicyFileOnHost, CUSTOM_DEBUG_POLICY_PATH);
 
         // This will fail because app wouldn't finish itself.
         // But let's run the app once and get logs.
@@ -327,7 +252,11 @@
         if (isFeatureEnabled("com.android.kvm.LLPVM_CHANGES")) {
             command = String.join(" ", command, "--instance-id-file", TEST_ROOT + "instance_id");
         }
-        return mAndroidDevice.executeShellV2Command(
-                command, CONSOLE_OUTPUT_WAIT_MS, TimeUnit.MILLISECONDS, /* retryAttempts= */ 0);
+        return getDevice()
+                .executeShellV2Command(
+                        command,
+                        CONSOLE_OUTPUT_WAIT_MS,
+                        TimeUnit.MILLISECONDS,
+                        /* retryAttempts= */ 0);
     }
 }
diff --git a/tests/pvmfw/java/com/android/pvmfw/test/PvmfwImgTest.java b/tests/pvmfw/java/com/android/pvmfw/test/PvmfwImgTest.java
index 9fbbd87..b68316d 100644
--- a/tests/pvmfw/java/com/android/pvmfw/test/PvmfwImgTest.java
+++ b/tests/pvmfw/java/com/android/pvmfw/test/PvmfwImgTest.java
@@ -16,124 +16,36 @@
 
 package com.android.pvmfw.test;
 
-import static com.android.tradefed.device.TestDevice.MicrodroidBuilder;
-
-import static com.google.common.truth.Truth.assertThat;
-
-import static org.junit.Assume.assumeTrue;
 import static org.junit.Assert.assertThrows;
 
-import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
-
-import com.android.microdroid.test.host.MicrodroidHostTestCaseBase;
 import com.android.pvmfw.test.host.Pvmfw;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.DeviceRuntimeException;
 import com.android.tradefed.device.ITestDevice;
-import com.android.tradefed.device.TestDevice;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
-import com.android.tradefed.util.FileUtil;
 
-import org.junit.After;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import java.io.File;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
-import java.util.Objects;
 
 /** Tests pvmfw.img and pvmfw */
 @RunWith(DeviceJUnit4ClassRunner.class)
-public class PvmfwImgTest extends MicrodroidHostTestCaseBase {
-    @NonNull private static final String PVMFW_FILE_NAME = "pvmfw_test.bin";
-    @NonNull private static final String BCC_FILE_NAME = "bcc.dat";
-    @NonNull private static final String PACKAGE_FILE_NAME = "MicrodroidTestApp.apk";
-    @NonNull private static final String PACKAGE_NAME = "com.android.microdroid.test";
-    @NonNull private static final String MICRODROID_DEBUG_FULL = "full";
-
-    @NonNull
-    private static final String MICRODROID_CONFIG_PATH = "assets/microdroid/vm_config_apex.json";
-
-    private static final int BOOT_COMPLETE_TIMEOUT_MS = 30000; // 30 seconds
-    private static final int BOOT_FAILURE_WAIT_TIME_MS = 10000; // 10 seconds
-
-    @NonNull private static final String CUSTOM_PVMFW_FILE_PREFIX = "pvmfw";
-    @NonNull private static final String CUSTOM_PVMFW_FILE_SUFFIX = ".bin";
-    @NonNull private static final String CUSTOM_PVMFW_IMG_PATH = TEST_ROOT + PVMFW_FILE_NAME;
-    @NonNull private static final String CUSTOM_PVMFW_IMG_PATH_PROP = "hypervisor.pvmfw.path";
-
-    @Nullable private static File mPvmfwBinFileOnHost;
-    @Nullable private static File mBccFileOnHost;
-
-    @Nullable private TestDevice mAndroidDevice;
-    @Nullable private ITestDevice mMicrodroidDevice;
-    @Nullable private File mCustomPvmfwBinFileOnHost;
-
-    @Before
-    public void setUp() throws Exception {
-        mAndroidDevice = (TestDevice) Objects.requireNonNull(getDevice());
-
-        // Check device capabilities
-        assumeDeviceIsCapable(mAndroidDevice);
-        assumeTrue(
-                "Skip if protected VMs are not supported",
-                mAndroidDevice.supportsMicrodroid(/* protectedVm= */ true));
-
-        // tradefed copies the test artfacts under /tmp when running tests,
-        // so we should *find* the artifacts with the file name.
-        mPvmfwBinFileOnHost =
-                getTestInformation().getDependencyFile(PVMFW_FILE_NAME, /* targetFirst= */ false);
-        mBccFileOnHost =
-                getTestInformation().getDependencyFile(BCC_FILE_NAME, /* targetFirst= */ false);
-
-        // Prepare for system properties for custom pvmfw.img.
-        // File will be prepared later in individual test and then pushed to device
-        // when launching with launchProtectedVmAndWaitForBootCompleted().
-        mCustomPvmfwBinFileOnHost =
-                FileUtil.createTempFile(CUSTOM_PVMFW_FILE_PREFIX, CUSTOM_PVMFW_FILE_SUFFIX);
-        setPropertyOrThrow(mAndroidDevice, CUSTOM_PVMFW_IMG_PATH_PROP, CUSTOM_PVMFW_IMG_PATH);
-
-        // Prepare for launching microdroid
-        mAndroidDevice.installPackage(findTestFile(PACKAGE_FILE_NAME), /* reinstall */ false);
-        prepareVirtualizationTestSetup(mAndroidDevice);
-        mMicrodroidDevice = null;
-    }
-
-    @After
-    public void shutdown() throws Exception {
-        if (!mAndroidDevice.supportsMicrodroid(/* protectedVm= */ true)) {
-            return;
-        }
-        if (mMicrodroidDevice != null) {
-            mAndroidDevice.shutdownMicrodroid(mMicrodroidDevice);
-            mMicrodroidDevice = null;
-        }
-        mAndroidDevice.uninstallPackage(PACKAGE_NAME);
-
-        // Cleanup for custom pvmfw.img
-        setPropertyOrThrow(mAndroidDevice, CUSTOM_PVMFW_IMG_PATH_PROP, "");
-        FileUtil.deleteFile(mCustomPvmfwBinFileOnHost);
-
-        cleanUpVirtualizationTestSetup(mAndroidDevice);
-    }
-
+public class PvmfwImgTest extends CustomPvmfwHostTestCaseBase {
     @Test
     public void testConfigVersion1_0_boots() throws Exception {
-        Pvmfw pvmfw =
-                new Pvmfw.Builder(mPvmfwBinFileOnHost, mBccFileOnHost).setVersion(1, 0).build();
-        pvmfw.serialize(mCustomPvmfwBinFileOnHost);
+        Pvmfw pvmfw = new Pvmfw.Builder(getPvmfwBinFile(), getBccFile()).setVersion(1, 0).build();
+        pvmfw.serialize(getCustomPvmfwFile());
 
         launchProtectedVmAndWaitForBootCompleted(BOOT_COMPLETE_TIMEOUT_MS);
     }
 
     @Test
     public void testConfigVersion1_1_boots() throws Exception {
-        Pvmfw pvmfw =
-                new Pvmfw.Builder(mPvmfwBinFileOnHost, mBccFileOnHost).setVersion(1, 1).build();
-        pvmfw.serialize(mCustomPvmfwBinFileOnHost);
+        Pvmfw pvmfw = new Pvmfw.Builder(getPvmfwBinFile(), getBccFile()).setVersion(1, 1).build();
+        pvmfw.serialize(getCustomPvmfwFile());
 
         launchProtectedVmAndWaitForBootCompleted(BOOT_COMPLETE_TIMEOUT_MS);
     }
@@ -153,7 +65,7 @@
                         new int[] {0xFFFF, 1},
                         new int[] {0xFFFF, 0xFFFF});
 
-        Pvmfw.Builder builder = new Pvmfw.Builder(mPvmfwBinFileOnHost, mBccFileOnHost);
+        Pvmfw.Builder builder = new Pvmfw.Builder(getPvmfwBinFile(), getBccFile());
 
         for (int[] pair : invalid_versions) {
             int major = pair[0];
@@ -161,7 +73,7 @@
             String version = "v" + major + "." + minor;
 
             Pvmfw pvmfw = builder.setVersion(major, minor).build();
-            pvmfw.serialize(mCustomPvmfwBinFileOnHost);
+            pvmfw.serialize(getCustomPvmfwFile());
 
             assertThrows(
                     "pvmfw shouldn't boot with invalid version " + version,
@@ -170,17 +82,9 @@
         }
     }
 
-    private ITestDevice launchProtectedVmAndWaitForBootCompleted(long adbTimeoutMs)
+    public ITestDevice launchProtectedVmAndWaitForBootCompleted(long adbTimeoutMs)
             throws DeviceNotAvailableException {
-        mMicrodroidDevice =
-                MicrodroidBuilder.fromDevicePath(
-                                getPathForPackage(PACKAGE_NAME), MICRODROID_CONFIG_PATH)
-                        .debugLevel(MICRODROID_DEBUG_FULL)
-                        .protectedVm(true)
-                        .addBootFile(mCustomPvmfwBinFileOnHost, PVMFW_FILE_NAME)
-                        .setAdbConnectTimeoutMs(adbTimeoutMs)
-                        .build(mAndroidDevice);
-        assertThat(mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT_MS)).isTrue();
-        return mMicrodroidDevice;
+        return launchProtectedVmAndWaitForBootCompleted(
+                MICRODROID_DEBUG_FULL, adbTimeoutMs, Collections.emptyMap());
     }
 }
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index 278365c..22bea58 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -434,11 +434,7 @@
                 None
             };
 
-        let debug_level = match config {
-            VirtualMachineConfig::AppConfig(config) => config.debugLevel,
-            _ => DebugLevel::NONE,
-        };
-        let debug_config = DebugConfig::new(debug_level);
+        let debug_config = DebugConfig::new(config);
 
         let ramdump = if debug_config.is_ramdump_needed() {
             Some(prepare_ramdump_file(&temporary_directory)?)
diff --git a/virtualizationmanager/src/debug_config.rs b/virtualizationmanager/src/debug_config.rs
index 902a6cb..451d1c6 100644
--- a/virtualizationmanager/src/debug_config.rs
+++ b/virtualizationmanager/src/debug_config.rs
@@ -15,17 +15,17 @@
 //! Functions for AVF debug policy and debug level
 
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
-    VirtualMachineAppConfig::DebugLevel::DebugLevel,
+    VirtualMachineAppConfig::DebugLevel::DebugLevel, VirtualMachineConfig::VirtualMachineConfig,
 };
 use anyhow::{anyhow, Context, Error, Result};
+use lazy_static::lazy_static;
+use libfdt::{Fdt, FdtError};
+use log::{info, warn};
+use rustutils::system_properties;
+use std::ffi::{CString, NulError};
 use std::fs;
 use std::io::ErrorKind;
 use std::path::{Path, PathBuf};
-use std::ffi::{CString, NulError};
-use log::{warn, info};
-use rustutils::system_properties;
-use libfdt::{Fdt, FdtError};
-use lazy_static::lazy_static;
 
 const CUSTOM_DEBUG_POLICY_OVERLAY_SYSPROP: &str =
     "hypervisor.virtualizationmanager.debug_policy.path";
@@ -156,7 +156,12 @@
 }
 
 impl DebugConfig {
-    pub fn new(debug_level: DebugLevel) -> Self {
+    pub fn new(config: &VirtualMachineConfig) -> Self {
+        let debug_level = match config {
+            VirtualMachineConfig::AppConfig(config) => config.debugLevel,
+            _ => DebugLevel::NONE,
+        };
+
         match system_properties::read(CUSTOM_DEBUG_POLICY_OVERLAY_SYSPROP).unwrap_or_default() {
             Some(path) if !path.is_empty() => {
                 match Self::from_custom_debug_overlay_policy(debug_level, Path::new(&path)) {
@@ -179,6 +184,11 @@
         }
 
         info!("Debug policy is disabled");
+        Self::new_with_debug_level(debug_level)
+    }
+
+    /// Creates a new DebugConfig with debug level. Only use this for test purpose.
+    pub fn new_with_debug_level(debug_level: DebugLevel) -> Self {
         Self {
             debug_level,
             debug_policy_log: false,
@@ -228,14 +238,6 @@
 #[cfg(test)]
 mod tests {
     use super::*;
-    use anyhow::ensure;
-
-    fn can_set_sysprop() -> bool {
-        if let Ok(Some(value)) = system_properties::read("ro.build.type") {
-            return "user".eq(&value);
-        }
-        false // if we're in doubt, skip test.
-    }
 
     #[test]
     fn test_read_avf_debug_policy_with_ramdump() -> Result<()> {
@@ -316,40 +318,4 @@
 
         Ok(())
     }
-
-    fn test_new_with_custom_policy_internal() -> Result<()> {
-        let debug_config = DebugConfig::new(DebugLevel::NONE);
-
-        ensure!(debug_config.debug_level == DebugLevel::NONE);
-        ensure!(!debug_config.debug_policy_log);
-        ensure!(!debug_config.debug_policy_ramdump);
-        ensure!(debug_config.debug_policy_adb);
-
-        Ok(())
-    }
-
-    #[test]
-    fn test_new_with_custom_policy() -> Result<()> {
-        if !can_set_sysprop() {
-            // Skip test if we can't override sysprop.
-            return Ok(());
-        }
-
-        // Setup
-        let old_sysprop = system_properties::read(CUSTOM_DEBUG_POLICY_OVERLAY_SYSPROP)
-            .context("Failed to read existing sysprop")?
-            .unwrap_or_default();
-        let file_name = "avf_debug_policy_with_adb.dtbo";
-        system_properties::write(CUSTOM_DEBUG_POLICY_OVERLAY_SYSPROP, file_name)
-            .context("Failed to set sysprop")?;
-
-        // Run test
-        let test_result = test_new_with_custom_policy_internal();
-
-        // Clean up.
-        system_properties::write(CUSTOM_DEBUG_POLICY_OVERLAY_SYSPROP, &old_sysprop)
-            .context("Failed to restore sysprop")?;
-
-        test_result
-    }
 }
diff --git a/virtualizationmanager/src/payload.rs b/virtualizationmanager/src/payload.rs
index 05626d3..9d0c7d6 100644
--- a/virtualizationmanager/src/payload.rs
+++ b/virtualizationmanager/src/payload.rs
@@ -631,7 +631,7 @@
             collect_apex_infos(
                 &apex_info_list,
                 &apex_configs,
-                &DebugConfig::new(DebugLevel::FULL)
+                &DebugConfig::new_with_debug_level(DebugLevel::FULL)
             )?,
             vec![
                 // Pass active/required APEXes
@@ -660,8 +660,11 @@
         };
         let apex_configs = vec![ApexConfig { name: "apex-vendor".to_string() }];
 
-        let ret =
-            collect_apex_infos(&apex_info_list, &apex_configs, &DebugConfig::new(DebugLevel::NONE));
+        let ret = collect_apex_infos(
+            &apex_info_list,
+            &apex_configs,
+            &DebugConfig::new_with_debug_level(DebugLevel::NONE),
+        );
         assert!(ret
             .is_err_and(|ret| ret.to_string()
                 == "Non-system APEX apex-vendor is not supported in Microdroid"));
@@ -687,7 +690,7 @@
             collect_apex_infos(
                 &apex_info_list,
                 &apex_configs,
-                &DebugConfig::new(DebugLevel::NONE)
+                &DebugConfig::new_with_debug_level(DebugLevel::NONE)
             )?,
             vec![&apex_info_list.list[0]]
         );