Reland "vm tool: rename --gki arg to more generic --os one"

Test: presubmit
Bug: 333730505
Change-Id: I867bb2edeb9f5a285fdee59dae2366e3afeb707d
diff --git a/android/vm/src/main.rs b/android/vm/src/main.rs
index 81ca8fa..110e0ca 100644
--- a/android/vm/src/main.rs
+++ b/android/vm/src/main.rs
@@ -155,10 +155,10 @@
     #[arg(long)]
     devices: Vec<PathBuf>,
 
-    /// Version of GKI to use. If set, use instead of microdroid kernel
-    #[cfg(vendor_modules)]
+    /// Version of OS to use. If not set, defaults to microdroid.
+    /// You can list all available OSes via `vm info` command.
     #[arg(long)]
-    gki: Option<String>,
+    os: Option<String>,
 }
 
 impl MicrodroidConfig {
@@ -172,16 +172,6 @@
         }
     }
 
-    fn gki(&self) -> Option<&str> {
-        cfg_if::cfg_if! {
-            if #[cfg(vendor_modules)] {
-                self.gki.as_deref()
-            } else {
-                None
-            }
-        }
-    }
-
     fn devices(&self) -> &[PathBuf] {
         cfg_if::cfg_if! {
             if #[cfg(device_assignment)] {
diff --git a/android/vm/src/run.rs b/android/vm/src/run.rs
index 0e1f4cc..b07a472 100644
--- a/android/vm/src/run.rs
+++ b/android/vm/src/run.rs
@@ -140,11 +140,7 @@
         bail!("Either --config-path or --payload-binary-name must be defined")
     };
 
-    let os_name = if let Some(ver) = config.microdroid.gki() {
-        format!("microdroid_gki-{ver}")
-    } else {
-        "microdroid".to_owned()
-    };
+    let os_name = if let Some(ref os) = config.microdroid.os { os } else { "microdroid" };
 
     let payload_config_str = format!("{:?}!{:?}", config.apk, payload);
 
@@ -192,7 +188,7 @@
         memoryMib: config.common.mem.unwrap_or(0) as i32, // 0 means use the VM default
         cpuTopology: config.common.cpu_topology,
         customConfig: Some(custom_config),
-        osName: os_name,
+        osName: os_name.to_string(),
         hugePages: config.common.hugepages,
         boostUclamp: config.common.boost_uclamp,
     });
diff --git a/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java b/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
index c5e0171..811d8db 100644
--- a/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
+++ b/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
@@ -40,11 +40,8 @@
 import java.io.File;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
-import java.util.HashSet;
 import java.util.List;
-import java.util.Set;
-import java.util.stream.Collectors;
+import java.util.Map;
 
 public abstract class MicrodroidHostTestCaseBase extends BaseHostJUnit4Test {
     protected static final String TEST_ROOT = "/data/local/tmp/virt/";
@@ -70,8 +67,13 @@
                             * 1000
                             / MICRODROID_COMMAND_RETRY_INTERVAL_MILLIS);
 
-    protected static final Set<String> SUPPORTED_GKI_VERSIONS =
-            Collections.unmodifiableSet(new HashSet(Arrays.asList("android15-6.6")));
+    // We use a map here because the parameterizer `DeviceParameterizedRunner` doesn't support "-"
+    // in test names.
+    // The key of the map is the name of the parameter while the value is the actual OS variant.
+    protected static final Map<String, String> SUPPORTED_OSES =
+            Map.ofEntries(
+                    Map.entry("microdroid", "microdroid"),
+                    Map.entry("android15_66", "microdroid_gki-android15-6.6"));
 
     /* Keep this sync with AssignableDevice.aidl */
     public static final class AssignableDevice {
@@ -261,13 +263,6 @@
         return parseStringArrayFieldsFromVmInfo("Available OS list: ");
     }
 
