diff --git a/nixos/lib/test-driver/setup.py b/nixos/lib/test-driver/setup.py
index 476c7b2dab2a..1719b988db68 100644
--- a/nixos/lib/test-driver/setup.py
+++ b/nixos/lib/test-driver/setup.py
@@ -4,6 +4,7 @@ setup(
   name="nixos-test-driver",
   version='1.1',
   packages=find_packages(),
+  package_data={"test_driver": ["py.typed"]},
   entry_points={
     "console_scripts": [
       "nixos-test-driver=test_driver:main",
diff --git a/nixos/lib/testing-python.nix b/nixos/lib/testing-python.nix
index a6868a708aaf..4bb1689ffd78 100644
--- a/nixos/lib/testing-python.nix
+++ b/nixos/lib/testing-python.nix
@@ -119,6 +119,7 @@ rec {
       {
         inherit testName;
         nativeBuildInputs = [ makeWrapper mypy ];
+        buildInputs = [ testDriver ];
         testScript = testScript';
         preferLocalBuild = true;
         passthru = passthru // {
@@ -138,13 +139,10 @@ rec {
           echo "${builtins.toString vlanNames}" >> testScriptWithTypes
           echo -n "$testScript" >> testScriptWithTypes
 
-          # set pythonpath so mypy knows where to find the imports. this requires the py.typed file.
-          export PYTHONPATH='${./test-driver}'
           mypy  --no-implicit-optional \
                 --pretty \
                 --no-color-output \
                 testScriptWithTypes
-          unset PYTHONPATH
         ''}
 
         echo -n "$testScript" >> $out/test-script