test: Add pytest option to apply custom CPU template to any microVM

Sometimes we would like to test a custom CPU template for debugging /
testing purposes. The new pytest option `--custom-cpu-template` allows
us to apply the given CPU template to any microVM in any test unless it
is overwritten by a test. Since `cpu_template_any` yields not only CPU
templates but also None, the given CPU template will be used in the None
case.

Signed-off-by: Takahiro Itazuri <itazur@amazon.com>

Author: zulinx86

tests/conftest.py

@@ -23,6 +23,7 @@
 """
 
 import inspect
+import json
 import os
 import re
 import shutil
@@ -66,6 +67,14 @@ def pytest_addoption(parser):
         help="use firecracker/jailer binaries from this directory instead of compiling from source",
     )
 
+    parser.addoption(
+        "--custom-cpu-template",
+        action="store",
+        help="Path to custom CPU template to be applied unless overwritten by a test",
+        default=None,
+        type=Path,
+    )
+
 
 @pytest.hookimpl(wrapper=True, tryfirst=True)
 def pytest_runtest_makereport(item, call):  # pylint:disable=unused-argument
@@ -270,9 +279,22 @@ def microvm_factory(request, record_property, results_dir):
         fc_binary_path, jailer_binary_path = build_tools.get_firecracker_binaries()
     record_property("firecracker_bin", str(fc_binary_path))
 
+    # If `--custom-cpu-template` option is provided, the given CPU template will
+    # be applied afterwards unless overwritten.
+    custom_cpu_template_path = request.config.getoption("--custom-cpu-template")
+    custom_cpu_template = (
+        {
+            "name": custom_cpu_template_path.stem,
+            "template": json.loads(custom_cpu_template_path.read_text("utf-8")),
+        }
+        if custom_cpu_template_path
+        else None
+    )
     # We could override the chroot base like so
     # jailer_kwargs={"chroot_base": "/srv/jailo"}
-    uvm_factory = MicroVMFactory(fc_binary_path, jailer_binary_path)
+    uvm_factory = MicroVMFactory(
+        fc_binary_path, jailer_binary_path, custom_cpu_template=custom_cpu_template
+    )
     yield uvm_factory
 
     # if the test failed, save important files from the root of the uVM into `test_results` for troubleshooting

tests/framework/microvm.py

@@ -188,6 +188,7 @@ def __init__(
         monitor_memory: bool = True,
         jailer_kwargs: Optional[dict] = None,
         numa_node=None,
+        custom_cpu_template: Path = None,
     ):
         """Set up microVM attributes, paths, and data structures."""
         # pylint: disable=too-many-statements
@@ -246,6 +247,9 @@ def __init__(
         self.vcpus_count = None
         self.mem_size_bytes = None
         self.cpu_template_name = None
+        # The given custom CPU template will be set in basic_config() but could
+        # be overwritten via set_cpu_template().
+        self.custom_cpu_template = custom_cpu_template
 
         self._connections = []
 
@@ -748,6 +752,9 @@ def basic_config(
         self.vcpus_count = vcpu_count
         self.mem_size_bytes = mem_size_mib * 2**20
 
+        if self.custom_cpu_template is not None:
+            self.set_cpu_template(self.custom_cpu_template)
+
         if cpu_template is not None:
             self.set_cpu_template(cpu_template)