feat(virtio-pmem): add integration tests

Add functional and API tests for virtio-pmem device
and its configuration fields

Signed-off-by: Egor Lazarchuk <[email protected]>

Author: ShadowCurse

tests/framework/http_api.py

@@ -132,4 +132,5 @@ def __init__(self, api_usocket_full_name, *, on_error=None):
         self.snapshot_load = Resource(self, "/snapshot/load")
         self.cpu_config = Resource(self, "/cpu-config")
         self.entropy = Resource(self, "/entropy")
+        self.pmem = Resource(self, "/pmem", "id")
         self.serial = Resource(self, "/serial")

tests/framework/microvm.py

@@ -973,6 +973,24 @@ def add_net_iface(self, iface=None, api=True, **kwargs):
 
         return iface
 
+    def add_pmem(
+        self,
+        pmem_id,
+        path_on_host,
+        root_device=False,
+        read_only=False,
+    ):
+        """Add a pmem device."""
+
+        path_on_jail = self.create_jailed_resource(path_on_host)
+        self.api.pmem.put(
+            id=pmem_id,
+            path_on_host=path_on_jail,
+            root_device=root_device,
+            read_only=read_only,
+        )
+        self.disks[pmem_id] = path_on_host
+
     def start(self):
         """Start the microvm.
 

tests/framework/vm_config.json

@@ -31,5 +31,6 @@
   "logger": null,
   "metrics": null,
   "mmds-config": null,
-  "entropy": null
+  "entropy": null,
+  "pmem": []
 }

tests/integration_tests/functional/test_api.py

@@ -1044,6 +1044,67 @@ def test_api_balloon(uvm_nano):
         test_microvm.api.balloon.patch(amount_mib=33554432)
 
 
+def test_pmem_api(uvm_plain_any, rootfs):
+    """
+    Test virtio-pmem API commands
+    """
+
+    vm = uvm_plain_any
+    vm.spawn()
+    vm.basic_config(add_root_device=False)
+
+    invalid_pmem_path_on_host = os.path.join(vm.fsfiles, "invalid_scratch")
+    utils.check_output(f"touch {invalid_pmem_path_on_host}")
+    invalid_pmem_file_path = vm.create_jailed_resource(str(invalid_pmem_path_on_host))
+
+    pmem_size_mb = 2
+    pmem_path_on_host = drive_tools.FilesystemFile(
+        os.path.join(vm.fsfiles, "scratch"), size=pmem_size_mb
+    )
+    pmem_file_path = vm.create_jailed_resource(pmem_path_on_host.path)
+
+    # Try to add pmem without setting `path_on_host`
+    expected_msg = re.escape(
+        "An error occurred when deserializing the json body of a request: missing field `path_on_host`"
+    )
+    with pytest.raises(RuntimeError, match=expected_msg):
+        vm.api.pmem.put(id="pmem")
+
+    # Try to add pmem with 0 sized backing file
+    expected_msg = re.escape("Error backing file size is 0")
+    with pytest.raises(RuntimeError, match=expected_msg):
+        vm.api.pmem.put(id="pmem", path_on_host=invalid_pmem_file_path)
+
+    # Try to add pmem as root while block is set as root
+    vm.api.drive.put(drive_id="drive", path_on_host=pmem_file_path, is_root_device=True)
+    expected_msg = re.escape(
+        "Attempt to add pmem as a root device while the root device defined as a block device"
+    )
+    with pytest.raises(RuntimeError, match=expected_msg):
+        vm.api.pmem.put(id="pmem", path_on_host=pmem_file_path, root_device=True)
+
+    # Reset block from being root
+    vm.api.drive.put(
+        drive_id="drive", path_on_host=pmem_file_path, is_root_device=False
+    )
+
+    # Try to add pmem as root twice
+    vm.api.pmem.put(id="pmem", path_on_host=pmem_file_path, root_device=True)
+    expected_msg = re.escape("A root pmem device already exist")
+    with pytest.raises(RuntimeError, match=expected_msg):
+        vm.api.pmem.put(id="pmem2", path_on_host=pmem_file_path, root_device=True)
+
+    # Reset pmem from being root
+    vm.api.pmem.put(id="pmem", path_on_host=pmem_file_path, root_device=False)
+
+    # Add a rootfs to boot a vm
+    vm.add_pmem("rootfs", rootfs, True, True)
+
+    # No post boot API calls to pmem
+    with pytest.raises(RuntimeError):
+        vm.api.pmem.put(id="pmem")
+
+
 def test_get_full_config_after_restoring_snapshot(microvm_factory, uvm_nano):
     """
     Test the configuration of a microVM after restoring from a snapshot.
@@ -1085,6 +1146,21 @@ def test_get_full_config_after_restoring_snapshot(microvm_factory, uvm_nano):
         }
     ]
 
+    uvm_nano.api.pmem.put(
+        id="pmem",
+        path_on_host="/" + uvm_nano.rootfs_file.name,
+        root_device=False,
+        read_only=False,
+    )
+    setup_cfg["pmem"] = [
+        {
+            "id": "pmem",
+            "path_on_host": "/" + uvm_nano.rootfs_file.name,
+            "root_device": False,
+            "read_only": False,
+        }
+    ]
+
     # Add a memory balloon device.
     uvm_nano.api.balloon.put(amount_mib=1, deflate_on_oom=True)
     setup_cfg["balloon"] = {
@@ -1196,6 +1272,21 @@ def test_get_full_config(uvm_plain):
         }
     ]
 
