aboutsummaryrefslogtreecommitdiff
path: root/tests/nvidia-preflight/test_nvidia_preflight_gate.py
blob: 1d766896b163aebbcca12151641af4d1a732fb6b (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
"""Tests for the nvidia_preflight policy wrapper in the archsetup installer.

nvidia_preflight_report (covered by test_nvidia_preflight.py) is the pure
detection core. nvidia_preflight is the policy around it: it decides whether
the check applies at all, runs the report, aborts on a too-old driver, and
confirms an otherwise-supported card with the user.

The check is Wayland-specific -- every line it prints names Wayland/Hyprland,
and its own failure hint tells the user to install with DESKTOP_ENV=dwm
instead. So it must not fire on a dwm (X11) or headless (none) install, and it
must not fire when --no-gpu-drivers / NO_GPU_DRIVERS=yes says the user handles
the driver themselves -- the same ruling install_gpu_drivers already makes.

These tests exercise the REAL function body, extracted from the `archsetup`
script at run time (not a copy), with nvidia_preflight_report stubbed to a
chosen verdict.

Run from repo root:
    python3 -m unittest tests.nvidia-preflight.test_nvidia_preflight_gate
"""

import os
import shutil
import subprocess
import tempfile
import unittest


REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
ARCHSETUP = os.path.join(REPO_ROOT, "archsetup")

# nvidia_preflight_report's documented return codes.
RC_NO_NVIDIA = 0
RC_NVIDIA_OK = 10
RC_NVIDIA_TOO_OLD = 11


class NvidiaPreflightGateHarness(unittest.TestCase):
    """Source nvidia_preflight out of the real archsetup script."""

    def setUp(self):
        self.tmp = tempfile.mkdtemp(prefix="nvidia-gate-test-")
        self.marker = os.path.join(self.tmp, "report-calls")
        self.wrapper = os.path.join(self.tmp, "run.sh")
        with open(self.wrapper, "w") as f:
            f.write(
                "#!/bin/bash\n"
                'ARCHSETUP="$1"; shift\n'
                "source <(sed -n "
                "'/^nvidia_preflight() {/,/^}/p' \"$ARCHSETUP\")\n"
                # Stub the detection core: record the call, answer with $STUB_RC.
                "nvidia_preflight_report() {\n"
                '    echo "called" >> "$MARKER"\n'
                '    return "${STUB_RC:-0}"\n'
                "}\n"
                "nvidia_preflight\n"
                'echo "PREFLIGHT_RETURNED"\n'
                # Whatever the prompt did not consume is still on stdin.
                'if IFS= read -r leftover; then echo "LEFTOVER:$leftover"; '
                'else echo "LEFTOVER:"; fi\n'
            )
        os.chmod(self.wrapper, 0o755)

    def tearDown(self):
        shutil.rmtree(self.tmp, ignore_errors=True)

    def run_gate(self, rc, desktop_env="hyprland", skip_gpu_drivers="false",
                 config_file="", stdin=""):
        env = dict(os.environ)
        env["MARKER"] = self.marker
        env["STUB_RC"] = str(rc)
        env["desktop_env"] = desktop_env
        env["skip_gpu_drivers"] = skip_gpu_drivers
        env["config_file"] = config_file
        return subprocess.run(
            ["bash", self.wrapper, ARCHSETUP],
            input=stdin, capture_output=True, text=True, env=env,
        )

    def report_calls(self):
        if not os.path.exists(self.marker):
            return 0
        with open(self.marker) as f:
            return len([line for line in f if line.strip()])

    def assertContinued(self, r):
        self.assertIn("PREFLIGHT_RETURNED", r.stdout,
                      "preflight aborted; stdout=%r stderr=%r" % (r.stdout, r.stderr))
        self.assertEqual(r.returncode, 0)

    def assertAborted(self, r):
        self.assertNotIn("PREFLIGHT_RETURNED", r.stdout)
        self.assertEqual(r.returncode, 1)

    # ---------------------------------------------------------- normal ----
    def test_hyprland_no_nvidia_continues_without_prompting(self):
        r = self.run_gate(RC_NO_NVIDIA, stdin="sentinel\n")
        self.assertContinued(r)
        self.assertEqual(self.report_calls(), 1)
        self.assertIn("LEFTOVER:sentinel", r.stdout)

    def test_hyprland_supported_driver_prompts_and_continues_on_yes(self):
        # bash emits `read -p`'s prompt only when stdin is a terminal, so the
        # prompt text is unassertable here. What is assertable -- and is the
        # thing that matters -- is that the read consumed exactly the answer
        # line and left the rest of stdin alone.
        r = self.run_gate(RC_NVIDIA_OK, stdin="y\nsentinel\n")
        self.assertContinued(r)
        self.assertIn("Continuing on NVIDIA", r.stdout)
        self.assertIn("LEFTOVER:sentinel", r.stdout)

    def test_hyprland_old_driver_still_aborts(self):
        # The check must keep working where it actually applies.
        r = self.run_gate(RC_NVIDIA_TOO_OLD)
        self.assertAborted(r)
        self.assertEqual(self.report_calls(), 1)

    # -------------------------------------------------------- boundary ----
    def test_dwm_old_driver_does_not_abort(self):
        # An X11 install does not care about the Wayland driver floor -- and
        # the abort's own hint tells the user to install with DESKTOP_ENV=dwm.
        r = self.run_gate(RC_NVIDIA_TOO_OLD, desktop_env="dwm")
        self.assertContinued(r)
        self.assertEqual(self.report_calls(), 0)

    def test_none_old_driver_does_not_abort(self):
        # A headless install never runs a compositor at all.
        r = self.run_gate(RC_NVIDIA_TOO_OLD, desktop_env="none")
        self.assertContinued(r)
        self.assertEqual(self.report_calls(), 0)

    def test_skip_gpu_drivers_old_driver_does_not_abort(self):
        # --no-gpu-drivers means the user installs the driver themselves, so
        # the repo's candidate version is not archsetup's call to veto.
        r = self.run_gate(RC_NVIDIA_TOO_OLD, skip_gpu_drivers="true")
        self.assertContinued(r)
        self.assertEqual(self.report_calls(), 0)

    def test_dwm_supported_driver_does_not_prompt(self):
        r = self.run_gate(RC_NVIDIA_OK, desktop_env="dwm", stdin="sentinel\n")
        self.assertContinued(r)
        self.assertIn("LEFTOVER:sentinel", r.stdout)

    def test_unattended_config_file_answers_the_prompt_itself(self):
        # --config-file is the documented unattended mode. The rc-10 prompt is
        # advisory (the driver requirement is met), so it must not block on a
        # read -- the same failure aur_install's --answerdiff/--answerclean
        # comment already rules out.
        r = self.run_gate(RC_NVIDIA_OK, config_file="/etc/archsetup.conf",
                          stdin="sentinel\n")
        self.assertContinued(r)
        self.assertIn("Continuing on NVIDIA", r.stdout)
        self.assertIn("LEFTOVER:sentinel", r.stdout)

    # ----------------------------------------------------------- error ----
    def test_hyprland_supported_driver_aborts_on_no(self):
        r = self.run_gate(RC_NVIDIA_OK, stdin="n\n")
        self.assertAborted(r)
        self.assertIn("Aborted at NVIDIA preflight.", r.stdout)

    def test_hyprland_supported_driver_aborts_on_capital_no(self):
        r = self.run_gate(RC_NVIDIA_OK, stdin="No\n")
        self.assertAborted(r)

    def test_unattended_old_driver_still_aborts(self):
        # Unattended does not mean "never fail" -- a driver below the floor is
        # still a hard stop on a hyprland install.
        r = self.run_gate(RC_NVIDIA_TOO_OLD, config_file="/etc/archsetup.conf")
        self.assertAborted(r)


if __name__ == "__main__":
    unittest.main()