diff --git a/tests/lib.nix b/tests/lib.nix
index 3fb3486..6d787ed 100644
--- a/tests/lib.nix
+++ b/tests/lib.nix
@@ -103,7 +103,7 @@
         virtualisation.emptyDiskImages = builtins.genList (_: 4096) num-disks;
       };
 
-      testScript = ''
+      testScript = { nodes, ... }: ''
         def disks(oldmachine, num_disks):
             disk_flags = ""
             for i in range(num_disks):
@@ -126,9 +126,9 @@
           machine.succeed("${tsp-mount}") # verify that the command is idempotent
         ''}
         ${lib.optionalString (testMode == "module") ''
-          machine.succeed("disko-create")
-          machine.succeed("disko-mount")
-          machine.succeed("disko-mount") # verify that the command is idempotent
+          machine.succeed("${nodes.machine.system.build.formatScript}")
+          machine.succeed("${nodes.machine.system.build.mountScript}")
+          machine.succeed("${nodes.machine.system.build.mountScript}") # verify that the command is idempotent
         ''}
         ${lib.optionalString (testMode == "cli") ''
           # TODO use the disko cli here