diff --git a/example/zfs.nix b/example/zfs.nix
index aebb6a3..e126d51 100644
--- a/example/zfs.nix
+++ b/example/zfs.nix
@@ -64,6 +64,7 @@
           "com.sun:auto-snapshot" = "false";
         };
         mountpoint = "/";
+        postCreateHook = "zfs snapshot zroot@blank";
 
         datasets = {
           zfs_fs = {
@@ -98,6 +99,9 @@
               keyformat = "passphrase";
               keylocation = "file:///tmp/secret.key";
             };
+            postCreateHook = ''
+              zfs set keylocation="prompt" "zroot/$name";
+            '';
           };
           "encrypted/test" = {
             zfs_type = "filesystem";
diff --git a/tests/zfs.nix b/tests/zfs.nix
index 75268ef..39c2ebe 100644
--- a/tests/zfs.nix
+++ b/tests/zfs.nix
@@ -8,9 +8,6 @@ makeDiskoTest {
     fileSystems."/zfs_legacy_fs".options = [ "nofail" ]; # TODO find out why we need this!
     boot.zfs.requestEncryptionCredentials = true;
   };
-  postDisko = ''
-    machine.succeed("zfs set keylocation=prompt zroot/encrypted")
-  '';
   enableOCR = true;
   bootCommands = ''
     machine.wait_for_text("passphrase for")
@@ -32,6 +29,9 @@ makeDiskoTest {
     assert_property("zroot/zfs_testvolume", "volsize", "10M")
     assert_property("zroot/zfs_unmounted_fs", "mountpoint", "none")
 
+    assert_property("zroot/encrypted", "keylocation", "prompt")
+    machine.succeed("zfs get name zroot@blank")
+
     machine.succeed("mountpoint /zfs_fs");
     machine.succeed("mountpoint /zfs_legacy_fs");
     machine.succeed("mountpoint /ext4onzfs");