python3Packages.pynvml: add a gpu test

pkgs/development/python-modules/pynvml/default.nix

@@ -1,6 +1,7 @@
 {
   lib,
   buildPythonPackage,
+  callPackage,
   fetchFromGitHub,
   substituteAll,
   pythonOlder,
@@ -50,6 +51,9 @@ buildPythonPackage rec {
   # OSError: /run/opengl-driver/lib/libnvidia-ml.so.1: cannot open shared object file: No such file or directory
   doCheck = false;
 
+  passthru.tests.nvmlInit = callPackage ./test-gpu.nix { };
+
+
   meta = with lib; {
     description = "Python bindings for the NVIDIA Management Library";
     homepage = "https://github.com/gpuopenanalytics/pynvml";

pkgs/development/python-modules/pynvml/test-gpu.nix

@@ -0,0 +1,23 @@
+{ runCommandNoCC
+, python
+}:
+
+runCommandNoCC "pynvml-gpu-test"
+{
+  nativeBuildInputs = [
+    (python.withPackages (ps: [ ps.pynvml ]))
+  ];
+  requiredSystemFeatures = [
+    "cuda"
+  ];
+} ''
+  python3 << EOF
+  import pynvml
+  from pynvml.smi import nvidia_smi
+
+  pynvml.nvmlInit()
+  EOF
+
+  touch $out
+''
+