From 8d52c10716d4753aaf811e7c283f9f69b453db94 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Fern=C3=A1ndez=20L=C3=B3pez?=
 <ereslibre@ereslibre.es>
Date: Fri, 16 Aug 2024 20:33:50 +0200
Subject: [PATCH] nvidia-container-toolkit: add initial set of tests to check
 closures

---
 .../nvidia-container-toolkit/default.nix      |   3 +
 nixos/tests/all-tests.nix                     |   1 +
 nixos/tests/nvidia-container-toolkit.nix      | 158 ++++++++++++++++++
 3 files changed, 162 insertions(+)
 create mode 100644 nixos/tests/nvidia-container-toolkit.nix

diff --git a/nixos/modules/services/hardware/nvidia-container-toolkit/default.nix b/nixos/modules/services/hardware/nvidia-container-toolkit/default.nix
index fe4999bfbdbbe..3936b2757c462 100644
--- a/nixos/modules/services/hardware/nvidia-container-toolkit/default.nix
+++ b/nixos/modules/services/hardware/nvidia-container-toolkit/default.nix
@@ -71,6 +71,8 @@
           /usr/local/nvidia/lib64.
         '';
       };
+
+      package = lib.mkPackageOption pkgs "nvidia-container-toolkit" { };
     };
 
   };
@@ -129,6 +131,7 @@
           let
             script = pkgs.callPackage ./cdi-generate.nix {
               inherit (config.hardware.nvidia-container-toolkit) mounts;
+              nvidia-container-toolkit = config.hardware.nvidia-container-toolkit.package;
               nvidia-driver = config.hardware.nvidia.package;
               deviceNameStrategy = config.hardware.nvidia-container-toolkit.device-name-strategy;
             };
