From a99736e399c68db364485f127ae86e7ce3970f75 Mon Sep 17 00:00:00 2001
From: tljuniper <48209000+tljuniper@users.noreply.github.com>
Date: Mon, 23 May 2022 11:08:13 +0200
Subject: nixos/test-driver: add option to add extra python packages to test
 code

---
 nixos/lib/test-driver/default.nix | 14 ++++++++++++--
 nixos/lib/testing-python.nix      |  8 +++++---
 2 files changed, 17 insertions(+), 5 deletions(-)

(limited to 'nixos')

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;
-- 
cgit 1.4.1


From 8858bf009e8ed2545205686945e64f0406dede9d Mon Sep 17 00:00:00 2001
From: tljuniper <48209000+tljuniper@users.noreply.github.com>
Date: Mon, 23 May 2022 11:08:51 +0200
Subject: nixos/test-driver: add test for extraPythonPackages

---
 nixos/tests/all-tests.nix             |  1 +
 nixos/tests/extra-python-packages.nix | 13 +++++++++++++
 2 files changed, 14 insertions(+)
 create mode 100644 nixos/tests/extra-python-packages.nix

(limited to 'nixos')

diff --git a/nixos/tests/all-tests.nix b/nixos/tests/all-tests.nix
index f4b6ee73562eb..bd0cdcb408aeb 100644
--- a/nixos/tests/all-tests.nix
+++ b/nixos/tests/all-tests.nix
@@ -152,6 +152,7 @@ in
   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.])")
+    '';
+  })
-- 
cgit 1.4.1


From 2a750c302669a59a13d8a2a6fa038cbc6e6cb134 Mon Sep 17 00:00:00 2001
From: tljuniper <48209000+tljuniper@users.noreply.github.com>
Date: Mon, 23 May 2022 13:56:18 +0200
Subject: nixos/manual: Add docs on extra python packages in tests

---
 .../development/writing-nixos-tests.section.md     | 22 ++++++++++++++++++
 .../development/writing-nixos-tests.section.xml    | 26 ++++++++++++++++++++++
 2 files changed, 48 insertions(+)

(limited to 'nixos')

diff --git a/nixos/doc/manual/development/writing-nixos-tests.section.md b/nixos/doc/manual/development/writing-nixos-tests.section.md
index e5ee1cb01ff16..fff8873e61d0d 100644
--- a/nixos/doc/manual/development/writing-nixos-tests.section.md
+++ b/nixos/doc/manual/development/writing-nixos-tests.section.md
@@ -380,3 +380,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 7ce3e4cb29065..36f5f00410fc3 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
@@ -652,4 +652,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>
-- 
cgit 1.4.1