move test code into a module

Signed-off-by: Zen <[email protected]>

Author: desultory

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ugrd"
-version = "1.16.0"
+version = "1.17.0"
 authors = [
   { name="Desultory", email="[email protected]" },
 ]

src/ugrd/base/test.py

@@ -0,0 +1,76 @@
+__version__ = "0.2.0"
+
+
+COPY_CONFIG = ['mounts', 'test_image_size', 'test_flag', 'out_dir', 'clean']
+
+
+def init_test_vars(self):
+    if not self.get("test_kernel"):
+        raise ValueError("test_kernel not set")
+    if not self['test_kernel'].exists():
+        raise FileNotFoundError("Test kernel not found: %s" % self['test_kernel'])
+
+
+def get_qemu_cmd_args(self):
+    """ Gets the qemu command from the configuration """
+    qemu_args = {'-m': self['test_memory'],
+                 '-cpu': self['test_cpu'],
+                 '-kernel': self['test_kernel'],
+                 '-initrd': self['_archive_out_path'],
+                 '-serial': 'mon:stdio',
+                 '-append': self['test_cmdline'],
+                 '-drive': 'file=%s,format=raw' % self['_test_rootfs']['_archive_out_path']}
+
+    qemu_bools = [f'-{item}' for item in self['qemu_bool_args']]
+
+    arglist = [f"qemu-system-{self['test_arch']}"] + qemu_bools
+    for key, value in qemu_args.items():
+        arglist.append(key)
+        arglist.append(value)
+
+    self['_qemu_cmd'] = ' '.join(str(arg) for arg in arglist)
+
+    return arglist
+
+
+def make_test_image(self):
+    """ Creates a new initramfs generator to creaate the test image """
+    from ugrd.initramfs_generator import InitramfsGenerator
+
+    kwargs = {'logger': self.logger,
+              'validate': False,
+              'NO_BASE': True,
+              'config': None,
+              'modules': 'ugrd.fs.test_image',
+              'out_file': self['test_rootfs_name'],
+              'build_dir': self['test_rootfs_build_dir'],
+              **{key: self[key] for key in COPY_CONFIG}}
+
+    target_fs = InitramfsGenerator(**kwargs)
+    target_fs.build()
+    self['_test_rootfs'] = target_fs
+
+
+def test_image(self):
+    qemu_cmd = get_qemu_cmd_args(self)
+    self.logger.info("Testing initramfs image: %s", self['_archive_out_path'])
+    self.logger.info("Test kernel: %s", self['test_kernel'])
+    self.logger.info("Test rootfs: %s", self['_test_rootfs']['_archive_out_path'])
+    self.logger.info("QEMU command: %s", self['_qemu_cmd'])
+
+    try:
+        results = self._run(qemu_cmd, timeout=self['test_timeout'])
+    except RuntimeError as e:
+        self.logger.error("Test failed: %s", e)
+        return False
+
+    stdout = results.stdout.decode('utf-8').split('\r\n')
+    self.logger.debug("QEMU output: %s", stdout)
+
+    if self['test_flag'] in stdout:
+        self.logger.info("Test passed")
+    else:
+        self.logger.error("Test failed")
+        self.logger.error("QEMU stdout:\n%s", stdout)
+
+

src/ugrd/base/test.toml