-    public List<String> getSupportedGKIVersions() throws Exception {
-        return getSupportedOSList().stream()
-                .filter(os -> os.startsWith("microdroid_gki-"))
-                .map(os -> os.replaceFirst("^microdroid_gki-", ""))
-                .collect(Collectors.toList());
-    }
-
     protected boolean isPkvmHypervisor() throws DeviceNotAvailableException {
         return "kvm.arm-protected".equals(getDevice().getProperty("ro.boot.hypervisor.version"));
     }
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index ee2da09..0f77f7a 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -78,7 +78,6 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -112,36 +111,18 @@
         }
     }
 
-    // This map is needed because the parameterizer `DeviceParameterizedRunner` doesn't support "-"
-    // in test names. The key is the test name, while the value is the actual kernel version.
-    private static HashMap<String, String> sGkiVersions = new HashMap<>();
-
-    private static void initGkiVersions() {
-        if (!sGkiVersions.isEmpty()) {
-            return;
-        }
-        sGkiVersions.put("null", null); /* use microdroid kernel */
-        // TODO(b/302465542): run only the latest GKI on presubmit to reduce running time
-        for (String gki : SUPPORTED_GKI_VERSIONS) {
-            String key = gki.split("-")[0];
-            assertThat(sGkiVersions.containsKey(key)).isFalse();
-            sGkiVersions.put(key, gki);
-        }
-    }
-
     public static List<Object[]> params() {
         List<Object[]> ret = new ArrayList<>();
-        for (Object[] gki : gkiVersions()) {
-            ret.add(new Object[] {true /* protectedVm */, gki[0]});
-            ret.add(new Object[] {false /* protectedVm */, gki[0]});
+        for (Object[] osKey : osVersions()) {
+            ret.add(new Object[] {true /* protectedVm */, osKey[0]});
+            ret.add(new Object[] {false /* protectedVm */, osKey[0]});
         }
         return ret;
     }
 
-    public static List<Object[]> gkiVersions() {
-        initGkiVersions();
-        return sGkiVersions.keySet().stream()
-                .map(gki -> new Object[] {gki})
+    public static List<Object[]> osVersions() {
+        return SUPPORTED_OSES.keySet().stream()
+                .map(osKey -> new Object[] {osKey})
                 .collect(Collectors.toList());
     }
 
@@ -304,10 +285,10 @@
             Map<String, File> keyOverrides,
             boolean isProtected,
             boolean updateBootconfigs,
-            String gki)
+            String os)
             throws Exception {
         CommandRunner android = new CommandRunner(getDevice());
-        gki = sGkiVersions.get(gki);
+        os = SUPPORTED_OSES.get(os);
 
         File virtApexDir = FileUtil.createTempDir("virt_apex");
 
@@ -359,7 +340,6 @@
         //   - its idsig
 
         // Load etc/microdroid.json
-        final String os = (gki == null) ? "microdroid" : "microdroid_gki-" + gki;
         File microdroidConfigFile = new File(virtApexEtcDir, os + ".json");
         JSONObject config = new JSONObject(FileUtil.readStringFromFile(microdroidConfigFile));
 
@@ -424,9 +404,6 @@
                         "--console " + CONSOLE_PATH,
                         "--log " + LOG_PATH,
                         configPath);
-        if (gki != null) {
-            args.add("--gki " + gki);
-        }
 
         PipedInputStream pis = new PipedInputStream();
         Process process = createRunUtil().runCmdInBackground(args, new PipedOutputStream(pis));
@@ -500,12 +477,12 @@
     }
 
     @Test
-    @Parameters(method = "gkiVersions")
-    @TestCaseName("{method}_gki_{0}")
+    @Parameters(method = "osVersions")
+    @TestCaseName("{method}_os_{0}")
     @CddTest(requirements = {"9.17/C-2-1", "9.17/C-2-2", "9.17/C-2-6"})
