diff --git a/nixos/lib/test-driver/default.nix b/nixos/lib/test-driver/default.nix
index 35471f74949c..f22744806d48 100644
--- a/nixos/lib/test-driver/default.nix
+++ b/nixos/lib/test-driver/default.nix
@@ -14,22 +14,12 @@
   extraPythonPackages ? (_: [ ]),
   nixosTests,
 }:
-let
-  fs = lib.fileset;
-in
 python3Packages.buildPythonApplication {
   pname = "nixos-test-driver";
   version = "1.1";
   pyproject = true;
 
-  src = fs.toSource {
-    root = ./.;
-    fileset = fs.unions [
-      ./pyproject.toml
-      ./test_driver
-      ./extract-docstrings.py
-    ];
-  };
+  src = ./src;
 
   build-system = with python3Packages; [
     setuptools
diff --git a/nixos/lib/test-driver/nixos-test-driver-docstrings.nix b/nixos/lib/test-driver/nixos-test-driver-docstrings.nix
index b76da9ac5fb2..b6450685e047 100644
--- a/nixos/lib/test-driver/nixos-test-driver-docstrings.nix
+++ b/nixos/lib/test-driver/nixos-test-driver-docstrings.nix
@@ -11,6 +11,6 @@ in
 
 runCommand "nixos-test-driver-docstrings" env ''
   mkdir $out
-  python3 ${./extract-docstrings.py} ${./test_driver/machine.py} \
+  python3 ${./src/extract-docstrings.py} ${./src/test_driver/machine.py} \
     > $out/machine-methods.md
 ''
diff --git a/nixos/lib/test-driver/extract-docstrings.py b/nixos/lib/test-driver/src/extract-docstrings.py
similarity index 100%
rename from nixos/lib/test-driver/extract-docstrings.py
rename to nixos/lib/test-driver/src/extract-docstrings.py
diff --git a/nixos/lib/test-driver/pyproject.toml b/nixos/lib/test-driver/src/pyproject.toml
similarity index 100%
rename from nixos/lib/test-driver/pyproject.toml
rename to nixos/lib/test-driver/src/pyproject.toml
diff --git a/nixos/lib/test-driver/test_driver/__init__.py b/nixos/lib/test-driver/src/test_driver/__init__.py
similarity index 100%
rename from nixos/lib/test-driver/test_driver/__init__.py
rename to nixos/lib/test-driver/src/test_driver/__init__.py
diff --git a/nixos/lib/test-driver/test_driver/driver.py b/nixos/lib/test-driver/src/test_driver/driver.py
similarity index 100%
rename from nixos/lib/test-driver/test_driver/driver.py
rename to nixos/lib/test-driver/src/test_driver/driver.py
diff --git a/nixos/lib/test-driver/test_driver/logger.py b/nixos/lib/test-driver/src/test_driver/logger.py
similarity index 100%
rename from nixos/lib/test-driver/test_driver/logger.py
rename to nixos/lib/test-driver/src/test_driver/logger.py
diff --git a/nixos/lib/test-driver/test_driver/machine.py b/nixos/lib/test-driver/src/test_driver/machine.py
similarity index 100%
rename from nixos/lib/test-driver/test_driver/machine.py
rename to nixos/lib/test-driver/src/test_driver/machine.py
diff --git a/nixos/lib/test-driver/test_driver/polling_condition.py b/nixos/lib/test-driver/src/test_driver/polling_condition.py
similarity index 100%
rename from nixos/lib/test-driver/test_driver/polling_condition.py
rename to nixos/lib/test-driver/src/test_driver/polling_condition.py
diff --git a/nixos/lib/test-driver/test_driver/py.typed b/nixos/lib/test-driver/src/test_driver/py.typed
similarity index 100%
rename from nixos/lib/test-driver/test_driver/py.typed
rename to nixos/lib/test-driver/src/test_driver/py.typed
diff --git a/nixos/lib/test-driver/test_driver/qmp.py b/nixos/lib/test-driver/src/test_driver/qmp.py
similarity index 100%
rename from nixos/lib/test-driver/test_driver/qmp.py
rename to nixos/lib/test-driver/src/test_driver/qmp.py
diff --git a/nixos/lib/test-driver/test_driver/vlan.py b/nixos/lib/test-driver/src/test_driver/vlan.py
similarity index 100%
rename from nixos/lib/test-driver/test_driver/vlan.py
rename to nixos/lib/test-driver/src/test_driver/vlan.py