1 files changed, 95 insertions, 0 deletions
diff --git a/tests/installer-steps/test_ensure_nvme_early_module.py b/tests/installer-steps/test_ensure_nvme_early_module.py
new file mode 100644
index 0000000..8e3f036
--- /dev/null
+++ b/tests/installer-steps/test_ensure_nvme_early_module.py
@@ -0,0 +1,95 @@
+"""Test the nvme early-module step.
+
+Two bugs the extracted ensure_nvme_early_module fixes:
+
+1. The MODULES=(... nvme) edit was only compiled into the initramfs by a
+   mkinitcpio -P that ran behind `if ! is_zfs_root`, so on ZFS-root machines
+   the early-load hardening never landed. The step now rebuilds whenever it
+   changed the file, unconditionally.
+2. The already-present check grepped the whole file for "nvme", so a comment
+   or an unrelated token (nvme_tcp) anywhere in mkinitcpio.conf skipped the
+   edit. The check now looks for the nvme word on the MODULES line only.
+
+Method: sed-extract ensure_nvme_early_module from the real `archsetup`, point
+it at a temp mkinitcpio.conf, and fake has_nvme_drives/backup_system_file/
+mkinitcpio/display/error_warn.
+
+Run from repo root:
+    python3 -m unittest tests.installer-steps.test_ensure_nvme_early_module
+"""
+
+import os
+import subprocess
+import tempfile
+import textwrap
+import unittest
+
+REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+ARCHSETUP = os.path.join(REPO_ROOT, "archsetup")
+
+
+def run(conf_body, has_nvme=True):
+    with tempfile.TemporaryDirectory() as d:
+        conf = os.path.join(d, "mkinitcpio.conf")
+        with open(conf, "w") as f:
+            f.write(conf_body)
+        script = textwrap.dedent(f"""\
+            logfile=/dev/null
+            action=""
+            display() {{ :; }}
+            backup_system_file() {{ :; }}
+            error_warn() {{ echo "WARN: $1"; return 1; }}
+            has_nvme_drives() {{ return {0 if has_nvme else 1}; }}
+            # stdout is redirected into $logfile at the call site, so the fake
+            # records its invocation in a side file instead.
+            mkinitcpio() {{ echo "MKINITCPIO $*" >> "{d}/mk.log"; }}
+            source <(sed -n '/^ensure_nvme_early_module() {{/,/^}}/p' "{ARCHSETUP}")
+            ensure_nvme_early_module "{conf}"
+            echo "RC=$?"
+            echo "CONF:[$(cat "{conf}")]"
+            [ -f "{d}/mk.log" ] && cat "{d}/mk.log"
+        """)
+        return subprocess.run(
+            ["bash", "-c", script], capture_output=True, text=True, timeout=10,
+        )
+
+
+class EnsureNvmeEarlyModule(unittest.TestCase):
+    def test_empty_modules_gets_nvme_and_rebuilds(self):
+        r = run("MODULES=()\nHOOKS=(base udev)\n")
+        self.assertIn("MODULES=(nvme)", r.stdout)
+        self.assertIn("MKINITCPIO -P", r.stdout,
+                      "the initramfs must rebuild after the MODULES edit")
+
+    def test_populated_modules_appends_nvme_and_rebuilds(self):
+        r = run("MODULES=(btrfs)\n")
+        self.assertIn("MODULES=(btrfs nvme)", r.stdout)
+        self.assertIn("MKINITCPIO -P", r.stdout)
+
+    def test_nvme_already_present_no_edit_no_rebuild(self):
+        r = run("MODULES=(nvme)\n")
+        self.assertIn("MODULES=(nvme)", r.stdout)
+        self.assertNotIn("MODULES=(nvme nvme)", r.stdout)
+        self.assertNotIn("MKINITCPIO", r.stdout,
+                         "an unchanged conf must not trigger a rebuild (resume idempotence)")
+
+    def test_nvme_mention_elsewhere_does_not_skip_the_edit(self):
+        # The old whole-file grep skipped the edit when any line mentioned
+        # nvme; a comment must not mask a missing MODULES entry.
+        r = run("# early nvme notes: nvme_load=YES\nMODULES=(btrfs)\n")
+        self.assertIn("MODULES=(btrfs nvme)", r.stdout)
+        self.assertIn("MKINITCPIO -P", r.stdout)
+
+    def test_nvme_tcp_module_does_not_count_as_nvme(self):
+        r = run("MODULES=(nvme_tcp)\n")
+        self.assertIn("MODULES=(nvme_tcp nvme)", r.stdout)
+
+    def test_no_nvme_drives_is_a_noop(self):
+        r = run("MODULES=()\n", has_nvme=False)
+        self.assertIn("RC=0", r.stdout)
+        self.assertIn("CONF:[MODULES=()", r.stdout)
+        self.assertNotIn("MKINITCPIO", r.stdout)
+
+
+if __name__ == "__main__":
+    unittest.main()