cuda.bindings.nvml: Correctly handle a test requiring root

Author: mdboom

cuda_bindings/tests/nvml/conftest.py

@@ -1,7 +1,9 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
+import os
 from collections import namedtuple
+from contextlib import contextmanager
 
 import pytest
 from cuda.bindings import _nvml as nvml
@@ -128,3 +130,12 @@ def pci_info(ngpus, handles):
     pci_info = [nvml.device_get_pci_info_v3(handles[i]) for i in range(ngpus)]
     assert len(pci_info) == ngpus
     return pci_info
+
+
+@contextmanager
+def requires_root():
+    if os.geteuid() == 0:
+        yield
+    else:
+        with pytest.raises(nvml.NoPermissionError):
+            yield

cuda_bindings/tests/nvml/test_compute_mode.py

@@ -7,6 +7,8 @@
 import pytest
 from cuda.bindings import _nvml as nvml
 
+from . import conftest
+
 COMPUTE_MODES = [
     nvml.ComputeMode.COMPUTEMODE_DEFAULT,
     nvml.ComputeMode.COMPUTEMODE_PROHIBITED,
@@ -25,7 +27,7 @@ def test_compute_mode_supported_nonroot(all_devices):
             continue
 
         for cm in COMPUTE_MODES:
-            with pytest.raises(nvml.NoPermissionError):
+            with conftest.requires_root():
                 nvml.device_set_compute_mode(device, cm)
             assert original_compute_mode == nvml.device_get_compute_mode(device), "Compute mode shouldn't have changed"