Migrate pkvm_perf_test.py test_boot_time_when_pkvm_enabled and
test_boot_time_when_pkvm_disabled functions to AOSP

Bug: 236799228
Test: atest AVFHostTestCases
Change-Id: I4fb8b47e4d0d9e647e73d679d68d1ef388dee949
diff --git a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
index b9d9d25..c6d8f33 100644
--- a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
+++ b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
@@ -22,6 +22,8 @@
 
 import static com.google.common.truth.Truth.assertWithMessage;
 
+import static org.junit.Assume.assumeTrue;
+
 import com.android.microdroid.test.CommandRunner;
 import com.android.microdroid.test.MicrodroidHostTestCaseBase;
 import com.android.tradefed.log.LogUtil.CLog;
@@ -56,7 +58,7 @@
     private static final int BOOT_COMPLETE_TIMEOUT_MS = 10 * 60 * 1000;
     private static final double NANOS_IN_SEC = 1_000_000_000.0;
     private static final int ROUND_COUNT = 5;
-    private static final String METRIC_PREFIX = "avf_perf/compos/";
+    private static final String METRIC_PREFIX = "avf_perf/hostside/";
 
     @Before
     public void setUp() throws Exception {
@@ -65,8 +67,9 @@
 
     @After
     public void tearDown() throws Exception {
-        // Reboot to prevent previous staged session.
-        rebootAndWaitBootCompleted();
+        // Set PKVM enable and reboot to prevent previous staged session.
+        setPKVMStatusWithRebootToBootloader(true);
+        rebootFromBootloaderAndWaitBootCompleted();
 
         CommandRunner android = new CommandRunner(getDevice());
 
@@ -75,6 +78,37 @@
     }
 
     @Test
+    public void testBootEnableAndDisablePKVM() throws Exception {
+
+        testPKVMStatusSwitchSupported();
+
+        double[] bootWithPKVMEnableTime = new double[ROUND_COUNT];
+        double[] bootWithoutPKVMEnableTime = new double[ROUND_COUNT];
+
+        for (int round = 0; round < ROUND_COUNT; ++round) {
+
+            setPKVMStatusWithRebootToBootloader(true);
+            long start = System.nanoTime();
+            rebootFromBootloaderAndWaitBootCompleted();
+            long elapsedWithPKVMEnable = System.nanoTime() - start;
+            double elapsedSec = elapsedWithPKVMEnable / NANOS_IN_SEC;
+            bootWithPKVMEnableTime[round] = elapsedSec;
+            CLog.i("Boot time with PKVM enable took " + elapsedSec + "s");
+
+            setPKVMStatusWithRebootToBootloader(false);
+            start = System.nanoTime();
+            rebootFromBootloaderAndWaitBootCompleted();
+            long elapsedWithoutPKVMEnable = System.nanoTime() - start;
+            elapsedSec = elapsedWithoutPKVMEnable / NANOS_IN_SEC;
+            bootWithoutPKVMEnableTime[round] = elapsedSec;
+            CLog.i("Boot time with PKVM disable took " + elapsedSec + "s");
+        }
+
+        reportMetric("boot_time_with_pkvm_enable", "s", bootWithPKVMEnableTime);
+        reportMetric("boot_time_with_pkvm_disable", "s", bootWithoutPKVMEnableTime);
+    }
+
+    @Test
     public void testBootWithAndWithoutCompOS() throws Exception {
 
         double[] bootWithCompOsTime = new double[ROUND_COUNT];
@@ -106,6 +140,18 @@
         reportMetric("boot_time_without_compos", "s", bootWithoutCompOsTime);
     }
 
+    private void testPKVMStatusSwitchSupported() throws Exception {
+        if (!getDevice().isStateBootloaderOrFastbootd()) {
+            getDevice().rebootIntoBootloader();
+        }
+        getDevice().waitForDeviceBootloader();
+
+        CommandResult result;
+        result = getDevice().executeFastbootCommand("oem", "pkvm", "status");
+        rebootFromBootloaderAndWaitBootCompleted();
+        assumeTrue(!result.getStderr().contains("Invalid oem command"));
+    }
+
     private void reportMetric(String name, String unit, double[] values) {
         double sum = 0;
         double min = Double.MAX_VALUE;
@@ -133,6 +179,41 @@
         metrics.addTestMetric(METRIC_PREFIX + name + "_stdev_" + unit, Double.toString(stdev));
     }
 
+    private void setPKVMStatusWithRebootToBootloader(boolean isEnable) throws Exception {
+
+        if (!getDevice().isStateBootloaderOrFastbootd()) {
+            getDevice().rebootIntoBootloader();
+        }
+        getDevice().waitForDeviceBootloader();
+
+        CommandResult result;
+        if (isEnable) {
+            result = getDevice().executeFastbootCommand("oem", "pkvm", "enable");
+        } else {
+            result = getDevice().executeFastbootCommand("oem", "pkvm", "disable");
+        }
+
+        result = getDevice().executeFastbootCommand("oem", "pkvm", "status");
+        CLog.i("Gets PKVM status : " + result);
+
+        String expectedOutput = "";
+
+        if (isEnable) {
+            expectedOutput = "pkvm is enabled";
+        } else {
+            expectedOutput = "pkvm is disabled";
+        }
+        assertWithMessage("Failed to set PKVM status. Reason: " + result)
+            .that(result.toString()).ignoringCase().contains(expectedOutput);
+    }
+
+    private void rebootFromBootloaderAndWaitBootCompleted() throws Exception {
+        getDevice().executeFastbootCommand("reboot");
+        getDevice().waitForDeviceOnline(BOOT_COMPLETE_TIMEOUT_MS);
+        getDevice().waitForBootComplete(BOOT_COMPLETE_TIMEOUT_MS);
+        getDevice().enableAdbRoot();
+    }
+
     private void rebootAndWaitBootCompleted() throws Exception {
         getDevice().nonBlockingReboot();
         getDevice().waitForDeviceOnline(BOOT_COMPLETE_TIMEOUT_MS);