-    public void protectedVmRunsPvmfw(String gki) throws Exception {
+    public void protectedVmRunsPvmfw(String os) throws Exception {
         // Arrange
-        assumeKernelSupported(gki);
+        assumeKernelSupported(os);
         assumeVmTypeSupported(true);
         final String configPath = "assets/vm_config_apex.json";
 
@@ -516,7 +493,7 @@
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
                         .protectedVm(true)
-                        .gki(sGkiVersions.get(gki))
+                        .os(SUPPORTED_OSES.get(os))
                         .name("protected_vm_runs_pvmfw")
                         .build(getAndroidDevice());
 
@@ -533,12 +510,12 @@
     }
 
     @Test
-    @Parameters(method = "gkiVersions")
-    @TestCaseName("{method}_gki_{0}")
+    @Parameters(method = "osVersions")
+    @TestCaseName("{method}_os_{0}")
     @CddTest(requirements = {"9.17/C-2-1", "9.17/C-2-2", "9.17/C-2-5", "9.17/C-2-6"})
-    public void protectedVmWithImageSignedWithDifferentKeyFailsToVerifyPayload(String gki)
+    public void protectedVmWithImageSignedWithDifferentKeyFailsToVerifyPayload(String os)
             throws Exception {
-        assumeKernelSupported(gki);
+        assumeKernelSupported(os);
         assumeVmTypeSupported(true);
         File key = findTestFile("test.com.android.virt.pem");
 
@@ -549,7 +526,7 @@
                         /* keyOverrides= */ Map.of(),
                         /* isProtected= */ true,
                         /* updateBootconfigs= */ true,
-                        gki);
+                        os);
 
         // Assert
         vmInfo.mProcess.waitFor(5L, TimeUnit.SECONDS);
@@ -562,13 +539,13 @@
     }
 
     @Test
-    @Parameters(method = "gkiVersions")
-    @TestCaseName("{method}_gki_{0}")
+    @Parameters(method = "osVersions")
+    @TestCaseName("{method}_os_{0}")
     @CddTest(requirements = {"9.17/C-2-2", "9.17/C-2-6"})
-    public void testBootSucceedsWhenNonProtectedVmStartsWithImagesSignedWithDifferentKey(String gki)
+    public void testBootSucceedsWhenNonProtectedVmStartsWithImagesSignedWithDifferentKey(String os)
             throws Exception {
         // Preconditions
-        assumeKernelSupported(gki);
+        assumeKernelSupported(os);
 
         File key = findTestFile("test.com.android.virt.pem");
         Map<String, File> keyOverrides = Map.of();
@@ -578,7 +555,7 @@
                         keyOverrides,
                         /* isProtected= */ false,
                         /* updateBootconfigs= */ true,
-                        gki);
+                        os);
         assertThatEventually(
                 100000,
                 () ->
@@ -590,12 +567,12 @@
     }
 
     @Test
-    @Parameters(method = "gkiVersions")
-    @TestCaseName("{method}_gki_{0}")
+    @Parameters(method = "osVersions")
+    @TestCaseName("{method}_os_{0}")
     @CddTest(requirements = {"9.17/C-2-2", "9.17/C-2-5", "9.17/C-2-6"})
