From 30421d96f0c8326a0bdd46cd77599f82121680ee Mon Sep 17 00:00:00 2001
From: Simon Gardling <titaniumtown@proton.me>
Date: Mon, 11 Aug 2025 16:18:10 -0700
Subject: [PATCH] move ensureZfsMounts

---
 flake.nix           |  5 +++-
 lib.nix             | 63 ++++++++++++---------------------------------
 overlays.nix        | 33 ++++++++++++++++++++++++
 tests/minecraft.nix |  5 +++-
 tests/zfs.nix       | 12 +++++++--
 5 files changed, 68 insertions(+), 50 deletions(-)
 create mode 100644 overlays.nix

diff --git a/flake.nix b/flake.nix
index 9eda8d5..6b4496a 100644
--- a/flake.nix
+++ b/flake.nix
@@ -189,7 +189,10 @@
           # get nix-minecraft working!
           nix-minecraft.nixosModules.minecraft-servers
           {
-            nixpkgs.overlays = [ nix-minecraft.overlay ];
+            nixpkgs.overlays = [
+              nix-minecraft.overlay
+              (import ./overlays.nix)
+            ];
           }
 
           lanzaboote.nixosModules.lanzaboote
diff --git a/lib.nix b/lib.nix
index 1437931..ef658f0 100644
--- a/lib.nix
+++ b/lib.nix
@@ -9,58 +9,29 @@ inputs.nixpkgs.lib.extend (
     lib = prev;
   in
   {
-    # Should probably be moved to a package later instead of this
-    ensureZfsMounts = pkgs.writeShellApplication {
-      name = "zfsEnsureMounted";
-      runtimeInputs = with pkgs; [
-        zfs
-        gnugrep
-        gawk
-        coreutils
-      ];
 
-      text = ''
-        #!/bin/sh -x
+    serviceMountDeps =
+      serviceName: dirs:
+      { pkgs, ... }:
+      {
+        systemd.services."${serviceName}_mounts" = {
+          wants = [ "zfs.target" ];
+          before = [ "${serviceName}.service" ];
 
-        if [[ "$#" -eq "0" ]]; then
-           echo "no arguments passed"
-           exit 1
-        fi
+          serviceConfig = {
+            Type = "oneshot";
+            RemainAfterExit = true;
+            ExecStart = "${lib.getExe pkgs.ensureZfsMounts} ${lib.strings.concatStringsSep " " dirs}";
+          };
+        };
 
-        TARGETS=$(echo "$@" | sort | uniq)
-        MOUNTED=$(zfs list -o mountpoint,mounted -H | awk '$2 == "yes" {print $1}' | sort | uniq)
-        NUM_MATCHED=$(echo "$MOUNTED" | grep -Ec "$(echo "$@" | tr ' ' '\|')") # does not properly handle paths with strings
-        if [[ "$NUM_MATCHED" -eq "$#" ]]; then
-           exit 0
-        fi
-
-        FOUND=$(printf "%s\n%s" "$TARGETS" "$MOUNTED" | sort | uniq -c | awk '$1 == "2" {print $2}' | sort)
-        MISSING=$(printf "%s\n%s" "$FOUND" "$TARGETS" | sort | uniq -u | sort)
-
-        echo "FAILURE, missing: $MISSING" 1>&2
-        exit 1
-      '';
-    };
-
-    serviceMountDeps = serviceName: dirs: {
-      systemd.services."${serviceName}_mounts" = {
-        wants = [ "zfs.target" ];
-        before = [ "${serviceName}.service" ];
-
-        serviceConfig = {
-          Type = "oneshot";
-          RemainAfterExit = true;
-          ExecStart = "${lib.getExe final.ensureZfsMounts} ${lib.strings.concatStringsSep " " dirs}";
+        systemd.services.${serviceName} = {
+          wants = [ "${serviceName}_mounts.service" ];
+          after = [ "${serviceName}_mounts.service" ];
+          requires = [ "${serviceName}_mounts.service" ];
         };
       };
 
-      systemd.services.${serviceName} = {
-        wants = [ "${serviceName}_mounts.service" ];
-        after = [ "${serviceName}_mounts.service" ];
-        requires = [ "${serviceName}_mounts.service" ];
-      };
-    };
-
     # stolen from: https://stackoverflow.com/a/42398526
     optimizeWithFlags =
       pkg: flags:
diff --git a/overlays.nix b/overlays.nix
new file mode 100644
index 0000000..e64025b
--- /dev/null
+++ b/overlays.nix
@@ -0,0 +1,33 @@
+final: prev: {
+  ensureZfsMounts = prev.writeShellApplication {
+    name = "zfsEnsureMounted";
+    runtimeInputs = with prev; [
+      zfs
+      gnugrep
+      gawk
+      coreutils
+    ];
+
+    text = ''
+      #!/bin/sh -x
+
+      if [[ "$#" -eq "0" ]]; then
+         echo "no arguments passed"
+         exit 1
+      fi
+
+      TARGETS=$(echo "$@" | sort | uniq)
+      MOUNTED=$(zfs list -o mountpoint,mounted -H | awk '$2 == "yes" {print $1}' | sort | uniq)
+      NUM_MATCHED=$(echo "$MOUNTED" | grep -Ec "$(echo "$@" | tr ' ' '\|')") # does not properly handle paths with strings
+      if [[ "$NUM_MATCHED" -eq "$#" ]]; then
+         exit 0
+      fi
+
+      FOUND=$(printf "%s\n%s" "$TARGETS" "$MOUNTED" | sort | uniq -c | awk '$1 == "2" {print $2}' | sort)
+      MISSING=$(printf "%s\n%s" "$FOUND" "$TARGETS" | sort | uniq -u | sort)
+
+      echo "FAILURE, missing: $MISSING" 1>&2
+      exit 1
+    '';
+  };
+}
diff --git a/tests/minecraft.nix b/tests/minecraft.nix
index 9b9c9df..a304158 100644
--- a/tests/minecraft.nix
+++ b/tests/minecraft.nix
@@ -10,7 +10,10 @@ let
   testPkgs = import inputs.nixpkgs {
     system = pkgs.system;
     config.allowUnfreePredicate = pkg: builtins.elem (lib.getName pkg) [ "minecraft-server" ];
-    overlays = [ inputs.nix-minecraft.overlay ];
+    overlays = [
+      inputs.nix-minecraft.overlay
+      (import ../overlays.nix)
+    ];
   };
 
   # Create a wrapper module that imports the actual minecraft service
diff --git a/tests/zfs.nix b/tests/zfs.nix
index 6bfabfc..ca9d92e 100644
--- a/tests/zfs.nix
+++ b/tests/zfs.nix
@@ -2,9 +2,17 @@
   config,
   lib,
   pkgs,
+  inputs,
   ...
 }:
-pkgs.testers.runNixOSTest {
+let
+  # Create pkgs with ensureZfsMounts overlay
+  testPkgs = import inputs.nixpkgs {
+    system = pkgs.system;
+    overlays = [ (import ../overlays.nix) ];
+  };
+in
+testPkgs.testers.runNixOSTest {
   name = "zfs folder dependency and mounting test";
 
   nodes.machine =
@@ -35,7 +43,7 @@ pkgs.testers.runNixOSTest {
 
       environment.systemPackages = [
         pkgs.parted
-        lib.ensureZfsMounts
+        pkgs.ensureZfsMounts
       ];
 
       systemd.services.foobar = {