diff --git a/flake.nix b/flake.nix
index 6f5b7af..94ad937 100644
--- a/flake.nix
+++ b/flake.nix
@@ -14,13 +14,17 @@
 
   outputs = { nixpkgs, nix-flake-tests, flake-utils, nmdsrc, ... }: flake-utils.lib.eachDefaultSystem (system:
     let
-      patches = [
-      ];
       originPkgs = nixpkgs.legacyPackages.${system};
+      patches = [
+        (originPkgs.fetchpatch {
+          url = "https://patch-diff.githubusercontent.com/raw/NixOS/nixpkgs/pull/315018.patch";
+          hash = "sha256-8jcGyO/d+htfv/ZajxXh89S3OiDZAr7/fsWC1JpGczM=";
+        })
+      ];
       patchedNixpkgs = originPkgs.applyPatches {
         name = "nixpkgs-patched";
         src = nixpkgs;
-        patches = map (p: originPkgs.writeText "patch" p) patches;
+        inherit patches;
       };
 
       pkgs = import patchedNixpkgs {
@@ -85,12 +89,14 @@
                 inherit (pkgs) lib;
               }
             );
+
+            shblib = pkgs.callPackage ./lib {};
           in (rec {
             all = mergeTests [
               modules
             ];
 
-            modules = nix-flake-tests.lib.check {
+            modules = shblib.check {
               inherit pkgs;
               tests =
                 mergeTests (importFiles [
diff --git a/lib/default.nix b/lib/default.nix
index 19831ca..de48d55 100644
--- a/lib/default.nix
+++ b/lib/default.nix
@@ -204,4 +204,47 @@ rec {
     (lib.attrsets.filterAttrs (name: v: name == from) attrset) // {
       ${to} = attrset.${from};
     };
+
+  # Taken from https://github.com/antifuchs/nix-flake-tests/blob/main/default.nix
+  # with a nicer diff display function.
+  check = { pkgs, tests }:
+    let
+      system = pkgs.stdenv.targetPlatform.system;
+      formatValue = val:
+        if (builtins.isList val || builtins.isAttrs val) then builtins.toJSON val
+        else builtins.toString val;
+      resultToString = { name, expected, result }:
+        pkgs.callPackage (pkgs.runCommand "nix-flake-tests-error" {
+          expected = formatValue expected;
+          result = formatValue result;
+          passAsFile = [ "expected" "result" ];
+        } ''
+          echo "${name} failed (- expected, + result)"
+          cp ''${expectedPath} ''${expectedPath}.json
+          cp ''${resultPath} ''${resultPath}.json
+          ${pkgs.deepdiff}/bin/deep diff ''${expectedPath}.json ''${resultPath}.json
+        '') {};
+
+
+
+      #   ''
+      #   ${name} failed: expected ${formatValue expected}, but got ${
+      #     formatValue result
+      #   }
+      # '';
+      results = pkgs.lib.runTests tests;
+    in
+    if results != [ ] then
+      builtins.throw (builtins.concatStringsSep "\n" (map resultToString results))
+    ## TODO: The derivation below is preferable but "nix flake check" hangs with it:
+    ##       (it's preferable because "examples/many-failures" would then show all errors.)
+    # pkgs.runCommand "nix-flake-tests-failure" { } ''
+    #   cat <<EOF
+    #   ${builtins.concatStringsSep "\n" (map resultToString results)}
+    #   EOF
+    #   exit 1
+    # ''
+    else
+      pkgs.runCommand "nix-flake-tests-success" { } "echo > $out";
+
 }