diff --git a/README.md b/README.md
index 27cd3f1..192129e 100644
--- a/README.md
+++ b/README.md
@@ -36,16 +36,16 @@ This is how your iso configuation may look like
 
 /etc/nixos/configuration.nix
 ```nix
-{ pkgs, ... }:
+{ pkgs, modulesPath, ... }:
 let
-  disko = (builtins.fetchGit {
+  disko = builtins.fetchGit {
     url = https://cgit.lassul.us/disko/;
     rev = "88f56a0b644dd7bfa8438409bea5377adef6aef4";
-  }) + "/lib";
+  };
   cfg = builtins.fromJSON ./tsp-disk.json;
 in {
   imports = [
-    <nixpkgs/nixos/modules/installer/cd-dvd/installation-cd-minimal.nix>
+    (modulesPath + "/installer/cd-dvd/installation-cd-minimal.nix")
   ];
   environment.systemPackages = with pkgs;[
     (pkgs.writeScriptBin "tsp-create" (disko.create cfg))
diff --git a/default.nix b/default.nix
index 6d08b8d..6858c40 100644
--- a/default.nix
+++ b/default.nix
@@ -1,3 +1,5 @@
-{
-  inherit (import ./lib) config create mount;
+{ lib ? (import <nixpkgs> {}).lib }: {
+  inherit (import ./lib {
+    inherit lib;
+  }) config create mount;
 }
diff --git a/example/default.nix b/example/default.nix
index aa03d27..f2826c9 100644
--- a/example/default.nix
+++ b/example/default.nix
@@ -1,6 +1,10 @@
 # usage: nix-instantiate --eval --json --strict example | jq -r .
 
-with import ../lib;
+let
+  # TODO: get rid of NIX_PATH dependency here
+  pkgs = import <nixpkgs> {};
+in
+with import ../lib { inherit (pkgs) lib;};
 
 {
   config = config (import ./config.nix);
diff --git a/lib/default.nix b/lib/default.nix
index 7793acf..beba1b0 100644
--- a/lib/default.nix
+++ b/lib/default.nix
@@ -1,4 +1,5 @@
-with import <nixpkgs/lib>;
+{ lib }:
+with lib;
 with builtins;
 
 let {
@@ -130,5 +131,4 @@ let {
 
   mount.table = q: x:
     foldl' recursiveUpdate {} (imap (index: mount-f (q // { inherit index; })) x.partitions);
-
 }
diff --git a/tests/test.nix b/tests/test.nix
index 76bc5ea..fa89208 100644
--- a/tests/test.nix
+++ b/tests/test.nix
@@ -1,5 +1,12 @@
-import <nixpkgs/nixos/tests/make-test-python.nix> ({ pkgs, ... }: let
-
+{ makeTest ? import <nixpkgs/nixos/tests/make-test-python.nix>
+, pkgs ? (import <nixpkgs> {})
+}:
+let
+  makeTest' = args:
+    makeTest args {
+      inherit pkgs;
+      inherit (pkgs) system;
+    };
   disko-config = {
     type = "devices";
     content = {
@@ -78,28 +85,28 @@ import <nixpkgs/nixos/tests/make-test-python.nix> ({ pkgs, ... }: let
       };
     };
   };
-
-in {
+in makeTest' {
   name = "disko";
 
   nodes.machine =
-    { config, pkgs, ... }:
+    { config, pkgs, modulesPath, ... }:
 
     {
       imports = [
-        <nixpkgs/nixos/modules/profiles/installation-device.nix>
-        <nixpkgs/nixos/modules/profiles/base.nix>
+        (modulesPath + "/profiles/installation-device.nix")
+        (modulesPath + "/profiles/base.nix")
       ];
 
+      # speed-up eval
+      documentation.enable = false;
+
       virtualisation.emptyDiskImages = [ 512 ];
     };
 
-  testScript =
-    ''
-      machine.succeed("echo 'secret' > /tmp/secret.key");
-      machine.succeed("${pkgs.writeScript "create" ((import ../lib).create disko-config)}");
-      machine.succeed("${pkgs.writeScript "mount" ((import ../lib).mount disko-config)}");
-      machine.succeed("test -b /dev/mapper/pool-raw");
-    '';
-
-})
+  testScript = ''
+    machine.succeed("echo 'secret' > /tmp/secret.key");
+    machine.succeed("${pkgs.writeScript "create" ((pkgs.callPackage ../. {}).create disko-config)}");
+    machine.succeed("${pkgs.writeScript "mount" ((pkgs.callPackage ../. {}).mount disko-config)}");
+    machine.succeed("test -b /dev/mapper/pool-raw");
+  '';
+}