diff --git a/nixos/doc/manual/development/writing-nixos-tests.section.md b/nixos/doc/manual/development/writing-nixos-tests.section.md
index e5ee1cb01ff1..583b8f712b41 100644
--- a/nixos/doc/manual/development/writing-nixos-tests.section.md
+++ b/nixos/doc/manual/development/writing-nixos-tests.section.md
@@ -332,6 +332,19 @@ repository):
'';
```
+Similarly, the type checking of test scripts can be disabled in the following
+way:
+
+```nix
+import ./make-test-python.nix {
+ skipTypeCheck = true;
+ nodes.machine =
+ { config, pkgs, ... }:
+ { configuration…
+ };
+}
+```
+
## Failing tests early {#ssec-failing-tests-early}
To fail tests early when certain invariables are no longer met (instead of waiting for the build to time out), the decorator `polling_condition` is provided. For example, if we are testing a program `foo` that should not quit after being started, we might write the following:
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 7ce3e4cb2906..79df3b9c3764 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
@@ -589,6 +589,19 @@ import ./make-test-python.nix {
Python code…
# fmt: on
'';
+</programlisting>
+ <para>
+ Similarly, the type checking of test scripts can be disabled in
+ the following way:
+ </para>
+ <programlisting language="bash">
+import ./make-test-python.nix {
+ skipTypeCheck = true;
+ nodes.machine =
+ { config, pkgs, ... }:
+ { configuration…
+ };
+}
</programlisting>
</section>
<section xml:id="ssec-failing-tests-early">
diff --git a/nixos/lib/test-driver/default.nix b/nixos/lib/test-driver/default.nix
index 3aee91343189..6cd0ffe222d2 100644
--- a/nixos/lib/test-driver/default.nix
+++ b/nixos/lib/test-driver/default.nix
@@ -25,6 +25,8 @@ python3Packages.buildPythonApplication rec {
checkPhase = ''
mypy --disallow-untyped-defs \
--no-implicit-optional \
+ --pretty \
+ --no-color-output \
--ignore-missing-imports ${src}/test_driver
pylint --errors-only --enable=unused-import ${src}/test_driver
black --check --diff ${src}/test_driver
diff --git a/nixos/lib/test-driver/test_driver/py.typed b/nixos/lib/test-driver/test_driver/py.typed
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/nixos/lib/test-script-prepend.py b/nixos/lib/test-script-prepend.py
new file mode 100644
index 000000000000..15e59ce01047
--- /dev/null
+++ b/nixos/lib/test-script-prepend.py
@@ -0,0 +1,42 @@
+# This file contains type hints that can be prepended to Nix test scripts so they can be type
+# checked.
+
+from test_driver.driver import Driver
+from test_driver.vlan import VLan
+from test_driver.machine import Machine
+from test_driver.logger import Logger
+from typing import Callable, Iterator, ContextManager, Optional, List, Dict, Any, Union
+from typing_extensions import Protocol
+from pathlib import Path
+
+
+class RetryProtocol(Protocol):
+ def __call__(self, fn: Callable, timeout: int = 900) -> None:
+ raise Exception("This is just type information for the Nix test driver")
+
+
+class PollingConditionProtocol(Protocol):
+ def __call__(
+ self,
+ fun_: Optional[Callable] = None,
+ *,
+ seconds_interval: float = 2.0,
+ description: Optional[str] = None,
+ ) -> Union[Callable[[Callable], ContextManager], ContextManager]:
+ raise Exception("This is just type information for the Nix test driver")
+
+
+start_all: Callable[[], None]
+subtest: Callable[[str], ContextManager[None]]
+retry: RetryProtocol
+test_script: Callable[[], None]
+machines: List[Machine]
+vlans: List[VLan]
+driver: Driver
+log: Logger
+create_machine: Callable[[Dict[str, Any]], Machine]
+run_tests: Callable[[], None]
+join_all: Callable[[], None]
+serial_stdout_off: Callable[[], None]
+serial_stdout_on: Callable[[], None]
+polling_condition: PollingConditionProtocol
diff --git a/nixos/lib/testing-python.nix b/nixos/lib/testing-python.nix
index c1015ec0aca0..8ba2d32ddda5 100644
--- a/nixos/lib/testing-python.nix
+++ b/nixos/lib/testing-python.nix
@@ -50,6 +50,7 @@ rec {
, qemu_pkg ? pkgs.qemu_test
, enableOCR ? false
, skipLint ? false
+ , skipTypeCheck ? false
, passthru ? {}
, interactive ? false
}:
@@ -85,7 +86,7 @@ rec {
nodeHostNames = let
nodesList = map (c: c.config.system.name) (lib.attrValues nodes);
- in nodesList ++ lib.optional (lib.length nodesList == 1) "machine";
+ in nodesList ++ lib.optional (lib.length nodesList == 1 && !lib.elem "machine" nodesList) "machine";
# TODO: This is an implementation error and needs fixing
# the testing famework cannot legitimately restrict hostnames further
@@ -100,6 +101,9 @@ rec {
then testScript { inherit nodes; }
else testScript;
+ uniqueVlans = lib.unique (builtins.concatLists vlans);
+ vlanNames = map (i: "vlan${toString i}: VLan;") uniqueVlans;
+ machineNames = map (name: "${name}: Machine;") nodeHostNames;
in
if lib.length invalidNodeNames > 0 then
throw ''
@@ -113,7 +117,7 @@ rec {
else lib.warnIf skipLint "Linting is disabled" (runCommand testDriverName
{
inherit testName;
- nativeBuildInputs = [ makeWrapper ];
+ nativeBuildInputs = [ makeWrapper mypy ];
testScript = testScript';
preferLocalBuild = true;
passthru = passthru // {
@@ -125,7 +129,25 @@ rec {
mkdir -p $out/bin
vmStartScripts=($(for i in ${toString vms}; do echo $i/bin/run-*-vm; done))
- echo -n "$testScript" > $out/test-script
+
+ ${lib.optionalString (!skipTypeCheck) ''
+ # prepend type hints so the test script can be type checked with mypy
+ cat "${./test-script-prepend.py}" >> testScriptWithTypes
+ echo "${builtins.toString machineNames}" >> testScriptWithTypes
+ 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
+
ln -s ${testDriver}/bin/nixos-test-driver $out/bin/nixos-test-driver
${testDriver}/bin/generate-driver-symbols
@@ -152,6 +174,7 @@ rec {
, testScript
, enableOCR ? false
, name ? "unnamed"
+ , skipTypeCheck ? false
# Skip linting (mainly intended for faster dev cycles)
, skipLint ? false
, passthru ? {}
@@ -213,13 +236,13 @@ rec {
);
driver = setupDriverForTest {
- inherit testScript enableOCR skipLint passthru;
+ inherit testScript enableOCR skipTypeCheck skipLint passthru;
testName = name;
qemu_pkg = pkgs.qemu_test;
nodes = mkNodes pkgs.qemu_test;
};
driverInteractive = setupDriverForTest {
- inherit testScript enableOCR skipLint passthru;
+ inherit testScript enableOCR skipTypeCheck skipLint passthru;
testName = name;
qemu_pkg = pkgs.qemu;
nodes = mkNodes pkgs.qemu;