From 4c88deef38678154f3043ca2e0f4bd84a3c6cbc5 Mon Sep 17 00:00:00 2001
From: Parker Hoyes <contact@parkerhoyes.com>
Date: Tue, 3 Sep 2024 17:27:56 +0000
Subject: [PATCH] Add tests for daemon with cgroups

---
 tests/nixos/cgroups/default.nix | 40 +++++++++++++++++++++++++++++++++
 tests/nixos/cgroups/hang.nix    | 10 +++++++++
 tests/nixos/default.nix         |  2 ++
 3 files changed, 52 insertions(+)
 create mode 100644 tests/nixos/cgroups/default.nix
 create mode 100644 tests/nixos/cgroups/hang.nix

diff --git a/tests/nixos/cgroups/default.nix b/tests/nixos/cgroups/default.nix
new file mode 100644
index 000000000..b8febbf4b
--- /dev/null
+++ b/tests/nixos/cgroups/default.nix
@@ -0,0 +1,40 @@
+{ nixpkgs, ... }:
+
+{
+  name = "cgroups";
+
+  nodes =
+    {
+      host =
+        { config, pkgs, ... }:
+        { virtualisation.additionalPaths = [ pkgs.stdenvNoCC ];
+          nix.extraOptions =
+            ''
+              extra-experimental-features = nix-command auto-allocate-uids cgroups
+              extra-system-features = uid-range
+            '';
+          nix.settings.use-cgroups = true;
+          nix.nixPath = [ "nixpkgs=${nixpkgs}" ];
+        };
+    };
+
+  testScript = { nodes }: ''
+    start_all()
+
+    host.wait_for_unit("multi-user.target")
+
+    # Start build in background
+    host.execute("NIX_REMOTE=daemon nix build --auto-allocate-uids --file ${./hang.nix} >&2 &")
+    service = "/sys/fs/cgroup/system.slice/nix-daemon.service"
+
+    # Wait for cgroups to be created
+    host.succeed(f"until [ -e {service}/nix-daemon ]; do sleep 1; done", timeout=30)
+    host.succeed(f"until [ -e {service}/nix-build-uid-* ]; do sleep 1; done", timeout=30)
+
+    # Check that there aren't processes where there shouldn't be, and that there are where there should be
+    host.succeed(f'[ -z "$(cat {service}/cgroup.procs)" ]')
+    host.succeed(f'[ -n "$(cat {service}/nix-daemon/cgroup.procs)" ]')
+    host.succeed(f'[ -n "$(cat {service}/nix-build-uid-*/cgroup.procs)" ]')
+  '';
+
+}
diff --git a/tests/nixos/cgroups/hang.nix b/tests/nixos/cgroups/hang.nix
new file mode 100644
index 000000000..cefe2d031
--- /dev/null
+++ b/tests/nixos/cgroups/hang.nix
@@ -0,0 +1,10 @@
+{ }:
+
+with import <nixpkgs> {};
+
+runCommand "hang"
+  { requiredSystemFeatures = "uid-range";
+  }
+  ''
+    sleep infinity
+  ''
diff --git a/tests/nixos/default.nix b/tests/nixos/default.nix
index 40d29b371..62fc6b10f 100644
--- a/tests/nixos/default.nix
+++ b/tests/nixos/default.nix
@@ -157,4 +157,6 @@ in
   s3-binary-cache-store = runNixOSTestFor "x86_64-linux" ./s3-binary-cache-store.nix;
 
   fsync = runNixOSTestFor "x86_64-linux" ./fsync.nix;
+
+  cgroups = runNixOSTestFor "x86_64-linux" ./cgroups;
 }