From 7ed2cba5e8882c93516f71e3de2bf8a828ccdb06 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Wed, 20 Sep 2023 07:02:35 +0300
Subject: [PATCH] python3Packages.pynvml: add a gpu test

---
 .../python-modules/pynvml/default.nix         |  4 ++++
 .../python-modules/pynvml/test-gpu.nix        | 23 +++++++++++++++++++
 2 files changed, 27 insertions(+)
 create mode 100644 pkgs/development/python-modules/pynvml/test-gpu.nix

diff --git a/pkgs/development/python-modules/pynvml/default.nix b/pkgs/development/python-modules/pynvml/default.nix
index a115cd723998..e76913a63df4 100644
--- a/pkgs/development/python-modules/pynvml/default.nix
+++ b/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";
diff --git a/pkgs/development/python-modules/pynvml/test-gpu.nix b/pkgs/development/python-modules/pynvml/test-gpu.nix
new file mode 100644
index 000000000000..c316d0b5094b
--- /dev/null
+++ b/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
+''
+