test: revamp .add_drive support

Track added drives in the Microvm class, and clean up some of the logic.

Signed-off-by: Pablo Barbáchano <pablob@amazon.com>

Author: pb8o

tests/framework/microvm.py

@@ -157,14 +157,12 @@ def __init__(
 
         # Initalize memory monitor
         self.memory_monitor = None
-
-        # iface dictionary
-        self.iface = {}
-
-        # Deal with memory monitoring.
         if monitor_memory:
             self.memory_monitor = mem_tools.MemoryMonitor()
 
+        # device dictionaries
+        self.iface = {}
+        self.disks = {}
         self.vcpus_count = None
 
         # External clone/exec tool, because Python can't into clone
@@ -616,20 +614,18 @@ def basic_config(
         ), response.text
 
         if add_root_device and self.rootfs_file != "":
-            jail_fn = self.create_jailed_resource
-            if self.jailer.uses_ramfs:
-                jail_fn = self.copy_to_jail_ramfs
-            # Add the root file system with rw permissions.
-            response = self.drive.put(
+            self.rootfs_file = Path(self.rootfs_file)
+            read_only = self.rootfs_file.suffix == ".squashfs"
+
+            # Add the root file system
+            self.add_drive(
                 drive_id="rootfs",
-                path_on_host=jail_fn(self.rootfs_file),
+                path_on_host=self.rootfs_file,
                 is_root_device=True,
-                is_read_only=False,
+                is_read_only=read_only,
                 io_engine=rootfs_io_engine,
+                use_ramdisk=self.jailer.uses_ramfs,
             )
-            assert self._api_session.is_status_no_content(
-                response.status_code
-            ), response.text
 
     def cpu_config(self, config):
         """Set CPU configuration."""
@@ -669,29 +665,32 @@ def daemonize_jailer(self, jailer_param_list):
     def add_drive(
         self,
         drive_id,
-        file_path,
-        root_device=False,
+        path_on_host,
+        is_root_device=False,
         is_read_only=False,
         partuuid=None,
         cache_type=None,
         io_engine=None,
         use_ramdisk=False,
     ):
         """Add a block device."""
+
+        if use_ramdisk:
+            path_on_jail = self.copy_to_jail_ramfs(path_on_host)
+        else:
+            path_on_jail = self.create_jailed_resource(path_on_host)
+
         response = self.drive.put(
             drive_id=drive_id,
-            path_on_host=(
-                self.copy_to_jail_ramfs(file_path)
-                if use_ramdisk
-                else self.create_jailed_resource(file_path)
-            ),
-            is_root_device=root_device,
+            path_on_host=path_on_jail,
+            is_root_device=is_root_device,
             is_read_only=is_read_only,
             partuuid=partuuid,
             cache_type=cache_type,
             io_engine=io_engine,
         )
         assert self.api_session.is_status_no_content(response.status_code)
+        self.disks[drive_id] = path_on_host
 
     def patch_drive(self, drive_id, file):
         """Modify/patch an existing block device."""
@@ -700,6 +699,7 @@ def patch_drive(self, drive_id, file):
             path_on_host=self.create_jailed_resource(file.path),
         )
         assert self.api_session.is_status_no_content(response.status_code)
+        self.disks[drive_id] = Path(file.path)
 
     def add_net_iface(self, iface=None, api=True, **kwargs):
         """Add a network interface"""

tests/integration_tests/functional/test_api.py

@@ -15,7 +15,7 @@
 import host_tools.drive as drive_tools
 import host_tools.network as net_tools
 from framework import utils_cpuid
-from framework.artifacts import SnapshotType, NetIfaceConfig
+from framework.artifacts import NetIfaceConfig, SnapshotType
 from framework.builder import MicrovmBuilder, SnapshotBuilder
 from framework.utils import get_firecracker_version_from_toml, is_io_uring_supported
 from framework.utils_cpu_templates import skip_on_arm
@@ -793,16 +793,14 @@ def test_drive_patch(test_microvm_with_api):
     # a root file system with the rw permission.
     test_microvm.basic_config(rootfs_io_engine="Sync")
 
-    # The drive to be patched.
     fs = drive_tools.FilesystemFile(os.path.join(test_microvm.fsfiles, "scratch"))
-    response = test_microvm.drive.put(
+    test_microvm.add_drive(
         drive_id="scratch",
-        path_on_host=test_microvm.create_jailed_resource(fs.path),
+        path_on_host=fs.path,
         is_root_device=False,
         is_read_only=False,
         io_engine="Async" if is_io_uring_supported() else "Sync",
     )
-    assert test_microvm.api_session.is_status_no_content(response.status_code)
 
     # Patching drive before boot is not allowed.
     response = test_microvm.drive.patch(drive_id="scratch", path_on_host="foo.bar")

tests/integration_tests/functional/test_drives.py

@@ -191,7 +191,7 @@ def test_partuuid_boot(test_microvm_with_partuuid):
     test_microvm.add_drive(
         "rootfs",
         test_microvm.rootfs_file,
-        root_device=True,
+        is_root_device=True,
         partuuid=PARTUUID[platform.machine()],
     )
 
@@ -219,14 +219,14 @@ def test_partuuid_update(test_microvm_with_api):
 
     # Add the root block device specified through PARTUUID.
     test_microvm.add_drive(
-        "rootfs", test_microvm.rootfs_file, root_device=True, partuuid="0eaa91a0-01"
+        "rootfs", test_microvm.rootfs_file, is_root_device=True, partuuid="0eaa91a0-01"
     )
 
     # Update the root block device to boot from /dev/vda.
     test_microvm.add_drive(
         "rootfs",
         test_microvm.rootfs_file,
-        root_device=True,
+        is_root_device=True,
     )
 
     test_microvm.start()
@@ -289,7 +289,7 @@ def test_no_flush(test_microvm_with_api):
     test_microvm.add_drive(
         "rootfs",
         test_microvm.rootfs_file,
-        root_device=True,
+        is_root_device=True,
     )
 
     # Configure the metrics.
@@ -331,7 +331,7 @@ def test_flush(test_microvm_with_api):
     test_microvm.add_drive(
         "rootfs",
         test_microvm.rootfs_file,
-        root_device=True,
+        is_root_device=True,
         cache_type="Writeback",
     )
 
@@ -369,7 +369,7 @@ def test_block_default_cache_old_version(test_microvm_with_api):
     test_microvm.add_drive(
         "rootfs",
         test_microvm.rootfs_file,
-        root_device=True,
+        is_root_device=True,
         cache_type="Writeback",
     )