nixos/test-driver: add option to add extra python packages to test code

2 files changed, 17 insertions, 5 deletions

diff --git a/nixos/lib/test-driver/default.nix b/nixos/lib/test-driver/default.nix
index 3aee913431890..a79207267abd9 100644
--- a/nixos/lib/test-driver/default.nix
+++ b/nixos/lib/test-driver/default.nix
@@ -10,6 +10,7 @@
 , socat
 , tesseract4
 , vde2
+, extraPythonPackages ? (_ : [])
 }:
 
 python3Packages.buildPythonApplication rec {
@@ -17,8 +18,17 @@ python3Packages.buildPythonApplication rec {
   version = "1.1";
   src = ./.;
 
-  propagatedBuildInputs = [ coreutils netpbm python3Packages.colorama python3Packages.ptpython qemu_pkg socat vde2 ]
-    ++ (lib.optionals enableOCR [ imagemagick_light tesseract4 ]);
+  propagatedBuildInputs = [
+    coreutils
+    netpbm
+    python3Packages.colorama
+    python3Packages.ptpython
+    qemu_pkg
+    socat
+    vde2
+  ]
+    ++ (lib.optionals enableOCR [ imagemagick_light tesseract4 ])
+    ++ extraPythonPackages python3Packages;
 
   doCheck = true;
   checkInputs = with python3Packages; [ mypy pylint black ];
diff --git a/nixos/lib/testing-python.nix b/nixos/lib/testing-python.nix
index 4ab3cf53045af..79b55f8d1d682 100644
--- a/nixos/lib/testing-python.nix
+++ b/nixos/lib/testing-python.nix
@@ -52,12 +52,13 @@ rec {
     , skipLint ? false
     , passthru ? {}
     , interactive ? false
+    , extraPythonPackages ? (_ :[])
   }:
     let
       # Reifies and correctly wraps the python test driver for
       # the respective qemu version and with or without ocr support
       testDriver = pkgs.callPackage ./test-driver {
-        inherit enableOCR;
+        inherit enableOCR extraPythonPackages;
         qemu_pkg = qemu_test;
         imagemagick_light = imagemagick_light.override { inherit libtiff; };
         tesseract4 = tesseract4.override { enableLanguages = [ "eng" ]; };
@@ -161,6 +162,7 @@ rec {
         (if meta.description or null != null
           then builtins.unsafeGetAttrPos "description" meta
           else builtins.unsafeGetAttrPos "testScript" t)
+    , extraPythonPackages ? (_ : [])
     } @ t:
     let
       mkNodes = qemu_pkg:
@@ -213,13 +215,13 @@ rec {
           );
 
       driver = setupDriverForTest {
-        inherit testScript enableOCR skipLint passthru;
+        inherit testScript enableOCR skipLint passthru extraPythonPackages;
         testName = name;
         qemu_pkg = pkgs.qemu_test;
         nodes = mkNodes pkgs.qemu_test;
       };
       driverInteractive = setupDriverForTest {
-        inherit testScript enableOCR skipLint passthru;
+        inherit testScript enableOCR skipLint passthru extraPythonPackages;
         testName = name;
         qemu_pkg = pkgs.qemu;
         nodes = mkNodes pkgs.qemu;