nixos/test-driver: Typecheck TestScript

4 files changed, 66 insertions, 3 deletions

diff --git a/nixos/lib/test-driver/default.nix b/nixos/lib/test-driver/default.nix
index 3aee9134318..6cd0ffe222d 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 00000000000..e69de29bb2d
--- /dev/null
+++ b/nixos/lib/test-driver/test_driver/py.typed
diff --git a/nixos/lib/test-script-prepend.py b/nixos/lib/test-script-prepend.py
new file mode 100644
index 00000000000..15e59ce0104
--- /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 c1015ec0aca..b76ce3bc322 100644
--- a/nixos/lib/testing-python.nix
+++ b/nixos/lib/testing-python.nix
@@ -85,7 +85,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 +100,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 +116,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 +128,23 @@ rec {
         mkdir -p $out/bin
 
         vmStartScripts=($(for i in ${toString vms}; do echo $i/bin/run-*-vm; done))
-        echo -n "$testScript" > $out/test-script
+
+        # 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