vmtest.vm: add --{build,insert}-test-kmod options

When running tests manually with vmtest.vm, it's tedious to set up the
test kernel module. Make vmtest.vm.run_in_vm() handle it and add command
line options.

Signed-off-by: Omar Sandoval <osandov@osandov.com>

setup.py

@@ -200,21 +200,14 @@ class test(Command):
         return make_check_success and test.result.wasSuccessful()
 
     def _run_vm(self, kernel):
-        from vmtest.kmod import build_kmod
         import vmtest.vm
 
         logger.info("running tests in VM on Linux %s", kernel.release)
 
-        try:
-            kmod = build_kmod(Path(self.vmtest_dir), kernel)
-        except subprocess.CalledProcessError:
-            return False
-
         command = rf"""
 set -e
 
 export PYTHON={shlex.quote(sys.executable)}
-export DRGN_TEST_KMOD={shlex.quote(str(kmod))}
 if [ -e /proc/vmcore ]; then
     "$PYTHON" -Bm unittest discover -t . -s tests/linux_kernel/vmcore {"-v" if self.verbose else ""}
 else
@@ -228,7 +221,11 @@ fi
 """
         try:
             returncode = vmtest.vm.run_in_vm(
-                command, kernel, Path("/"), Path(self.vmtest_dir)
+                command,
+                kernel,
+                Path("/"),
+                Path(self.vmtest_dir),
+                test_kmod=vmtest.vm.TestKmodMode.BUILD,
             )
         except vmtest.vm.LostVMError:
             logger.exception("error on Linux %s", kernel.release)

vmtest/__main__.py

@@ -24,9 +24,8 @@ from vmtest.download import (
     DownloadKernel,
     download_in_thread,
 )
-from vmtest.kmod import build_kmod
 from vmtest.rootfsbuild import build_drgn_in_rootfs
-from vmtest.vm import LostVMError, run_in_vm
+from vmtest.vm import LostVMError, TestKmodMode, run_in_vm
 
 logger = logging.getLogger(__name__)
 
@@ -296,7 +295,6 @@ chroot "$1" sh -c 'cd /mnt && pytest -v --ignore=tests/linux_kernel'
         for kernel in downloads:
             if not isinstance(kernel, Kernel):
                 continue
-            kmod = build_kmod(args.directory, kernel)
 
             # Skip excessively slow tests when emulating.
             if kernel.arch is HOST_ARCHITECTURE:
@@ -317,7 +315,6 @@ chroot "$1" sh -c 'cd /mnt && pytest -v --ignore=tests/linux_kernel'
 set -e
 
 export PYTHON={shlex.quote(sys.executable)}
-export DRGN_TEST_KMOD={shlex.quote(str(kmod))}
 export DRGN_RUN_LINUX_KERNEL_TESTS=1
 if [ -e /proc/vmcore ]; then
     "$PYTHON" -Bm pytest -v tests/linux_kernel/vmcore
@@ -333,6 +330,7 @@ fi
                     kernel,
                     args.directory / kernel.arch.name / "rootfs",
                     args.directory,
+                    test_kmod=TestKmodMode.BUILD,
                 )
             except LostVMError as e:
                 print("error:", e, file=sys.stderr)

vmtest/vm.py

@@ -1,6 +1,7 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # SPDX-License-Identifier: LGPL-2.1-or-later
 
+import enum
 import os
 from pathlib import Path
 import re
@@ -20,6 +21,7 @@ from vmtest.download import (
     download,
     download_kernel_argparse_type,
 )
+from vmtest.kmod import build_kmod
 
 # Script run as init in the virtual machine.
 _INIT_TEMPLATE = r"""#!/bin/sh
@@ -138,6 +140,7 @@ if [ -z "$vport" ]; then
 fi
 
 cd {cwd}
+{test_kmod}
 set +e
 setsid -c sh -c {command}
 rc=$?
@@ -192,12 +195,23 @@ def _build_onoatimehack(dir: Path) -> Path:
     return onoatimehack_so
 
 
+class TestKmodMode(enum.Enum):
+    NONE = 0
+    BUILD = 1
+    INSERT = 2
+
+
 class LostVMError(Exception):
     pass
 
 
 def run_in_vm(
-    command: str, kernel: Kernel, root_dir: Optional[Path], build_dir: Path
+    command: str,
+    kernel: Kernel,
+    root_dir: Optional[Path],
+    build_dir: Path,
+    *,
+    test_kmod: TestKmodMode = TestKmodMode.NONE,
 ) -> int:
     if root_dir is None:
         if kernel.arch is HOST_ARCHITECTURE:
@@ -205,6 +219,14 @@ def run_in_vm(
         else:
             root_dir = build_dir / kernel.arch.name / "rootfs"
 
+    if test_kmod == TestKmodMode.NONE:
+        test_kmod_command = ""
+    else:
+        kmod = build_kmod(build_dir, kernel)
+        test_kmod_command = f"export DRGN_TEST_KMOD={shlex.quote(str(kmod))}"
+        if test_kmod == TestKmodMode.INSERT:
+            test_kmod_command += '\ninsmod "$DRGN_TEST_KMOD"'
+
     qemu_exe = "qemu-system-" + kernel.arch.name
     match = re.search(
         r"QEMU emulator version ([0-9]+(?:\.[0-9]+)*)",
@@ -278,6 +300,7 @@ def run_in_vm(
                     ),
                     command=shlex.quote(command),
                     kdump_needs_nosmp="" if kvm_args else "export KDUMP_NEEDS_NOSMP=1",
+                    test_kmod=test_kmod_command,
                 )
             )
         init_path.chmod(0o755)
@@ -403,6 +426,22 @@ if __name__ == "__main__":
         type=Path,
         help="directory to use as root directory in VM (default: / for the host architecture, $directory/$arch/rootfs otherwise)",
     )
+    parser.add_argument(
+        "--build-test-kmod",
+        dest="test_kmod",
+        action="store_const",
+        const=TestKmodMode.BUILD,
+        default=argparse.SUPPRESS,
+        help="build the drgn test kernel module and define the DRGN_TEST_KMOD environment variable in the VM",
+    )
+    parser.add_argument(
+        "--insert-test-kmod",
+        dest="test_kmod",
+        action="store_const",
+        const=TestKmodMode.INSERT,
+        default=argparse.SUPPRESS,
+        help="insert the drgn test kernel module. Implies --build-test-kmod",
+    )
     parser.add_argument(
         "command",
         type=str,
@@ -420,6 +459,8 @@ if __name__ == "__main__":
         kernel = next(download(args.directory, [args.kernel]))  # type: ignore[assignment]
     if not hasattr(args, "root_directory"):
         args.root_directory = None
+    if not hasattr(args, "test_kmod"):
+        args.test_kmod = TestKmodMode.NONE
 
     try:
         command = (
@@ -427,7 +468,15 @@ if __name__ == "__main__":
             if args.command
             else "sh -i"
         )
-        sys.exit(run_in_vm(command, kernel, args.root_directory, args.directory))
+        sys.exit(
+            run_in_vm(
+                command,
+                kernel,
+                args.root_directory,
+                args.directory,
+                test_kmod=args.test_kmod,
+            )
+        )
     except LostVMError as e:
         print("error:", e, file=sys.stderr)
         sys.exit(args.lost_status)