Tests: create new fixtures to make writing booted/restored simpler

tests/README.md

@@ -306,10 +306,14 @@ that are pre-initialized with specific guest kernels and rootfs:
   24.04 squashfs as rootfs,
 - `uvm_plain` yields a Firecracker process pre-initialized with a 5.10 kernel
   and the same Ubuntu 24.04 squashfs.
-
-Generally, tests should use the former if you are testing some interaction
-between the guest and Firecracker, while the latter should be used if
-Firecracker functionality unrelated to the guest is being tested.
+- `uvm_any` yields started microvms, parametrized by all supported kernels, all
+  CPU templates (static, custom and none), and either booted or restored from a
+  snapshot.
+- `uvm_any_booted` works the same as `uvm_any`, but only for booted VMs.
+
+Generally, tests should use `uvm_plain_any` if you are testing some interaction
+between the guest and Firecracker, and `uvm_plain` should be used if Firecracker
+functionality unrelated to the guest is being tested.
 
 ### Markers
 

tests/conftest.py

@@ -35,7 +35,7 @@
 
 import host_tools.cargo_build as build_tools
 from framework import defs, utils
-from framework.artifacts import kernel_params, rootfs_params
+from framework.artifacts import disks, kernel_params
 from framework.microvm import MicroVMFactory
 from framework.properties import global_props
 from framework.utils_cpu_templates import (
@@ -292,13 +292,6 @@ def microvm_factory(request, record_property, results_dir):
     uvm_factory.kill()
 
 
-@pytest.fixture(params=static_cpu_templates_params())
-def cpu_template(request, record_property):
-    """Return all static CPU templates supported by the vendor."""
-    record_property("static_cpu_template", request.param)
-    return request.param
-
-
 @pytest.fixture(params=custom_cpu_templates_params())
 def custom_cpu_template(request, record_property):
     """Return all dummy custom CPU templates supported by the vendor."""
@@ -361,13 +354,6 @@ def guest_kernel_fxt(request, record_property):
     return kernel
 
 
-def rootfs_fxt(request, record_property):
-    """Return all supported rootfs."""
-    fs = request.param
-    record_property("rootfs", fs.name)
-    return fs
-
-
 # Fixtures for all guest kernels, and specific versions
 guest_kernel = pytest.fixture(guest_kernel_fxt, params=kernel_params("vmlinux-*"))
 guest_kernel_acpi = pytest.fixture(
@@ -387,9 +373,17 @@ def rootfs_fxt(request, record_property):
     params=kernel_params("vmlinux-6.1*"),
 )
 
-# Fixtures for all Ubuntu rootfs, and specific versions
-rootfs = pytest.fixture(rootfs_fxt, params=rootfs_params("ubuntu-24*.squashfs"))
-rootfs_rw = pytest.fixture(rootfs_fxt, params=rootfs_params("*.ext4"))
+
+@pytest.fixture
+def rootfs():
+    """Return an Ubuntu 24.04 read-only rootfs"""
+    return disks("ubuntu-24.04.squashfs")[0]
+
+
+@pytest.fixture
+def rootfs_rw():
+    """Return an Ubuntu 24.04 ext4 rootfs"""
+    return disks("ubuntu-24.04.ext4")[0]
 
 
 @pytest.fixture
@@ -459,3 +453,80 @@ def uvm_with_initrd(
     uvm = microvm_factory.build(guest_kernel_linux_5_10)
     uvm.initrd_file = fs
     yield uvm
+
+
+@pytest.fixture
+def vcpu_count():
+    """Return default vcpu_count. Use indirect parametrization to override."""
+    return 2
+
+
+@pytest.fixture
+def mem_size_mib():
+    """Return memory size. Use indirect parametrization to override."""
+    return 256
+
+
+def uvm_booted(
+    microvm_factory, guest_kernel, rootfs, cpu_template, vcpu_count=2, mem_size_mib=256
+):
+    """Return a booted uvm"""
+    uvm = microvm_factory.build(guest_kernel, rootfs)
+    uvm.spawn()
+    uvm.basic_config(vcpu_count=vcpu_count, mem_size_mib=mem_size_mib)
+    uvm.set_cpu_template(cpu_template)
+    uvm.add_net_iface()
+    uvm.start()
+    return uvm
+
+
+def uvm_restored(microvm_factory, guest_kernel, rootfs, cpu_template, **kwargs):
+    """Return a restored uvm"""
+    uvm = uvm_booted(microvm_factory, guest_kernel, rootfs, cpu_template, **kwargs)
+    snapshot = uvm.snapshot_full()
+    uvm.kill()
+    uvm2 = microvm_factory.build_from_snapshot(snapshot)
+    uvm2.cpu_template_name = uvm.cpu_template_name
+    return uvm2
+
+
+@pytest.fixture(params=[uvm_booted, uvm_restored])
+def uvm_ctor(request):
+    """Fixture to return uvms with different constructors"""
+    return request.param
+
+
+@pytest.fixture
+def uvm_any(
+    microvm_factory,
+    uvm_ctor,
+    guest_kernel,
+    rootfs,
+    cpu_template_any,
+    vcpu_count,
+    mem_size_mib,
+):
+    """Return booted and restored uvms"""
+    return uvm_ctor(
+        microvm_factory,
+        guest_kernel,
+        rootfs,
+        cpu_template_any,
+        vcpu_count=vcpu_count,
+        mem_size_mib=mem_size_mib,
+    )
+
+
+@pytest.fixture
+def uvm_any_booted(
+    microvm_factory, guest_kernel, rootfs, cpu_template_any, vcpu_count, mem_size_mib
+):
+    """Return booted uvms"""
+    return uvm_booted(
+        microvm_factory,
+        guest_kernel,
+        rootfs,
+        cpu_template_any,
+        vcpu_count=vcpu_count,
+        mem_size_mib=mem_size_mib,
+    )

tests/framework/ab_test.py

@@ -21,7 +21,6 @@
 of both invocations is the same, the test passes (with us being alerted to this situtation via a special pipeline that
 does not block PRs). If not, it fails, preventing PRs from introducing new vulnerable dependencies.
 """
-import os
 import statistics
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -31,14 +30,14 @@
 
 from framework import utils
 from framework.defs import FC_WORKSPACE_DIR
-from framework.microvm import Microvm
+from framework.properties import global_props
 from framework.utils import CommandReturn
 from framework.with_filelock import with_filelock
-from host_tools.cargo_build import DEFAULT_TARGET_DIR, get_firecracker_binaries
+from host_tools.cargo_build import DEFAULT_TARGET_DIR
 
 # Locally, this will always compare against main, even if we try to merge into, say, a feature branch.
 # We might want to do a more sophisticated way to determine a "parent" branch here.
-DEFAULT_A_REVISION = os.environ.get("BUILDKITE_PULL_REQUEST_BASE_BRANCH") or "main"
+DEFAULT_A_REVISION = global_props.buildkite_revision_a or "main"
 
 
 T = TypeVar("T")
@@ -120,11 +119,6 @@ def binary_ab_test(
     return result_a, result_b, comparator(result_a, result_b)
 
 
-def is_pr() -> bool:
-    """Returns `True` iff we are executing in the context of a build kite run on a pull request"""
-    return os.environ.get("BUILDKITE_PULL_REQUEST", "false") != "false"
-
-
 def git_ab_test_host_command_if_pr(
     command: str,
     *,
@@ -134,7 +128,7 @@ def git_ab_test_host_command_if_pr(
     """Runs the given bash command as an A/B-Test if we're in a pull request context (asserting that its stdout and
     stderr did not change across the PR). Otherwise runs the command, asserting it returns a zero exit code
     """
-    if is_pr():
+    if global_props.buildkite_pr:
         git_ab_test_host_command(command, comparator=comparator)
         return None
 
@@ -176,56 +170,6 @@ def set_did_not_grow_comparator(
     )
 
 
-def precompiled_ab_test_guest_command(
-    microvm_factory: Callable[[Path, Path], Microvm],
-    command: str,
-    *,
-    comparator: Callable[[CommandReturn, CommandReturn], bool] = default_comparator,
-    a_revision: str = DEFAULT_A_REVISION,
-    b_revision: Optional[str] = None,
-):
-    """The same as git_ab_test_command, but via SSH. The closure argument should setup a microvm using the passed
-    paths to firecracker and jailer binaries."""
-    b_directory = (
-        DEFAULT_B_DIRECTORY
-        if b_revision is None
-        else FC_WORKSPACE_DIR / "build" / b_revision
-    )
-
-    def test_runner(bin_dir, _is_a: bool):
-        microvm = microvm_factory(bin_dir / "firecracker", bin_dir / "jailer")
-        return microvm.ssh.run(command)
-
-    (_, old_out, old_err), (_, new_out, new_err), the_same = binary_ab_test(
-        test_runner,
-        comparator,
-        a_directory=FC_WORKSPACE_DIR / "build" / a_revision,
-        b_directory=b_directory,
-    )
-
-    assert (
-        the_same
-    ), f"The output of running command `{command}` changed:\nOld:\nstdout:\n{old_out}\nstderr\n{old_err}\n\nNew:\nstdout:\n{new_out}\nstderr:\n{new_err}"
-
-
-def precompiled_ab_test_guest_command_if_pr(
-    microvm_factory: Callable[[Path, Path], Microvm],
-    command: str,
-    *,
-    comparator=default_comparator,
-    check_in_nonpr=True,
-):
-    """The same as git_ab_test_command_if_pr, but via SSH"""
-    if is_pr():
-        precompiled_ab_test_guest_command(
-            microvm_factory, command, comparator=comparator
-        )
-        return None
-
-    microvm = microvm_factory(*get_firecracker_binaries())
-    return microvm.ssh.run(command, check=check_in_nonpr)
-
-
 def check_regression(
     a_samples: List[float], b_samples: List[float], *, n_resamples: int = 9999
 ):

tests/framework/artifacts.py

@@ -44,9 +44,9 @@ def kernels(glob, artifact_dir: Path = ARTIFACT_DIR) -> Iterator:
                 break
 
 
-def disks(glob) -> Iterator:
+def disks(glob) -> list:
     """Return supported rootfs"""
-    yield from sorted(ARTIFACT_DIR.glob(glob))
+    return sorted(ARTIFACT_DIR.glob(glob))
 
 
 def kernel_params(
@@ -57,12 +57,6 @@ def kernel_params(
         yield pytest.param(kernel, id=kernel.name)
 
 
-def rootfs_params(glob="ubuntu-*.squashfs") -> Iterator:
-    """Return supported rootfs as pytest parameters"""
-    for rootfs in disks(glob=glob):
-        yield pytest.param(rootfs, id=rootfs.name)
-
-
 @dataclass(frozen=True, repr=True)
 class FirecrackerArtifact:
     """Utility class for Firecracker binary artifacts."""

tests/framework/microvm.py

@@ -75,6 +75,7 @@ class Snapshot:
     disks: dict
     ssh_key: Path
     snapshot_type: SnapshotType
+    meta: dict
 
     @property
     def is_diff(self) -> bool:
@@ -110,6 +111,7 @@ def copy_to_chroot(self, chroot) -> "Snapshot":
             disks=self.disks,
             ssh_key=self.ssh_key,
             snapshot_type=self.snapshot_type,
+            meta=self.meta,
         )
 
     @classmethod
@@ -125,6 +127,7 @@ def load_from(cls, src: Path) -> "Snapshot":
             disks={dsk: src / p for dsk, p in obj["disks"].items()},
             ssh_key=src / obj["ssh_key"],
             snapshot_type=SnapshotType(obj["snapshot_type"]),
+            meta=obj["meta"],
         )
 
     def save_to(self, dst: Path):
@@ -241,6 +244,7 @@ def __init__(
         self.disks_vhost_user = {}
         self.vcpus_count = None
         self.mem_size_bytes = None
+        self.cpu_template_name = None
 
         self._pre_cmd = []
         if numa_node:
@@ -732,12 +736,14 @@ def basic_config(
             smt=smt,
             mem_size_mib=mem_size_mib,
             track_dirty_pages=track_dirty_pages,
-            cpu_template=cpu_template,
             huge_pages=huge_pages,
         )
         self.vcpus_count = vcpu_count
         self.mem_size_bytes = mem_size_mib * 2**20
 
+        if cpu_template is not None:
+            self.set_cpu_template(cpu_template)
+
         if self.memory_monitor:
             self.memory_monitor.start()
 
@@ -770,6 +776,19 @@ def basic_config(
         if enable_entropy_device:
             self.enable_entropy_device()
 
+    def set_cpu_template(self, cpu_template):
+        """Set guest CPU template."""
+        if cpu_template is None:
+            return
+        # static CPU template
+        if isinstance(cpu_template, str):
+            self.api.machine_config.patch(cpu_template=cpu_template)
+            self.cpu_template_name = cpu_template.lower()
+        # custom CPU template
+        elif isinstance(cpu_template, dict):
+            self.api.cpu_config.put(**cpu_template["template"])
+            self.cpu_template_name = cpu_template["name"].lower()
+
     def add_drive(
         self,
         drive_id,
@@ -917,6 +936,9 @@ def make_snapshot(
             net_ifaces=[x["iface"] for ifname, x in self.iface.items()],
             ssh_key=self.ssh_key,
             snapshot_type=snapshot_type,
+            meta={
+                "kernel_file": self.kernel_file,
+            },
         )
 
     def snapshot_diff(self, *, mem_path: str = "mem", vmstate_path="vmstate"):
@@ -954,6 +976,9 @@ def restore_from_snapshot(
         if uffd_path is not None:
             mem_backend = {"backend_type": "Uffd", "backend_path": str(uffd_path)}
 
+        for key, value in snapshot.meta.items():
+            setattr(self, key, value)
+
         self.api.snapshot_load.put(
             mem_backend=mem_backend,
             snapshot_path=str(jailed_vmstate),
@@ -1059,6 +1084,13 @@ def build(self, kernel=None, rootfs=None, **kwargs):
             vm.ssh_key = ssh_key
         return vm
 
+    def build_from_snapshot(self, snapshot: Snapshot):
+        """Build a microvm from a snapshot"""
+        vm = self.build()
+        vm.spawn()
+        vm.restore_from_snapshot(snapshot, resume=True)
+        return vm
+
     def kill(self):
         """Clean up all built VMs"""
         for vm in self.vms: