aboutsummaryrefslogtreecommitdiff
path: root/tests/nvidia-preflight/test_nvidia_preflight.py
blob: bdacfd5e69588e7272b2202e268128e2afc09f41 (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
"""Tests for the nvidia_preflight_report helper in the archsetup installer.

nvidia_preflight_report is the pure core of the NVIDIA/Wayland preflight
check: it scans DRM (then PCI display-class) modalias files for the NVIDIA
vendor id, and when one matches it prints the Wayland warning + required
environment variables and checks the repo's candidate nvidia-utils major
version. Return codes: 0 = no NVIDIA GPU, 10 = NVIDIA and the driver
requirement (535+) is met, 11 = NVIDIA and the requirement is not met
(driver too old or unknown). The interactive continue/abort prompt lives in
preflight_checks, not here, so this core is unit testable.

These tests exercise the REAL function body, extracted from the `archsetup`
script at run time (not a copy), against temp modalias trees and a fake
pacman on PATH.

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

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_MODALIAS = "pci:v000010DEd00002684sv00001043sd000088E2bc03sc00i00"
NVIDIA_MODALIAS_LOWER = "pci:v000010ded00002684sv00001043sd000088e2bc03sc00i00"
AMD_MODALIAS = "pci:v00001002d0000164Esv00001462sd00007D78bc03sc80i00"
NON_DISPLAY_NVIDIA = "pci:v000010DEd00002684sv00001043sd000088E2bc0Csc03i30"


class NvidiaPreflightHarness(unittest.TestCase):
    """Source nvidia_preflight_report out of the real archsetup script."""

    def setUp(self):
        self.tmp = tempfile.mkdtemp(prefix="nvidia-preflight-test-")
        self.drm = os.path.join(self.tmp, "drm")
        self.pci = os.path.join(self.tmp, "pci")
        os.makedirs(self.drm)
        os.makedirs(self.pci)
        self.fakebin = os.path.join(self.tmp, "bin")
        os.makedirs(self.fakebin)
        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_report() {/,/^}/p' \"$ARCHSETUP\")\n"
                "nvidia_preflight_report\n"
            )
        os.chmod(self.wrapper, 0o755)

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

    def fake_pacman(self, version=None, fail=False):
        """A pacman stub answering `pacman -Si nvidia-utils`."""
        path = os.path.join(self.fakebin, "pacman")
        with open(path, "w") as f:
            if fail:
                f.write("#!/bin/sh\nexit 1\n")
            else:
                f.write(
                    "#!/bin/sh\n"
                    "printf 'Repository      : extra\\n'\n"
                    "printf 'Name            : nvidia-utils\\n'\n"
                    "printf 'Version         : %s\\n'\n" % version
                )
        os.chmod(path, 0o755)

    def add_modalias(self, root, subdir, content):
        d = os.path.join(root, subdir)
        os.makedirs(d, exist_ok=True)
        with open(os.path.join(d, "modalias"), "w") as f:
            f.write(content + "\n")

    def run_check(self):
        env = dict(os.environ)
        env["PATH"] = self.fakebin + os.pathsep + env["PATH"]
        env["NVIDIA_DRM_GLOB"] = os.path.join(self.drm, "card*", "modalias")
        env["NVIDIA_PCI_GLOB"] = os.path.join(self.pci, "*", "modalias")
        return subprocess.run(
            ["bash", self.wrapper, ARCHSETUP],
            capture_output=True, text=True, env=env,
        )

    # ---------------------------------------------------------- normal ----
    def test_no_gpu_files_returns_zero_and_silent(self):
        self.fake_pacman(version="575.51.02-1")
        r = self.run_check()
        self.assertEqual(r.returncode, 0)
        self.assertNotIn("NVIDIA", r.stdout)

    def test_amd_only_returns_zero(self):
        self.fake_pacman(version="575.51.02-1")
        self.add_modalias(self.drm, "card0", AMD_MODALIAS)
        r = self.run_check()
        self.assertEqual(r.returncode, 0)
        self.assertNotIn("NVIDIA", r.stdout)

    def test_nvidia_with_modern_driver_returns_ten_with_guidance(self):
        self.fake_pacman(version="575.51.02-1")
        self.add_modalias(self.drm, "card0", NVIDIA_MODALIAS)
        r = self.run_check()
        self.assertEqual(r.returncode, 10)
        self.assertIn("NVIDIA GPU detected", r.stdout)
        self.assertIn("LIBVA_DRIVER_NAME=nvidia", r.stdout)
        self.assertIn("GBM_BACKEND=nvidia-drm", r.stdout)
        self.assertIn("__GLX_VENDOR_LIBRARY_NAME=nvidia", r.stdout)
        self.assertIn("575.51.02-1", r.stdout)

    # -------------------------------------------------------- boundary ----
    def test_lowercase_vendor_id_detected(self):
        self.fake_pacman(version="575.51.02-1")
        self.add_modalias(self.drm, "card0", NVIDIA_MODALIAS_LOWER)
        r = self.run_check()
        self.assertEqual(r.returncode, 10)

    def test_exactly_535_meets_requirement(self):
        self.fake_pacman(version="535.216.01-1")
        self.add_modalias(self.drm, "card0", NVIDIA_MODALIAS)
        r = self.run_check()
        self.assertEqual(r.returncode, 10)

    def test_pci_fallback_display_class_only(self):
        # No DRM entries; PCI holds a display-class NVIDIA device -> detected.
        self.fake_pacman(version="575.51.02-1")
        self.add_modalias(self.pci, "0000:01:00.0", NVIDIA_MODALIAS)
        r = self.run_check()
        self.assertEqual(r.returncode, 10)

    def test_pci_non_display_nvidia_ignored(self):
        # An NVIDIA audio/usb function (bc0C) must not trigger the check.
        self.fake_pacman(version="575.51.02-1")
        self.add_modalias(self.pci, "0000:01:00.1", NON_DISPLAY_NVIDIA)
        r = self.run_check()
        self.assertEqual(r.returncode, 0)

    # ----------------------------------------------------------- error ----
    def test_old_driver_returns_eleven_with_error(self):
        self.fake_pacman(version="470.256.02-1")
        self.add_modalias(self.drm, "card0", NVIDIA_MODALIAS)
        r = self.run_check()
        self.assertEqual(r.returncode, 11)
        self.assertIn("535", r.stdout)
        self.assertIn("470.256.02-1", r.stdout)

    def test_pacman_failure_returns_eleven_unknown(self):
        self.fake_pacman(fail=True)
        self.add_modalias(self.drm, "card0", NVIDIA_MODALIAS)
        r = self.run_check()
        self.assertEqual(r.returncode, 11)
        self.assertIn("unknown", r.stdout)


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