@@ -0,0 +1,29 @@
+test_rootfs_name = 'ugrd-test-rootfs'
+test_rootfs_build_dir = '/tmp/initramfs_test_rootfs'
+test_image_size = 16
+test_flag = "UGRD TEST SUCCESSFUL"
+
+test_memory = '256M'
+test_cpu = 'host'
+test_arch = 'x86_64'
+test_timeout = 15
+test_cmdline = 'console=ttyS0,115200 panic=1'
+qemu_bool_args = ['nographic', 'no-reboot', 'enable-kvm']
+
+[imports.tests]
+"ugrd.base.test" = ["init_test_vars", "make_test_image", "test_image"]
+
+[custom_parameters]
+test_kernel = "Path"  # Define the kernel to use for the test
+test_memory = "str"  # Define the amount of memory to use for the test image (passed to qemu)
+test_cpu = "str"  # Define the CPU to use for the test image (passed to qemu)
+test_arch = "str"  # Define the qemu arch (added to the qemu-system- command)
+test_cmdline = "str"  # Define the kernel command line for the test image
+test_timeout = "int"  # Define the timeout for the test
+test_image_size = "int"  # Define the size of the test image, in MB
+test_flag = "str"  # Define the success flag for the test
+test_rootfs_name = "str"  # Define the name of the rootfs image
+test_rootfs_build_dir = "Path"  # Define the build directory for the rootfs image
+qemu_bool_args = "NoDupFlatList"  # Define the qemu boolean arguments
+_test_rootfs = "InitramfsGenerator"  # Define the target rootfs generator
+_qemu_cmd = "str"  # Define the qemu command to use for the test

src/ugrd/fs/test_image.py

@@ -1,12 +1,7 @@
 __version__ = "0.3.0"
 
 
-def init_test_vars(self):
-    if not self.get('test_kernel'):
-        raise ValueError("No test kernel specified")
-    elif not self['test_kernel'].exists():
-        raise FileNotFoundError("Test kernel not found: %s" % self['test_kernel'])
-
+def init_banner(self):
     self['banner'] = f"echo {self['test_flag']}"
 
 

src/ugrd/fs/test_image.toml

@@ -1,36 +1,23 @@
 modules = ['ugrd.fs.cpio']
 paths = ['dev', 'sys', 'proc']
 
-
-binaries = ['halt', 'shutdown', 'ls', 'top', 'ps']
-shebang = "#!/bin/bash"
-
-qemu_arch = 'x86_64'
 test_image_size = 16
-test_memory = '256M'
-test_cpu = 'host'
-test_cmdline = 'console=ttyS0,115200 panic=1'
 test_flag = "UGRD TEST SUCCESSFUL"
-test_timeout = 15
+shebang = "#!/bin/bash"
+
 
 [masks]
 pack = "make_cpio"
 
 [imports.build_pre]
-"ugrd.fs.test_image" = ["init_test_vars"]
+"ugrd.fs.test_image" = ["init_banner"]
 
 [imports.pack]
 "ugrd.fs.test_image" = ["make_test_image"]
 
 [custom_parameters]
-shebang = "str"  # Add the shebang property, because the test mode disables the base module
-test_kernel = "Path"  # Define the kernel to use for the test image
-test_cmdline = "str"  # Define the kernel command line for the test image
-test_image_size = "int"  # Define the size of the test image in MiB
-test_memory = "str"  # Define the amount of memory to use for the test image (passed to qemu)
-test_cpu = "str"  # Define the CPU to use for the test image (passed to qemu)
-qemu_arch = "str"  # Define the qemu arch (added to the qemu-system- command)
+shebang = "str"  # Add the shebang property, because the test mode disables the base moduleo
 mounts = "dict"  # We only need the mounts dict, not the whole module
+test_image_size = "int"  # Define the size of the test image in MiB
 test_flag = "str"  # Define the success flag used to determine if the test was successful
-test_timeout = "int"  # Define the timeout for the test
 

src/ugrd/initramfs_dict.py

@@ -74,6 +74,8 @@ def handle_parameter(self, key: str, value) -> None:
         for d in (self.builtin_parameters, self['custom_parameters']):
             expected_type = d.get(key)
             if expected_type:
+                if expected_type.__name__ == "InitramfsGenerator":
+                    return super().__setitem__(key, value)
                 break
         else:
             raise KeyError("Parameter not registered: %s" % key)
