Merge pull request #174441 from tljuniper/172325-nixostest-override-python-pkgs-additional-param

nixos/test-driver: additional Python packages in test driver

6 files changed, 79 insertions, 5 deletions

diff --git a/nixos/doc/manual/development/writing-nixos-tests.section.md b/nixos/doc/manual/development/writing-nixos-tests.section.md
index 583b8f712b412..f4f4056ad9889 100644
--- a/nixos/doc/manual/development/writing-nixos-tests.section.md
+++ b/nixos/doc/manual/development/writing-nixos-tests.section.md
@@ -393,3 +393,25 @@ with foo_running:
     def foo_running():
         machine.succeed("pgrep -x foo")
     ```
+
+## Adding Python packages to the test script {#ssec-python-packages-in-test-script}
+
+When additional Python libraries are required in the test script, they can be
+added using the parameter `extraPythonPackages`. For example, you could add
+`numpy` like this:
+
+```nix
+import ./make-test-python.nix
+{
+  extraPythonPackages = p: [ p.numpy ];
+
+  nodes = { };
+
+  testScript = ''
+    import numpy as np
+    assert str(np.zeros(4) == "array([0., 0., 0., 0.])")
+  '';
+}
+```
+
+In that case, `numpy` is chosen from the generic `python3Packages`.
diff --git a/nixos/doc/manual/from_md/development/writing-nixos-tests.section.xml b/nixos/doc/manual/from_md/development/writing-nixos-tests.section.xml
index 79df3b9c3764a..46367bdd345d1 100644
--- a/nixos/doc/manual/from_md/development/writing-nixos-tests.section.xml
+++ b/nixos/doc/manual/from_md/development/writing-nixos-tests.section.xml
@@ -665,4 +665,30 @@ def foo_running():
 ```
 </programlisting>
   </section>
+  <section xml:id="ssec-python-packages-in-test-script">
+    <title>Adding Python packages to the test script</title>
+    <para>
+      When additional Python libraries are required in the test script,
+      they can be added using the parameter
+      <literal>extraPythonPackages</literal>. For example, you could add
+      <literal>numpy</literal> like this:
+    </para>
+    <programlisting language="bash">
+import ./make-test-python.nix
+{
+  extraPythonPackages = p: [ p.numpy ];
+
+  nodes = { };
+
+  testScript = ''
+    import numpy as np
+    assert str(np.zeros(4) == &quot;array([0., 0., 0., 0.])&quot;)
+  '';
+}
+</programlisting>
+    <para>
+      In that case, <literal>numpy</literal> is chosen from the generic
+      <literal>python3Packages</literal>.
+    </para>
+  </section>
 </section>
diff --git a/nixos/lib/test-driver/default.nix b/nixos/lib/test-driver/default.nix
index 6cd0ffe222d20..e3786622c3c58 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 8ba2d32ddda54..a6868a708aaf3 100644
--- a/nixos/lib/testing-python.nix
+++ b/nixos/lib/testing-python.nix
@@ -53,12 +53,13 @@ rec {
     , skipTypeCheck ? 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" ]; };
@@ -184,6 +185,7 @@ rec {
         (if meta.description or null != null
           then builtins.unsafeGetAttrPos "description" meta
           else builtins.unsafeGetAttrPos "testScript" t)
+    , extraPythonPackages ? (_ : [])
     } @ t:
     let
       mkNodes = qemu_pkg:
@@ -236,13 +238,13 @@ rec {
           );
 
       driver = setupDriverForTest {
-        inherit testScript enableOCR skipTypeCheck skipLint passthru;
+        inherit testScript enableOCR skipTypeCheck skipLint passthru extraPythonPackages;
         testName = name;
         qemu_pkg = pkgs.qemu_test;
         nodes = mkNodes pkgs.qemu_test;
       };
       driverInteractive = setupDriverForTest {
-        inherit testScript enableOCR skipTypeCheck skipLint passthru;
+        inherit testScript enableOCR skipTypeCheck skipLint passthru extraPythonPackages;
         testName = name;
         qemu_pkg = pkgs.qemu;
         nodes = mkNodes pkgs.qemu;
diff --git a/nixos/tests/all-tests.nix b/nixos/tests/all-tests.nix
index f9e32592e95bc..f534888592ffc 100644
--- a/nixos/tests/all-tests.nix
+++ b/nixos/tests/all-tests.nix
@@ -166,6 +166,7 @@ let
   etcd-cluster = handleTestOn ["x86_64-linux"] ./etcd-cluster.nix {};
   etebase-server = handleTest ./etebase-server.nix {};
   etesync-dav = handleTest ./etesync-dav.nix {};
+  extra-python-packages = handleTest ./extra-python-packages.nix {};
   fancontrol = handleTest ./fancontrol.nix {};
   fcitx = handleTest ./fcitx {};
   fenics = handleTest ./fenics.nix {};
diff --git a/nixos/tests/extra-python-packages.nix b/nixos/tests/extra-python-packages.nix
new file mode 100644
index 0000000000000..7a48077cf98bc
--- /dev/null
+++ b/nixos/tests/extra-python-packages.nix
@@ -0,0 +1,13 @@
+import ./make-test-python.nix ({ ... }:
+  {
+    name = "extra-python-packages";
+
+    extraPythonPackages = p: [ p.numpy ];
+
+    nodes = { };
+
+    testScript = ''
+      import numpy as np
+      assert str(np.zeros(4) == "array([0., 0., 0., 0.])")
+    '';
+  })