+    test_microvm.api.pmem.put(
+        id="pmem",
+        path_on_host="/" + test_microvm.rootfs_file.name,
+        root_device=False,
+        read_only=False,
+    )
+    expected_cfg["pmem"] = [
+        {
+            "id": "pmem",
+            "path_on_host": "/" + test_microvm.rootfs_file.name,
+            "root_device": False,
+            "read_only": False,
+        }
+    ]
+
     # Add a memory balloon device.
     test_microvm.api.balloon.put(amount_mib=1, deflate_on_oom=True)
     expected_cfg["balloon"] = {

tests/integration_tests/functional/test_pmem.py

@@ -0,0 +1,128 @@
+# Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the virtio-pmem device."""
+
+import json
+import os
+
+import host_tools.drive as drive_tools
+
+ALIGNMENT = 2 << 20
+
+
+def align(size: int) -> int:
+    """
+    Align the value to ALIGNMENT
+    """
+    return (size + ALIGNMENT - 1) & ~(ALIGNMENT - 1)
+
+
+def check_pmem_exist(vm, index, root, read_only, size, extension):
+    """
+    Check the pmem exist with correct parameters
+    """
+    code, _, _ = vm.ssh.run(f"ls /dev/pmem{index}")
+    assert code == 0
+
+    if root:
+        code, stdout, _ = vm.ssh.run("mount")
+        assert code == 0
+        if read_only:
+            assert f"/dev/pmem0 on / type {extension} (ro" in stdout
+        else:
+            assert f"/dev/pmem0 on / type {extension} (rw" in stdout
+
+    code, stdout, _ = vm.ssh.run("lsblk -J")
+    assert code == 0
+
+    j = json.loads(stdout)
+    blocks = j["blockdevices"]
+    for block in blocks:
+        if block["name"] == f"pmem{index}":
+            assert block["size"][-1] == "M"
+            block_size_mb = int(block["size"][:-1])
+            assert int(block_size_mb << 20) == size
+            if root:
+                assert "/" in block["mountpoints"]
+            return
+    assert False
+
+
+def test_pmem_add(uvm_plain_any, microvm_factory):
+    """
+    Test addition of a single non root pmem device
+    """
+
+    vm = uvm_plain_any
+    vm.spawn()
+    vm.basic_config(add_root_device=True)
+    vm.add_net_iface()
+
+    # Pmem should work with non 2MB aligned files as well
+    pmem_size_mb_1 = 1
+    fs_1 = drive_tools.FilesystemFile(
+        os.path.join(vm.fsfiles, "scratch_1"), size=pmem_size_mb_1
+    )
+    pmem_size_mb_2 = 2
+    fs_2 = drive_tools.FilesystemFile(
+        os.path.join(vm.fsfiles, "scratch_2"), size=pmem_size_mb_2
+    )
+    vm.add_pmem("pmem_1", fs_1.path, False, False)
+    vm.add_pmem("pmem_2", fs_2.path, False, True)
+    vm.start()
+
+    # Both 1MB and 2MB block will show as 2MB because of
+    # the aligment
+    check_pmem_exist(vm, 0, False, False, align(pmem_size_mb_1 << 20), "ext4")
+    check_pmem_exist(vm, 1, False, True, align(pmem_size_mb_2 << 20), "ext4")
+
+    snapshot = vm.snapshot_full()
+    restored_vm = microvm_factory.build_from_snapshot(snapshot)
+    check_pmem_exist(restored_vm, 0, False, False, align(pmem_size_mb_1 << 20), "ext4")
+    check_pmem_exist(restored_vm, 1, False, True, align(pmem_size_mb_2 << 20), "ext4")
+
+
+def test_pmem_add_as_root_rw(uvm_plain_any, rootfs_rw, microvm_factory):
+    """
+    Test addition of a single root pmem device in read-write mode
+    """
+
+    vm = uvm_plain_any
+    vm.memory_monitor = None
+    vm.monitors = []
+    vm.spawn()
+    vm.basic_config(add_root_device=False)
+    vm.add_net_iface()
+
+    rootfs_size = os.path.getsize(rootfs_rw)
+    vm.add_pmem("pmem", rootfs_rw, True, False)
+    vm.start()
+
+    check_pmem_exist(vm, 0, True, False, align(rootfs_size), "ext4")
+
+    snapshot = vm.snapshot_full()
+    restored_vm = microvm_factory.build_from_snapshot(snapshot)
+    check_pmem_exist(restored_vm, 0, True, False, align(rootfs_size), "ext4")
+
+
+def test_pmem_add_as_root_ro(uvm_plain_any, rootfs, microvm_factory):
+    """
+    Test addition of a single root pmem device in read-only mode
+    """
+
+    vm = uvm_plain_any
+    vm.memory_monitor = None
+    vm.monitors = []
+    vm.spawn()
+    vm.basic_config(add_root_device=False)
+    vm.add_net_iface()
+
+    rootfs_size = os.path.getsize(rootfs)
+    vm.add_pmem("pmem", rootfs, True, True)
+    vm.start()
+
+    check_pmem_exist(vm, 0, True, True, align(rootfs_size), "squashfs")
+
+    snapshot = vm.snapshot_full()
+    restored_vm = microvm_factory.build_from_snapshot(snapshot)
+    check_pmem_exist(restored_vm, 0, True, True, align(rootfs_size), "squashfs")