@@ -125,6 +127,8 @@ def _process_custom_parameters(self, parameter_name: str, parameter_type: type)
         If the parameter is in the processing queue, process the queued values.
         """
         from pycpio import PyCPIO
+        from .initramfs_generator import InitramfsGenerator
+
         self['custom_parameters'][parameter_name] = eval(parameter_type)
         self.logger.debug("Registered custom parameter '%s' with type: %s" % (parameter_name, parameter_type))
 
@@ -139,12 +143,12 @@ def _process_custom_parameters(self, parameter_name: str, parameter_type: type)
                 super().__setitem__(parameter_name, 0)
             case "float":
                 super().__setitem__(parameter_name, 0.0)
-            case "PyCPIO":
-                super().__setitem__(parameter_name, PyCPIO(logger=self.logger, _log_init=False, _log_bump=10))
             case "str":
                 super().__setitem__(parameter_name, "")
             case "Path":
                 super().__setitem__(parameter_name, Path())
+            case "PyCPIO":
+                super().__setitem__(parameter_name, PyCPIO(logger=self.logger, _log_init=False, _log_bump=10))
             case _:  # For strings and things, don't init them so they are None
                 self.logger.warning("Leaving '%s' as None" % parameter_name)
                 super().__setitem__(parameter_name, None)

src/ugrd/initramfs_generator.py

@@ -76,6 +76,7 @@ def build(self) -> None:
         self.generate_init()
         self.run_hook('build_final')
         self.pack_build()
+        self.run_tests()
 
     def run_func(self, function, force_include=False) -> list[str]:
         """ Runs a function, If force_include is set, forces the function to be included in the bash source file. """
@@ -206,3 +207,14 @@ def pack_build(self) -> None:
         else:
             self.logger.warning("No pack functions specified, the final build is present in: %s" % self.build_dir)
 
+    def run_tests(self) -> None:
+        """ Runs tests if defined in self['imports']['tests']. """
+        if test_output := self.run_hook('tests'):
+            self.logger.info("Completed tests:\n%s", test_output)
+        else:
+            self.logger.debug("No tests executed.")
+
+
+
+
+

src/ugrd/initramfs_tester.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python
 
 from ugrd.initramfs_generator import InitramfsGenerator
-from ugrd.initramfs_tester import InitramfsTester
 from zenlib.util import get_kwargs_from_args, get_args_n_logger
 
 
@@ -47,12 +46,12 @@ def main():
     kwargs.pop('print_config', None)  # This is not a valid kwarg for InitramfsGenerator
     kwargs.pop('print_init', None)  # This is not a valid kwarg for InitramfsGenerator
     test = kwargs.pop('test', False)
-    test_kernel = kwargs.pop('test_kernel', None)
 
     if test:
         logger.warning("TEST MODE ENABLED")
         logger.info("Disabling DM autodetection")
         kwargs['autodetect_root_dm'] = False
+        kwargs['modules'] = 'ugrd.base.test'
 
     logger.debug(f"Using the following kwargs: {kwargs}")
     generator = InitramfsGenerator(**kwargs)
@@ -76,10 +75,6 @@ def main():
             for func in generator.imports[runlevel]:
                 print(f"    {func.__name__}")
 
-    if test:
-        tester = InitramfsTester(generator, test_kernel=test_kernel, logger=logger, _log_init=False)
-        tester.test()
-
 
 if __name__ == '__main__':
     main()

src/ugrd/main.py

@@ -1,6 +1,7 @@
 # NOP
 
 modules = [
+  "ugrd.base.test",
   "ugrd.kmod.standard_mask",
 ]
 
@@ -9,6 +10,10 @@ out_dir = "/tmp/initramfs_test"
 
 kernel_version = "6.6.35-gentoo-dist"
 test_kernel = "/boot/vmlinuz-6.6.35-gentoo-dist"
+cpio_compression = false
+
+autodetect_root_dm = false
+
 
 # Optionally supply a kernel version, uses the current kernel version if not specified
 #kernel_version = "6.1.53-gentoo-dist"