diff --git a/pkgs/development/python-modules/pyinfra-testgen/default.nix b/pkgs/development/python-modules/pyinfra-testgen/default.nix
new file mode 100644
index 000000000000..9a426e492ca6
--- /dev/null
+++ b/pkgs/development/python-modules/pyinfra-testgen/default.nix
@@ -0,0 +1,43 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchPypi,
+  uv-build,
+  pyyaml,
+}:
+
+buildPythonPackage rec {
+  pname = "pyinfra-testgen";
+  version = "0.1.1";
+  pyproject = true;
+
+  # no tags on GitHub
+  src = fetchPypi {
+    pname = "pyinfra_testgen";
+    inherit version;
+    hash = "sha256-c5pZ0SfRXC50vJZfnnf0HQgImf7hi2oQ5/XKMVNzlpc=";
+  };
+
+  postPatch = ''
+    substituteInPlace pyproject.toml \
+      --replace-fail "uv_build>=0.8.14,<0.9.0" uv_build
+  '';
+
+  build-system = [ uv-build ];
+
+  dependencies = [
+    pyyaml
+  ];
+
+  pythonImportsCheck = [ "testgen" ];
+
+  # upstream has no tests
+  doCheck = false;
+
+  meta = {
+    description = "Generate Python unit tests from JSON and YAML files";
+    homepage = "https://github.com/pyinfra-dev/testgen";
+    license = lib.licenses.mit;
+    maintainers = [ lib.maintainers.dotlambda ];
+  };
+}
diff --git a/pkgs/development/python-modules/pyinfra/default.nix b/pkgs/development/python-modules/pyinfra/default.nix
index a73a8d10b9b5..6eda1baa56b8 100644
--- a/pkgs/development/python-modules/pyinfra/default.nix
+++ b/pkgs/development/python-modules/pyinfra/default.nix
@@ -4,33 +4,54 @@
   click,
   distro,
   fetchFromGitHub,
+  fetchpatch,
+  freezegun,
   gevent,
-  importlib-metadata,
+  hatchling,
   jinja2,
   packaging,
   paramiko,
+  pydantic,
+  pyinfra-testgen,
+  pytest-testinfra,
   pytestCheckHook,
   python-dateutil,
   pythonOlder,
-  pywinrm,
-  setuptools,
   typeguard,
   typing-extensions,
+  uv-dynamic-versioning,
 }:
 
 buildPythonPackage rec {
   pname = "pyinfra";
-  version = "3.4.1";
+  version = "3.6";
   pyproject = true;
 
   src = fetchFromGitHub {
     owner = "Fizzadar";
     repo = "pyinfra";
     tag = "v${version}";
-    hash = "sha256-7bNkDm5SyIgVkrGQ95/q7AiY/JnxtWx+jkDO/rJQ2WQ=";
+    hash = "sha256-CTeGn9aN5voyCUL5LuTErLgTgC1Z/qTS7SB9TNfq7mc=";
   };
 
-  build-system = [ setuptools ];
+  patches = [
+    # paramiko v4 compat
+    # https://github.com/pyinfra-dev/pyinfra/pull/1525
+    (fetchpatch {
+      name = "remove-DSSKey.patch";
+      url = "https://github.com/pyinfra-dev/pyinfra/commit/a655bdf425884055145cfd0011c3b444c9a3ada2.patch";
+      hash = "sha256-puHcA4+KigltCL2tUYRMc9OT3kxvTeW77bbFbxgkcTs=";
+    })
+  ];
+
+  build-system = [
+    hatchling
+    uv-dynamic-versioning
+  ];
+
+  pythonRelaxDeps = [
+    "paramiko"
+  ];
 
   dependencies = [
     click
@@ -39,15 +60,18 @@ buildPythonPackage rec {
     jinja2
     packaging
     paramiko
+    pydantic
     python-dateutil
-    pywinrm
-    setuptools
     typeguard
   ]
-  ++ lib.optionals (pythonOlder "3.11") [ typing-extensions ]
-  ++ lib.optionals (pythonOlder "3.10") [ importlib-metadata ];
+  ++ lib.optionals (pythonOlder "3.11") [ typing-extensions ];
 
-  nativeCheckInputs = [ pytestCheckHook ];
+  nativeCheckInputs = [
+    freezegun
+    pyinfra-testgen
+    pytest-testinfra
+    pytestCheckHook
+  ];
 
   pythonImportsCheck = [ "pyinfra" ];
 
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 85382c4d8fad..77401c0a50f6 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -13645,6 +13645,8 @@ self: super: with self; {
 
   pyinfra = callPackage ../development/python-modules/pyinfra { };
 
+  pyinfra-testgen = callPackage ../development/python-modules/pyinfra-testgen { };
+
   pyinotify = callPackage ../development/python-modules/pyinotify { };
 
   pyinputevent = callPackage ../development/python-modules/pyinputevent { };