-    public void testBootFailsWhenVbMetaDigestDoesNotMatchBootconfig(String gki) throws Exception {
+    public void testBootFailsWhenVbMetaDigestDoesNotMatchBootconfig(String os) throws Exception {
         // protectedVmWithImageSignedWithDifferentKeyRunsPvmfw() is the protected case.
-        assumeKernelSupported(gki);
+        assumeKernelSupported(os);
 
         // Sign everything with key1 except vbmeta
         File key = findTestFile("test.com.android.virt.pem");
@@ -606,7 +583,7 @@
                         Map.of(),
                         /* isProtected= */ false,
                         /* updateBootconfigs= */ false,
-                        gki);
+                        os);
         // Wait so that init can print errors to console (time in cuttlefish >> in real device)
         assertThatEventually(
                 100000,
@@ -654,7 +631,7 @@
     }
 
     private boolean isTombstoneGeneratedWithCmd(
-            boolean protectedVm, String gki, String configPath, String... crashCommand)
+            boolean protectedVm, String os, String configPath, String... crashCommand)
             throws Exception {
         CommandRunner android = new CommandRunner(getDevice());
         String testStartTime = android.runWithTimeout(1000, "date", "'+%Y-%m-%d %H:%M:%S.%N'");
@@ -665,7 +642,7 @@
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
                         .protectedVm(protectedVm)
-                        .gki(sGkiVersions.get(gki))
+                        .os(SUPPORTED_OSES.get(os))
                         .build(getAndroidDevice());
         mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
         mMicrodroidDevice.enableAdbRoot();
@@ -682,18 +659,18 @@
 
     @Test
     @Parameters(method = "params")
-    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
-    public void testTombstonesAreGeneratedUponUserspaceCrash(boolean protectedVm, String gki)
+    @TestCaseName("{method}_protectedVm_{0}_os_{1}")
+    public void testTombstonesAreGeneratedUponUserspaceCrash(boolean protectedVm, String os)
             throws Exception {
         // Preconditions
-        assumeKernelSupported(gki);
+        assumeKernelSupported(os);
         assumeVmTypeSupported(protectedVm);
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
         assertThat(
                         isTombstoneGeneratedWithCmd(
                                 protectedVm,
-                                gki,
+                                os,
                                 "assets/vm_config.json",
                                 "kill",
                                 "-SIGSEGV",
@@ -703,18 +680,18 @@
 
     @Test
     @Parameters(method = "params")
-    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    @TestCaseName("{method}_protectedVm_{0}_os_{1}")
     public void testTombstonesAreNotGeneratedIfNotExportedUponUserspaceCrash(
-            boolean protectedVm, String gki) throws Exception {
+            boolean protectedVm, String os) throws Exception {
         // Preconditions
-        assumeKernelSupported(gki);
+        assumeKernelSupported(os);
         assumeVmTypeSupported(protectedVm);
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
         assertThat(
                         isTombstoneGeneratedWithCmd(
                                 protectedVm,
-                                gki,
+                                os,
                                 "assets/vm_config_no_tombstone.json",
                                 "kill",
                                 "-SIGSEGV",
@@ -724,12 +701,12 @@
 
     @Test
     @Parameters(method = "params")
-    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    @TestCaseName("{method}_protectedVm_{0}_os_{1}")
     @Ignore("b/341087884") // TODO(b/341087884): fix & re-enable
-    public void testTombstonesAreGeneratedUponKernelCrash(boolean protectedVm, String gki)
+    public void testTombstonesAreGeneratedUponKernelCrash(boolean protectedVm, String os)
             throws Exception {
         // Preconditions
-        assumeKernelSupported(gki);
+        assumeKernelSupported(os);
         assumeVmTypeSupported(protectedVm);
         assumeFalse("Cuttlefish is not supported", isCuttlefish());
         assumeFalse("Skipping test because ramdump is disabled on user build", isUserBuild());
@@ -738,7 +715,7 @@
         assertThat(
                         isTombstoneGeneratedWithCmd(
                                 protectedVm,
-                                gki,
+                                os,
                                 "assets/vm_config.json",
                                 "echo",
                                 "c",
@@ -748,12 +725,12 @@
     }
 
     private boolean isTombstoneGeneratedWithVmRunApp(
-            boolean protectedVm, String gki, boolean debuggable, String... additionalArgs)
+            boolean protectedVm, String os, boolean debuggable, String... additionalArgs)
             throws Exception {
         // we can't use microdroid builder as it wants ADB connection (debuggable)
         CommandRunner android = new CommandRunner(getDevice());
         String testStartTime = android.runWithTimeout(1000, "date", "'+%Y-%m-%d %H:%M:%S.%N'");
-        gki = sGkiVersions.get(gki);
+        os = SUPPORTED_OSES.get(os);
 
         android.run("rm", "-rf", TEST_ROOT + "*");
         android.run("mkdir", "-p", TEST_ROOT + "*");
@@ -780,10 +757,8 @@
         if (protectedVm) {
             cmd.add("--protected");
         }
-        if (gki != null) {
-            cmd.add("--gki");
-            cmd.add(gki);
-        }
+        cmd.add("--os");
+        cmd.add(os);
         Collections.addAll(cmd, additionalArgs);
 
         android.run(cmd.toArray(new String[0]));
@@ -791,10 +766,10 @@
     }
 
     private boolean isTombstoneGeneratedWithCrashPayload(
-            boolean protectedVm, String gki, boolean debuggable) throws Exception {
+            boolean protectedVm, String os, boolean debuggable) throws Exception {
         return isTombstoneGeneratedWithVmRunApp(
                 protectedVm,
-                gki,
+                os,
                 debuggable,
                 "--payload-binary-name",
                 "MicrodroidCrashNativeLib.so");
@@ -802,76 +777,76 @@
 
     @Test
     @Parameters(method = "params")
-    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
-    public void testTombstonesAreGeneratedWithCrashPayload(boolean protectedVm, String gki)
+    @TestCaseName("{method}_protectedVm_{0}_os_{1}")
+    public void testTombstonesAreGeneratedWithCrashPayload(boolean protectedVm, String os)
             throws Exception {
         // Preconditions
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assumeKernelSupported(gki);
+        assumeKernelSupported(os);
         assumeVmTypeSupported(protectedVm);
 
         // Act
-        assertThat(isTombstoneGeneratedWithCrashPayload(protectedVm, gki, /* debuggable= */ true))
+        assertThat(isTombstoneGeneratedWithCrashPayload(protectedVm, os, /* debuggable= */ true))
                 .isTrue();
     }
 
     @Test
     @Parameters(method = "params")
-    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    @TestCaseName("{method}_protectedVm_{0}_os_{1}")
     public void testTombstonesAreNotGeneratedWithCrashPayloadWhenNonDebuggable(
-            boolean protectedVm, String gki) throws Exception {
+            boolean protectedVm, String os) throws Exception {
         // Preconditions
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assumeKernelSupported(gki);
+        assumeKernelSupported(os);
         assumeVmTypeSupported(protectedVm);
 
         // Act
-        assertThat(isTombstoneGeneratedWithCrashPayload(protectedVm, gki, /* debuggable= */ false))
+        assertThat(isTombstoneGeneratedWithCrashPayload(protectedVm, os, /* debuggable= */ false))
                 .isFalse();
     }
 
     private boolean isTombstoneGeneratedWithCrashConfig(
-            boolean protectedVm, String gki, boolean debuggable) throws Exception {
+            boolean protectedVm, String os, boolean debuggable) throws Exception {
         return isTombstoneGeneratedWithVmRunApp(
-                protectedVm, gki, debuggable, "--config-path", "assets/vm_config_crash.json");
+                protectedVm, os, debuggable, "--config-path", "assets/vm_config_crash.json");
     }
 
     @Test
     @Parameters(method = "params")
-    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
-    public void testTombstonesAreGeneratedWithCrashConfig(boolean protectedVm, String gki)
+    @TestCaseName("{method}_protectedVm_{0}_os_{1}")
+    public void testTombstonesAreGeneratedWithCrashConfig(boolean protectedVm, String os)
             throws Exception {
         // Preconditions
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assumeKernelSupported(gki);
+        assumeKernelSupported(os);
         assumeVmTypeSupported(protectedVm);
 
         // Act
-        assertThat(isTombstoneGeneratedWithCrashConfig(protectedVm, gki, /* debuggable= */ true))
+        assertThat(isTombstoneGeneratedWithCrashConfig(protectedVm, os, /* debuggable= */ true))
                 .isTrue();
     }
 
     @Test
     @Parameters(method = "params")
-    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    @TestCaseName("{method}_protectedVm_{0}_os_{1}")
     public void testTombstonesAreNotGeneratedWithCrashConfigWhenNonDebuggable(
-            boolean protectedVm, String gki) throws Exception {
+            boolean protectedVm, String os) throws Exception {
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assumeKernelSupported(gki);
+        assumeKernelSupported(os);
         assumeVmTypeSupported(protectedVm);
-        assertThat(isTombstoneGeneratedWithCrashConfig(protectedVm, gki, /* debuggable= */ false))
+        assertThat(isTombstoneGeneratedWithCrashConfig(protectedVm, os, /* debuggable= */ false))
                 .isFalse();
     }
 
     @Test
     @Parameters(method = "params")
-    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
-    public void testTelemetryPushedAtoms(boolean protectedVm, String gki) throws Exception {
-        assumeKernelSupported(gki);
+    @TestCaseName("{method}_protectedVm_{0}_os_{1}")
+    public void testTelemetryPushedAtoms(boolean protectedVm, String os) throws Exception {
+        assumeKernelSupported(os);
         assumeVmTypeSupported(protectedVm);
         // Reset statsd config and report before the test
         ConfigUtils.removeConfig(getDevice());
@@ -894,7 +869,7 @@
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
                         .protectedVm(protectedVm)
-                        .gki(sGkiVersions.get(gki))
+                        .os(SUPPORTED_OSES.get(os))
                         .name("test_telemetry_pushed_atoms")
                         .build(device);
         microdroid.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
@@ -1031,11 +1006,11 @@
 
     @Test
     @Parameters(method = "params")
-    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    @TestCaseName("{method}_protectedVm_{0}_os_{1}")
     @CddTest(requirements = {"9.17/C-1-1", "9.17/C-1-2", "9.17/C/1-3"})
-    public void testMicrodroidBoots(boolean protectedVm, String gki) throws Exception {
+    public void testMicrodroidBoots(boolean protectedVm, String os) throws Exception {
         // Preconditions
-        assumeKernelSupported(gki);
+        assumeKernelSupported(os);
         assumeVmTypeSupported(protectedVm);
 
         final String configPath = "assets/vm_config.json"; // path inside the APK
@@ -1046,34 +1021,34 @@
                         .cpuTopology("match_host")
                         .protectedVm(protectedVm)
                         .name("test_microdroid_boots")
-                        .gki(sGkiVersions.get(gki)));
+                        .os(SUPPORTED_OSES.get(os)));
     }
 
     @Test
-    public void testMicrodroidRamUsage_protectedVm_true_gki_null() throws Exception {
-        checkMicrodroidRamUsage(/* protectedVm= */ true, /* gki= */ "null");
+    public void testMicrodroidRamUsage_protectedVm_true_os_microdroid() throws Exception {
+        checkMicrodroidRamUsage(/* protectedVm= */ true, /* os= */ "microdroid");
     }
 
     @Test
-    public void testMicrodroidRamUsage_protectedVm_false_gki_null() throws Exception {
-        checkMicrodroidRamUsage(/* protectedVm= */ false, /* gki= */ "null");
+    public void testMicrodroidRamUsage_protectedVm_false_os_microdroid() throws Exception {
+        checkMicrodroidRamUsage(/* protectedVm= */ false, /* os= */ "microdroid");
     }
 
     @Test
-    public void testMicrodroidRamUsage_protectedVm_true_gki_android15() throws Exception {
-        checkMicrodroidRamUsage(/* protectedVm= */ true, /* gki= */ "android15");
+    public void testMicrodroidRamUsage_protectedVm_true_os_android15_66() throws Exception {
+        checkMicrodroidRamUsage(/* protectedVm= */ true, /* os= */ "android15_66");
     }
 
     @Test
-    public void testMicrodroidRamUsage_protectedVm_false_gki_android15() throws Exception {
-        checkMicrodroidRamUsage(/* protectedVm= */ false, /* gki= */ "android15");
+    public void testMicrodroidRamUsage_protectedVm_false_os_android15_66() throws Exception {
+        checkMicrodroidRamUsage(/* protectedVm= */ false, /* os= */ "android15_66");
     }
 
     // TODO(b/209036125): Upgrade this function to a parameterized test once metrics can be
     // collected with tradefed parameterizer.
-    void checkMicrodroidRamUsage(boolean protectedVm, String gki) throws Exception {
+    void checkMicrodroidRamUsage(boolean protectedVm, String os) throws Exception {
         // Preconditions
-        assumeKernelSupported(gki);
+        assumeKernelSupported(os);
         assumeVmTypeSupported(protectedVm);
 
         final String configPath = "assets/vm_config.json";
@@ -1083,7 +1058,7 @@
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
                         .protectedVm(protectedVm)
-                        .gki(sGkiVersions.get(gki))
+                        .os(SUPPORTED_OSES.get(os))
                         .name("test_microdroid_ram_usage")
                         .build(getAndroidDevice());
         mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
@@ -1266,10 +1241,10 @@
 
     @Test
     @Parameters(method = "params")
-    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
-    public void testDeviceAssignment(boolean protectedVm, String gki) throws Exception {
+    @TestCaseName("{method}_protectedVm_{0}_os_{1}")
+    public void testDeviceAssignment(boolean protectedVm, String os) throws Exception {
         // Preconditions
-        assumeKernelSupported(gki);
+        assumeKernelSupported(os);
         assumeVmTypeSupported(protectedVm);
         assumeVfioPlatformSupported();
 
@@ -1280,7 +1255,7 @@
 
         // Try assign devices one by one
         for (AssignableDevice device : devices) {
-            launchWithDeviceAssignment(device.node, protectedVm, gki);
+            launchWithDeviceAssignment(device.node, protectedVm, os);
 
             String dtPath =
                     new CommandRunner(mMicrodroidDevice)
@@ -1302,7 +1277,7 @@
         }
     }
 
-    private void launchWithDeviceAssignment(String device, boolean protectedVm, String gki)
+    private void launchWithDeviceAssignment(String device, boolean protectedVm, String os)
             throws Exception {
         Objects.requireNonNull(device);
         final String configPath = "assets/vm_config.json";
@@ -1313,7 +1288,7 @@
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
                         .protectedVm(protectedVm)
-                        .gki(sGkiVersions.get(gki))
+                        .os(SUPPORTED_OSES.get(os))
                         .addAssignableDevice(device)
                         .build(getAndroidDevice());
 
@@ -1322,20 +1297,18 @@
     }
 
     @Test
-    public void testGkiVersions() throws Exception {
-        for (String gki : getSupportedGKIVersions()) {
-            assertTrue(
-                    "Unknown gki \"" + gki + "\". Supported gkis: " + SUPPORTED_GKI_VERSIONS,
-                    SUPPORTED_GKI_VERSIONS.contains(gki));
+    public void testOsVersions() throws Exception {
+        for (String os : getSupportedOSList()) {
+            assertWithMessage("Unknown OS \"%s\"", os).that(SUPPORTED_OSES.values()).contains(os);
         }
     }
 
     @Test
     @Parameters(method = "params")
-    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
-    public void testHugePages(boolean protectedVm, String gki) throws Exception {
+    @TestCaseName("{method}_protectedVm_{0}_os_{1}")
+    public void testHugePages(boolean protectedVm, String os) throws Exception {
         // Preconditions
-        assumeKernelSupported(gki);
+        assumeKernelSupported(os);
         assumeVmTypeSupported(protectedVm);
 
         ITestDevice device = getDevice();
@@ -1359,7 +1332,7 @@
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
                         .protectedVm(protectedVm)
-                        .gki(sGkiVersions.get(gki))
+                        .os(SUPPORTED_OSES.get(os))
                         .hugePages(true)
                         .name("test_huge_pages")
                         .build(getAndroidDevice());
@@ -1439,13 +1412,11 @@
         return runUtil;
     }
 
-    private void assumeKernelSupported(String gki) throws Exception {
-        String gkiVersion = sGkiVersions.get(gki);
-        if (gkiVersion != null) {
-            assumeTrue(
-                    "Skipping test as the GKI is not supported: " + gkiVersion,
-                    getSupportedGKIVersions().contains(gkiVersion));
-        }
+    private void assumeKernelSupported(String osKey) throws Exception {
+        String os = SUPPORTED_OSES.get(osKey);
+        assumeTrue(
+                "Skipping test as OS \"" + os + "\" is not supported",
+                getSupportedOSList().contains(os));
     }
 
     private void assumeVmTypeSupported(boolean protectedVm) throws Exception {