diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..b2be92b
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+result
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/flake.lock b/flake.lock
new file mode 100644
index 0000000..d99ebe5
--- /dev/null
+++ b/flake.lock
@@ -0,0 +1,27 @@
+{
+  "nodes": {
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1660646295,
+        "narHash": "sha256-V4G+egGRc3elXPTr7QLJ7r7yrYed0areIKDiIAlMLC8=",
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "762b003329510ea855b4097a37511eb19c7077f0",
+        "type": "github"
+      },
+      "original": {
+        "owner": "NixOS",
+        "ref": "nixos-unstable",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "nixpkgs": "nixpkgs"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}
diff --git a/flake.nix b/flake.nix
new file mode 100644
index 0000000..3e7263a
--- /dev/null
+++ b/flake.nix
@@ -0,0 +1,19 @@
+{
+  description = "Description for the project";
+
+  inputs.nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
+
+  outputs = { self, nixpkgs, ... }: {
+    lib = import ./. {
+      inherit (nixpkgs) lib;
+    };
+    checks.x86_64-linux = let
+      pkgs = nixpkgs.legacyPackages.x86_64-linux;
+    in {
+      # Run tests: nix flake check -L
+      nixos-test = pkgs.callPackage ./tests/test.nix {
+        makeTest = import (pkgs.path + "/nixos/tests/make-test-python.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");
+  '';
+}