diff --git a/pkgs/applications/version-management/nbstripout/default.nix b/pkgs/applications/version-management/nbstripout/default.nix
index 77540deb909f..a201ce6ddbb4 100644
--- a/pkgs/applications/version-management/nbstripout/default.nix
+++ b/pkgs/applications/version-management/nbstripout/default.nix
@@ -1,20 +1,17 @@
-{lib, python, git, mercurial, coreutils}:
+{ lib
+, python3
+, coreutils
+, git
+, mercurial
+}:
 
-with python.pkgs;
-buildPythonApplication rec {
-  version = "0.6.0";
+python3.pkgs.buildPythonApplication rec {
+  version = "0.6.1";
   pname = "nbstripout";
 
-  # Mercurial should be added as a build input but because it's a Python
-  # application, it would mess up the Python environment. Thus, don't add it
-  # here, instead add it to PATH when running unit tests
-  checkInputs = [ pytest pytest-flake8 git ];
-  nativeBuildInputs = [ pytest-runner ];
-  propagatedBuildInputs = [ ipython nbformat ];
-
-  src = fetchPypi {
+  src = python3.pkgs.fetchPypi {
     inherit pname version;
-    sha256 = "sha256-TWxDAhVqskaMyOcgLvKPNN2RhFFOIeRDQLzShpaMgss=";
+    hash = "sha256-kGW83RSIs4bk88CB/8HUj0UTovjYv00NmiggjF2v6dM=";
   };
 
   # for some reason, darwin uses /bin/sh echo native instead of echo binary, so
@@ -23,9 +20,24 @@ buildPythonApplication rec {
     substituteInPlace tests/test-git.t --replace "echo" "${coreutils}/bin/echo"
   '';
 
-  # ignore flake8 tests for the nix wrapped setup.py
-  checkPhase = ''
-    PATH=$PATH:$out/bin:${mercurial}/bin pytest .
+  propagatedBuildInputs = with python3.pkgs; [
+    ipython
+    nbformat
+  ];
+
+  checkInputs = [
+    coreutils
+    git
+    mercurial
+  ] ++ (with python3.pkgs; [
+    pytest-cram
+    pytestCheckHook
+  ]);
+
+  preCheck = ''
+    export HOME=$(mktemp -d)
+    export PATH=$out/bin:$PATH
+    git config --global init.defaultBranch main
   '';
 
   meta = {
diff --git a/pkgs/top-level/all-packages.nix b/pkgs/top-level/all-packages.nix
index 0f5d426c71d5..3a13a4a8fa98 100644
--- a/pkgs/top-level/all-packages.nix
+++ b/pkgs/top-level/all-packages.nix
@@ -30642,7 +30642,7 @@ with pkgs;
     inherit (darwin) libiconv;
   };
 
-  nbstripout = callPackage ../applications/version-management/nbstripout { python = python3; };
+  nbstripout = callPackage ../applications/version-management/nbstripout { };
 
   ncmpc = callPackage ../applications/audio/ncmpc { };