"""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()