Merge commit '8e2b14aceb1d40c7e8b84c03a7c78955359872bb'

2 files changed, 48 insertions, 7 deletions

diff --git a/nixpkgs/nixos/lib/test-driver/Logger.pm b/nixpkgs/nixos/lib/test-driver/Logger.pm
index 080310ea34e..a3384084a0e 100644
--- a/nixpkgs/nixos/lib/test-driver/Logger.pm
+++ b/nixpkgs/nixos/lib/test-driver/Logger.pm
@@ -8,17 +8,17 @@ use Time::HiRes qw(clock_gettime CLOCK_MONOTONIC);
 
 sub new {
     my ($class) = @_;
-    
+
     my $logFile = defined $ENV{LOGFILE} ? "$ENV{LOGFILE}" : "/dev/null";
     my $log = new XML::Writer(OUTPUT => new IO::File(">$logFile"));
-    
+
     my $self = {
         log => $log,
         logQueue => Thread::Queue->new()
     };
-    
+
     $self->{log}->startTag("logfile");
-    
+
     bless $self, $class;
     return $self;
 }
diff --git a/nixpkgs/nixos/lib/test-driver/test-driver.py b/nixpkgs/nixos/lib/test-driver/test-driver.py
index e7b05968b07..7b8d5803aa5 100644
--- a/nixpkgs/nixos/lib/test-driver/test-driver.py
+++ b/nixpkgs/nixos/lib/test-driver/test-driver.py
@@ -3,7 +3,10 @@ from contextlib import contextmanager, _GeneratorContextManager
 from queue import Queue, Empty
 from typing import Tuple, Any, Callable, Dict, Iterator, Optional, List
 from xml.sax.saxutils import XMLGenerator
+import queue
+import io
 import _thread
+import argparse
 import atexit
 import base64
 import codecs
@@ -19,6 +22,7 @@ import subprocess
 import sys
 import tempfile
 import time
+import traceback
 import unicodedata
 
 CHAR_TO_KEY = {
@@ -671,6 +675,22 @@ class Machine:
         with self.nested("waiting for {} to appear on screen".format(regex)):
             retry(screen_matches)
 
+    def wait_for_console_text(self, regex: str) -> None:
+        self.log("waiting for {} to appear on console".format(regex))
+        # Buffer the console output, this is needed
+        # to match multiline regexes.
+        console = io.StringIO()
+        while True:
+            try:
+                console.write(self.last_lines.get())
+            except queue.Empty:
+                self.sleep(1)
+                continue
+            console.seek(0)
+            matches = re.search(regex, console.read())
+            if matches is not None:
+                return
+
     def send_key(self, key: str) -> None:
         key = CHAR_TO_KEY.get(key, key)
         self.send_monitor_command("sendkey {}".format(key))
@@ -734,11 +754,16 @@ class Machine:
         self.monitor, _ = self.monitor_socket.accept()
         self.shell, _ = self.shell_socket.accept()
 
+        # Store last serial console lines for use
+        # of wait_for_console_text
+        self.last_lines: Queue = Queue()
+
         def process_serial_output() -> None:
             assert self.process.stdout is not None
             for _line in self.process.stdout:
                 # Ignore undecodable bytes that may occur in boot menus
                 line = _line.decode(errors="ignore").replace("\r", "").rstrip()
+                self.last_lines.put(line)
                 eprint("{} # {}".format(self.name, line))
                 self.logger.enqueue({"msg": line, "machine": self.name})
 
@@ -751,6 +776,11 @@ class Machine:
 
         self.log("QEMU running (pid {})".format(self.pid))
 
+    def cleanup_statedir(self) -> None:
+        self.log("delete the VM state directory")
+        if os.path.isfile(self.state_dir):
+            shutil.rmtree(self.state_dir)
+
     def shutdown(self) -> None:
         if not self.booted:
             return
@@ -863,7 +893,8 @@ def run_tests() -> None:
             try:
                 exec(tests, globals())
             except Exception as e:
-                eprint("error: {}".format(str(e)))
+                eprint("error: ")
+                traceback.print_exc()
                 sys.exit(1)
     else:
         ptpython.repl.embed(locals(), globals())
@@ -889,6 +920,15 @@ def subtest(name: str) -> Iterator[None]:
 
 
 if __name__ == "__main__":
+    arg_parser = argparse.ArgumentParser()
+    arg_parser.add_argument(
+        "-K",
+        "--keep-vm-state",
+        help="re-use a VM state coming from a previous run",
+        action="store_true",
+    )
+    (cli_args, vm_scripts) = arg_parser.parse_known_args()
+
     log = Logger()
 
     vlan_nrs = list(dict.fromkeys(os.environ.get("VLANS", "").split()))
@@ -896,8 +936,10 @@ if __name__ == "__main__":
     for nr, vde_socket, _, _ in vde_sockets:
         os.environ["QEMU_VDE_SOCKET_{}".format(nr)] = vde_socket
 
-    vm_scripts = sys.argv[1:]
     machines = [create_machine({"startCommand": s}) for s in vm_scripts]
+    for machine in machines:
+        if not cli_args.keep_vm_state:
+            machine.cleanup_statedir()
     machine_eval = [
         "{0} = machines[{1}]".format(m.name, idx) for idx, m in enumerate(machines)
     ]
@@ -911,7 +953,6 @@ if __name__ == "__main__":
                     continue
                 log.log("killing {} (pid {})".format(machine.name, machine.pid))
                 machine.process.kill()
-
             for _, _, process, _ in vde_sockets:
                 process.terminate()
         log.close()