tests: refactor test_vulnerabilities

Some further simplifications:
- Simplify the "_host" tests as they will provide the same result in
  both A/B situations.
- Add a second MicrovmFactory fixture with the "A" firecracker.
- Add a uvm_any fixture that includes all CPU templates and also a
  boot/restored dimension.

This decreases the need for separate tests. For example the guest test
test_check_vulnerability_files_ab can now test all variants:

2 (restored/booted) * 3 (kernels) * 9 (cpu templates) = 54 tests

Signed-off-by: Pablo Barbáchano <[email protected]>

Author: pb8o

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

@@ -459,3 +459,38 @@ def uvm_with_initrd(
     uvm = microvm_factory.build(guest_kernel_linux_5_10)
     uvm.initrd_file = fs
     yield uvm
+
+
+def uvm_booted(microvm_factory, guest_kernel, rootfs, cpu_template):
+    """Return a booted uvm"""
+    uvm = microvm_factory.build(guest_kernel, rootfs)
+    uvm.spawn()
+    uvm.basic_config(vcpu_count=2, mem_size_mib=256)
+    uvm.set_cpu_template(cpu_template)
+    uvm.add_net_iface()
+    uvm.start()
+    return uvm
+
+
+def uvm_restored(microvm_factory, guest_kernel, rootfs, cpu_template):
+    """Return a restored uvm"""
+    uvm = uvm_booted(microvm_factory, guest_kernel, rootfs, cpu_template)
+    snapshot = uvm.snapshot_full()
+    uvm.kill()
+    uvm2 = microvm_factory.build()
+    uvm2.spawn()
+    uvm2.restore_from_snapshot(snapshot, resume=True)
+    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):
+    """Return booted and restored uvms"""
+    return uvm_ctor(microvm_factory, guest_kernel, rootfs, cpu_template_any)

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/microvm.py

@@ -244,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:
@@ -735,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()
 
@@ -773,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,

tests/framework/properties.py

@@ -72,11 +72,14 @@ def __init__(self):
         # major.minor.patch
         self.host_linux_patch = get_kernel_version(2)
         self.os = get_os_version()
-        self.host_os = get_host_os()
+        self.host_os = get_host_os() or "NA"
         self.libc_ver = "-".join(platform.libc_ver())
         self.rust_version = run_cmd("rustc --version |awk '{print $2}'")
+        # Buildkite/PR information
         self.buildkite_pipeline_slug = os.environ.get("BUILDKITE_PIPELINE_SLUG")
         self.buildkite_build_number = os.environ.get("BUILDKITE_BUILD_NUMBER")
+        self.buildkite_pr = os.environ.get("BUILDKITE_PULL_REQUEST", "false") != "false"
+        self.buildkite_revision_a = os.environ.get("BUILDKITE_PULL_REQUEST_BASE_BRANCH")
 
         if self._in_git_repo():
             self.git_commit_id = run_cmd("git rev-parse HEAD")

tests/framework/utils_cpu_templates.py

@@ -3,6 +3,8 @@
 
 """Utilities for CPU template related functionality."""
 
+# pylint:disable=too-many-return-statements
+
 import json
 from pathlib import Path
 
@@ -23,9 +25,7 @@ def get_supported_cpu_templates():
     """
     Return the list of CPU templates supported by the platform.
     """
-    # pylint:disable=too-many-return-statements
     host_linux = global_props.host_linux_version_tpl
-
     match get_cpu_vendor(), global_props.cpu_codename:
         # T2CL template is only supported on Cascade Lake and newer CPUs.
         case CpuVendor.INTEL, CpuModel.INTEL_SKYLAKE:

tests/integration_tests/functional/test_net.py

@@ -84,14 +84,23 @@ def test_multi_queue_unsupported(uvm_plain):
         )
 
 
-def run_udp_offload_test(vm):
+@pytest.fixture
+def uvm_any(microvm_factory, uvm_ctor, guest_kernel, rootfs):
+    """Return booted and restored uvm with no CPU templates"""
+    return uvm_ctor(microvm_factory, guest_kernel, rootfs, None)
+
+
+def test_tap_offload(uvm_any):
     """
+    Verify that tap offload features are configured for a booted/restored VM.
+
     - Start a socat UDP server in the guest.
     - Try to send a UDP message with UDP offload enabled.
 
     If tap offload features are not configured, an attempt to send a message will fail with EIO "Input/output error".
     More info (search for "TUN_F_CSUM is a must"): https://blog.cloudflare.com/fr-fr/virtual-networking-101-understanding-tap/
     """
+    vm = uvm_any
     port = "81"
     out_filename = "/tmp/out.txt"
     message = "x"
@@ -112,35 +121,3 @@ def run_udp_offload_test(vm):
     # Check that the server received the message
     ret = vm.ssh.run(f"cat {out_filename}")
     assert ret.stdout == message, f"{ret.stdout=} {ret.stderr=}"
-
-
-def test_tap_offload_booted(uvm_plain_any):
-    """
-    Verify that tap offload features are configured for a booted VM.
-    """
-    vm = uvm_plain_any
-    vm.spawn()
-    vm.basic_config()
-    vm.add_net_iface()
-    vm.start()
-
-    run_udp_offload_test(vm)
-
-
-def test_tap_offload_restored(microvm_factory, guest_kernel, rootfs):
-    """
-    Verify that tap offload features are configured for a restored VM.
-    """
-    src = microvm_factory.build(guest_kernel, rootfs, monitor_memory=False)
-    src.spawn()
-    src.basic_config()
-    src.add_net_iface()
-    src.start()
-    snapshot = src.snapshot_full()
-    src.kill()
-
-    dst = microvm_factory.build()
-    dst.spawn()
-    dst.restore_from_snapshot(snapshot, resume=True)
-
-    run_udp_offload_test(dst)