diff --git a/nixos/tests/all-tests.nix b/nixos/tests/all-tests.nix
index e00c6b2b39d3d..efe14ed699b69 100644
--- a/nixos/tests/all-tests.nix
+++ b/nixos/tests/all-tests.nix
@@ -699,6 +699,7 @@ in {
   ntfy-sh = handleTest ./ntfy-sh.nix {};
   ntfy-sh-migration = handleTest ./ntfy-sh-migration.nix {};
   ntpd-rs = handleTest ./ntpd-rs.nix {};
+  nvidia-container-toolkit = handleTest ./nvidia-container-toolkit.nix {};
   nvmetcfg = handleTest ./nvmetcfg.nix {};
   nzbget = handleTest ./nzbget.nix {};
   nzbhydra2 = handleTest ./nzbhydra2.nix {};
diff --git a/nixos/tests/nvidia-container-toolkit.nix b/nixos/tests/nvidia-container-toolkit.nix
new file mode 100644
index 0000000000000..c6c96732d5963
--- /dev/null
+++ b/nixos/tests/nvidia-container-toolkit.nix
@@ -0,0 +1,158 @@
+import ./make-test-python.nix (
+  {
+    pkgs,
+    lib,
+    system,
+    ...
+  }:
+  let
+    unfreeAndInsecurePkgs = import ../.. {
+      inherit system;
+      config = {
+        allowUnfree = true;
+        permittedInsecurePackages = [ "openssl-1.1.1w" ];
+      };
+    };
+    testContainerImage =
+      let
+        testCDIScript = pkgs.writeShellScriptBin "test-cdi" ''
+          die() {
+            echo "$1"
+            exit 1
+          }
+
+          check_file_referential_integrity() {
+            echo "checking $file referential integrity"
+            files=$(set -o pipefail && \
+                    ${pkgs.glibc.bin}/bin/ldd "$1" | \
+                    ${pkgs.gnugrep}/bin/grep '=>' | \
+                    ${pkgs.gnused}/bin/sed "s/.* => //" | \
+                    ${pkgs.gnused}/bin/sed "s/ (.*//") || exit 1
+
+            for file in $files; do
+              if [ ! -f "$file" ]; then
+                die "$file does not exist in the container filesystem"
+              fi
+            done
+          }
+
+          check_directory_referential_integrity() {
+            ${pkgs.findutils}/bin/find "$1" -type f -print0 | while read -d $'\0' file; do
+              if [[ $(${pkgs.file}/bin/file "$file" | ${pkgs.gnugrep}/bin/grep ELF) ]]; then
+                check_file_referential_integrity "$file" || exit 1
+              else
+                echo "skipping $file"
+              fi
+            done
+          }
+
+          check_directory_referential_integrity "/usr/bin" || exit 1
+          check_directory_referential_integrity "${pkgs.addDriverRunpath.driverLink}" || exit 1
+          check_directory_referential_integrity "/usr/local/nvidia" || exit 1
+        '';
+      in
+      pkgs.dockerTools.buildImage {
+        name = "cdi-test";
+        tag = "latest";
+        config = {
+          Cmd = [ "${testCDIScript}/bin/test-cdi" ];
+          Env = [
+            "LD_LIBRARY_PATH=${unfreeAndInsecurePkgs.linuxPackages.nvidia_x11}/lib:${lib.getLib unfreeAndInsecurePkgs.openssl_1_1}/lib"
+          ];
+        };
+        copyToRoot = (
+          with pkgs.dockerTools;
+          [
+            usrBinEnv
+            binSh
+          ]
+        );
+      };
+  in
+  {
+    name = "nvidia-container-toolkit";
+    meta = with lib.maintainers; {
+      maintainers = [ ereslibre ];
+    };
+    nodes = {
+      no-nvidia-gpus = {
+        environment.systemPackages = with pkgs; [ jq ];
+        hardware.nvidia-container-toolkit.enable = true;
+        nixpkgs.config.allowUnfree = true;
+      };
+      nvidia-one-gpu =
+        { pkgs, ... }:
+        let
+          emptyCDISpec = ''
+            #! ${pkgs.runtimeShell}
+            cat <<CDI_DOCUMENT
+              {
+                "cdiVersion": "0.5.0",
+                "kind": "nvidia.com/gpu",
+                "devices": [
+                  {
+                    "name": "all",
+                    "containerEdits": {
+                      "deviceNodes": [
+                        {
+                          "path": "/dev/urandom"
+                        }
+                      ],
+                      "hooks": [],
+                      "mounts": []
+                    }
+                  }
+                ],
+                "containerEdits": {
+                  "deviceNodes": [],
+                  "hooks": [],
+                  "mounts": []
+                }
+              }
+            CDI_DOCUMENT
+          '';
+        in
+        {
+          virtualisation.diskSize = 10240;
+          environment.systemPackages = with pkgs; [
+            jq
+            podman
+          ];
+          hardware = {
+            nvidia-container-toolkit = {
+              enable = true;
+              package = pkgs.stdenv.mkDerivation {
+                name = "nvidia-ctk-dummy";
+                version = "1.0.0";
+                phases = "installPhase";
+                installPhase = ''
+                  mkdir -p $out/bin
+                  cat <<EOF > $out/bin/nvidia-ctk
+                    ${emptyCDISpec}
+                  EOF
+                  chmod +x $out/bin/nvidia-ctk
+                '';
+              };
+            };
+            opengl.enable = true;
+          };
+          nixpkgs.config.allowUnfree = true;
+          services.xserver.videoDrivers = [ "nvidia" ];
+          virtualisation.containers.enable = true;
+        };
+    };
+    testScript = ''
+      start_all()
+
+      with subtest("Generate the CDI spec (empty) for a machine with no Nvidia GPU"):
+        no_nvidia_gpus.wait_for_unit("nvidia-container-toolkit-cdi-generator.service")
+        no_nvidia_gpus.succeed("cat /var/run/cdi/nvidia-container-toolkit.json | jq")
+
+      with subtest("Generate the CDI spec for a machine with an Nvidia GPU"):
+        nvidia_one_gpu.wait_for_unit("nvidia-container-toolkit-cdi-generator.service")
+        nvidia_one_gpu.succeed("cat /var/run/cdi/nvidia-container-toolkit.json | jq")
+        nvidia_one_gpu.succeed("podman load < ${testContainerImage}")
+        print(nvidia_one_gpu.succeed("podman run --pull=never --device=nvidia.com/gpu=all -v /run/opengl-driver:/run/opengl-driver:ro cdi-test:latest"))
+    '';
+  }
+)