nixos/test-driver: fix #153766

nixos/lib/test-driver/default.nix

@@ -14,7 +14,7 @@
 
 python3Packages.buildPythonApplication rec {
   pname = "nixos-test-driver";
-  version = "1.1";
+  version = "1.1.1";
   src = ./.;
 
   propagatedBuildInputs = [ coreutils netpbm python3Packages.colorama python3Packages.ptpython qemu_pkg socat vde2 ]

nixos/lib/test-driver/test_driver/__init__.py

@@ -35,6 +35,13 @@ def __call__(self, parser, namespace, values, option_string=None):  # type: igno
 
 def main() -> None:
     arg_parser = argparse.ArgumentParser(prog="nixos-test-driver")
+    arg_parser.add_argument(
+        "-o",
+        "--output-dir",
+        help="specify the directory where the output (screenshots, etc.) should be placed",
+        default=None,
+        type=Path,
+    )
     arg_parser.add_argument(
         "-K",
         "--keep-vm-state",
@@ -77,7 +84,11 @@ def main() -> None:
         rootlog.info("Machine state will be reset. To keep it, pass --keep-vm-state")
 
     with Driver(
-        args.start_scripts, args.vlans, args.testscript.read_text(), args.keep_vm_state
+        start_scripts=args.start_scripts,
+        vlans=args.vlans,
+        tests=args.testscript.read_text(),
+        keep_vm_state=args.keep_vm_state,
+        out_dir=args.output_dir,
     ) as driver:
         if args.interactive:
             ptpython.repl.embed(driver.test_symbols(), {})

nixos/lib/test-driver/test_driver/driver.py

@@ -1,6 +1,6 @@
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Any, Dict, Iterator, List, Union, Optional, Callable, ContextManager
+from typing import Any, Callable, ContextManager, Dict, Iterator, List, Optional, Union
 import os
 import tempfile
 
@@ -18,19 +18,32 @@ class Driver:
     vlans: List[VLan]
     machines: List[Machine]
     polling_conditions: List[PollingCondition]
+    out_dir: Path
 
     def __init__(
         self,
         start_scripts: List[str],
         vlans: List[int],
         tests: str,
         keep_vm_state: bool = False,
+        out_dir: Optional[Path] = None,
     ):
         self.tests = tests
 
         tmp_dir = Path(os.environ.get("TMPDIR", tempfile.gettempdir()))
+        tmp_dir = tmp_dir.resolve()
         tmp_dir.mkdir(mode=0o700, exist_ok=True)
 
+        if out_dir is None:
+            out_dir = Path(".")
+        out_dir = out_dir.resolve()
+        self.out_dir = out_dir
+
+        # Check that these are 1. a directory 2. writable
+        for (p, name) in [(tmp_dir, "tmp_dir"), (out_dir, "out_dir")]:
+            if not (p.is_dir() and os.access(p, os.W_OK)):
+                raise RuntimeError(f"{name} ({p}) is not a writable directory")
+
         with rootlog.nested("start all VLans"):
             self.vlans = [VLan(nr, tmp_dir) for nr in vlans]
 
@@ -47,6 +60,7 @@ def cmd(scripts: List[str]) -> Iterator[NixStartScript]:
                 name=cmd.machine_name,
                 tmp_dir=tmp_dir,
                 callbacks=[self.check_polling_conditions],
+                out_dir=out_dir,
             )
             for cmd in cmd(start_scripts)
         ]
@@ -158,6 +172,7 @@ def create_machine(self, args: Dict[str, Any]) -> Machine:
             name=name,
             keep_vm_state=args.get("keep_vm_state", False),
             allow_reboot=args.get("allow_reboot", False),
+            out_dir=self.out_dir,
         )
 
     def serial_stdout_on(self) -> None:

nixos/lib/test-driver/test_driver/machine.py

@@ -302,6 +302,7 @@ class Machine:
     state_dir: Path
     monitor_path: Path
     shell_path: Path
+    out_dir: Path
 
     start_command: StartCommand
     keep_vm_state: bool
@@ -327,6 +328,7 @@ def __init__(
         self,
         tmp_dir: Path,
         start_command: StartCommand,
+        out_dir: Path,
         name: str = "machine",
         keep_vm_state: bool = False,
         allow_reboot: bool = False,
@@ -337,6 +339,7 @@ def __init__(
         self.allow_reboot = allow_reboot
         self.name = name
         self.start_command = start_command
+        self.out_dir = out_dir
         self.callbacks = callbacks if callbacks is not None else []
 
         # set up directories
@@ -702,19 +705,21 @@ def connect(self) -> None:
             self.connected = True
 
     def screenshot(self, filename: str) -> None:
-        out_dir = os.environ.get("out", os.getcwd())
-        word_pattern = re.compile(r"^\w+$")
+        word_pattern = re.compile(r"^[\w\-]+$")
         if word_pattern.match(filename):
-            filename = os.path.join(out_dir, "{}.png".format(filename))
-        tmp = "{}.ppm".format(filename)
+            filename = f"{filename}.png"
+
+        # If filename is already absolute, out_dir is ignored.
+        out_file = (self.out_dir / Path(filename)).resolve()
+        tmp = self.tmp_dir / (out_file.stem + ".ppm")
 
         with self.nested(
-            "making screenshot {}".format(filename),
-            {"image": os.path.basename(filename)},
+            f"making screenshot {out_file}",
+            {"image": out_file.name},
         ):
-            self.send_monitor_command("screendump {}".format(tmp))
-            ret = subprocess.run("pnmtopng {} > {}".format(tmp, filename), shell=True)
-            os.unlink(tmp)
+            self.send_monitor_command(f"screendump {tmp}")
+            ret = subprocess.run(f"pnmtopng {tmp} > {out_file}", shell=True)
+            tmp.unlink()
             if ret.returncode != 0:
                 raise Exception("Cannot convert screenshot")
 
@@ -756,7 +761,6 @@ def copy_from_vm(self, source: str, target_dir: str = "") -> None:
         all the VMs (using a temporary directory).
         """
         # Compute the source, target, and intermediate shared file names
-        out_dir = Path(os.environ.get("out", os.getcwd()))
         vm_src = Path(source)
         with tempfile.TemporaryDirectory(dir=self.shared_dir) as shared_td:
             shared_temp = Path(shared_td)
@@ -766,7 +770,7 @@ def copy_from_vm(self, source: str, target_dir: str = "") -> None:
             # Copy the file to the shared directory inside VM
             self.succeed(make_command(["mkdir", "-p", vm_shared_temp]))
             self.succeed(make_command(["cp", "-r", vm_src, vm_intermediate]))
-            abs_target = out_dir / target_dir / vm_src.name
+            abs_target = self.out_dir / target_dir / vm_src.name
             abs_target.parent.mkdir(exist_ok=True, parents=True)
             # Copy the file from the shared directory outside VM
             if intermediate.is_dir():

nixos/lib/testing-python.nix

@@ -30,7 +30,7 @@ rec {
           # effectively mute the XMLLogger
           export LOGFILE=/dev/null
 
-          ${driver}/bin/nixos-test-driver
+          ${driver}/bin/nixos-test-driver --output-dir "$out"
         '';
 
       passthru = driver.passthru // {