diff --git a/nixos/lib/test-driver/src/extract-docstrings.py b/nixos/lib/test-driver/src/extract-docstrings.py index 64850ca711f3..030ed0189704 100644 --- a/nixos/lib/test-driver/src/extract-docstrings.py +++ b/nixos/lib/test-driver/src/extract-docstrings.py @@ -51,7 +51,7 @@ def main() -> None: class_definitions = (node for node in module.body if isinstance(node, ast.ClassDef)) - machine_class = next(filter(lambda x: x.name == "Machine", class_definitions)) + machine_class = next(filter(lambda x: x.name == "BaseMachine", class_definitions)) assert machine_class is not None function_definitions = [ diff --git a/nixos/lib/test-driver/src/test_driver/__init__.py b/nixos/lib/test-driver/src/test_driver/__init__.py index 557cb1ff8142..35c5e3b11fb6 100755 --- a/nixos/lib/test-driver/src/test_driver/__init__.py +++ b/nixos/lib/test-driver/src/test_driver/__init__.py @@ -71,10 +71,18 @@ def main() -> None: help="Enable interactive debugging breakpoints for sandboxed runs", ) arg_parser.add_argument( - "--start-scripts", - metavar="START-SCRIPT", + "--vm-names", + metavar="VM-NAME", action=EnvDefault, - envvar="startScripts", + envvar="vmNames", + nargs="*", + help="names of participating virtual machines", + ) + arg_parser.add_argument( + "--vm-start-scripts", + metavar="VM-START-SCRIPT", + action=EnvDefault, + envvar="vmStartScripts", nargs="*", help="start scripts for participating virtual machines", ) @@ -138,14 +146,20 @@ def main() -> None: if args.debug_hook_attach is not None: debugger = Debug(logger, args.debug_hook_attach) + if args.vm_names is not None and args.vm_start_scripts is not None: + assert len(args.vm_names) == len(args.vm_start_scripts), ( + f"the number of vm names and vm start scripts must be the same: {args.vm_names} vs. {args.vm_start_scripts}" + ) + with Driver( - args.start_scripts, - args.vlans, - args.testscript.read_text(), - output_directory, - logger, - args.keep_vm_state, - args.global_timeout, + vm_names=args.vm_names, + vm_start_scripts=args.vm_start_scripts or [], + vlans=args.vlans, + tests=args.testscript.read_text(), + out_dir=output_directory, + logger=logger, + keep_vm_state=args.keep_vm_state, + global_timeout=args.global_timeout, debug=debugger, ) as driver: if offset := args.dump_vsocks: @@ -170,7 +184,14 @@ def generate_driver_symbols() -> None: in user's test scripts. That list is then used by pyflakes to lint those scripts. """ - d = Driver([], [], "", Path(), CompositeLogger([])) + d = Driver( + vm_names=[], + vm_start_scripts=[], + vlans=[], + tests="", + out_dir=Path(), + logger=CompositeLogger([]), + ) test_symbols = d.test_symbols() with open("driver-symbols", "w") as fp: fp.write(",".join(test_symbols.keys())) diff --git a/nixos/lib/test-driver/src/test_driver/driver.py b/nixos/lib/test-driver/src/test_driver/driver.py index c4f268404cbc..c8694d2ea09a 100644 --- a/nixos/lib/test-driver/src/test_driver/driver.py +++ b/nixos/lib/test-driver/src/test_driver/driver.py @@ -16,7 +16,7 @@ from colorama import Style from test_driver.debug import DebugAbstract, DebugNop from test_driver.errors import MachineError, RequestedAssertionFailed from test_driver.logger import AbstractLogger -from test_driver.machine import Machine, NixStartScript, retry +from test_driver.machine import BaseMachine, QemuMachine, retry from test_driver.polling_condition import PollingCondition from test_driver.vlan import VLan @@ -63,7 +63,7 @@ class Driver: tests: str vlans: list[VLan] - machines: list[Machine] + vm_machines: list[QemuMachine] polling_conditions: list[PollingCondition] global_timeout: int race_timer: threading.Timer @@ -72,7 +72,8 @@ class Driver: def __init__( self, - start_scripts: list[str], + vm_names: list[str] | None, + vm_start_scripts: list[str], vlans: list[int], tests: str, out_dir: Path, @@ -94,25 +95,30 @@ class Driver: vlans = list(set(vlans)) self.vlans = [VLan(nr, tmp_dir, self.logger) for nr in vlans] - def cmd(scripts: list[str]) -> Iterator[NixStartScript]: - for s in scripts: - yield NixStartScript(s) - self.polling_conditions = [] - self.machines = [ - Machine( - start_command=cmd, + self.vm_machines = [ + QemuMachine( + name=name, + start_command=vm_start_script, keep_vm_state=keep_vm_state, - name=cmd.machine_name, tmp_dir=tmp_dir, callbacks=[self.check_polling_conditions], out_dir=self.out_dir, logger=self.logger, ) - for cmd in cmd(start_scripts) + for name, vm_start_script in zip( + vm_names or (len(vm_start_scripts) * [None]), vm_start_scripts + ) ] + @property + def machines(self) -> list[QemuMachine]: + machines = self.vm_machines + # Sort the machines by name for consistency with `nodes` in . + machines.sort(key=lambda machine: machine.name) + return machines + def __enter__(self) -> "Driver": return self @@ -148,7 +154,7 @@ class Driver: general_symbols = dict( start_all=self.start_all, test_script=self.test_script, - machines=self.machines, + vm_machines=self.vm_machines, vlans=self.vlans, driver=self, log=self.logger, @@ -161,7 +167,7 @@ class Driver: serial_stdout_off=self.serial_stdout_off, serial_stdout_on=self.serial_stdout_on, polling_condition=self.polling_condition, - Machine=Machine, # for typing + BaseMachine=BaseMachine, # for typing t=AssertionTester(), debug=self.debug, ) @@ -186,14 +192,14 @@ class Driver: def dump_machine_ssh(self, offset: int) -> None: print("SSH backdoor enabled, the machines can be accessed like this:") print( - f"{Style.BRIGHT}Note:{Style.RESET_ALL} this requires {Style.BRIGHT}systemd-ssh-proxy(1){Style.RESET_ALL} to be enabled (default on NixOS 25.05 and newer)." + f"{Style.BRIGHT}Note:{Style.RESET_ALL} vsocks require {Style.BRIGHT}systemd-ssh-proxy(1){Style.RESET_ALL} to be enabled (default on NixOS 25.05 and newer)." ) - names = [machine.name for machine in self.machines] - longest_name = len(max(names, key=len)) - for num, name in enumerate(names, start=offset + 1): + longest_name = len(max((machine.name for machine in self.machines), key=len)) + for index, machine in enumerate(self.machines, start=offset + 1): + name = machine.name spaces = " " * (longest_name - len(name) + 2) print( - f" {name}:{spaces}{Style.BRIGHT}ssh -o User=root vsock/{num}{Style.RESET_ALL}" + f" {name}:{spaces}{Style.BRIGHT}{machine.ssh_backdoor_command(index)}{Style.RESET_ALL}" ) def test_script(self) -> None: @@ -280,16 +286,13 @@ class Driver: *, name: str | None = None, keep_vm_state: bool = False, - ) -> Machine: + ) -> QemuMachine: tmp_dir = get_tmp_dir() - cmd = NixStartScript(start_command) - name = name or cmd.machine_name - - return Machine( + return QemuMachine( + start_command=start_command, tmp_dir=tmp_dir, out_dir=self.out_dir, - start_command=cmd, name=name, keep_vm_state=keep_vm_state, logger=self.logger, diff --git a/nixos/lib/test-driver/src/test_driver/machine/__init__.py b/nixos/lib/test-driver/src/test_driver/machine/__init__.py index f722d36ae40e..fecdeda493b5 100644 --- a/nixos/lib/test-driver/src/test_driver/machine/__init__.py +++ b/nixos/lib/test-driver/src/test_driver/machine/__init__.py @@ -13,6 +13,7 @@ import sys import tempfile import threading import time +from abc import ABC, abstractmethod from collections.abc import Callable, Generator from contextlib import _GeneratorContextManager, contextmanager, nullcontext from pathlib import Path @@ -114,15 +115,30 @@ def retry(fn: Callable, timeout_seconds: int = 900) -> None: ) -class StartCommand: - """The Base Start Command knows how to append the necessary +class QemuStartCommand: + """This class knows how to append the necessary runtime qemu options as determined by a particular test driver - run. Any such start command is expected to happily receive and - append additional qemu args. + run. """ _cmd: str + def __init__(self, script: str): + self._cmd = script + + @property + def machine_name(self) -> str: + """A start script from nixos/modules/virtualiation/qemu-vm.nix. + These Nix commands have the particular characteristic that the + machine name can be extracted out of them via a regex match. + (Admittedly a _very_ implicit contract, evtl. TODO fix) + """ + match = re.search("run-(.+)-vm$", self._cmd) + name = "machine" + if match: + name = match.group(1) + return name + def cmd( self, monitor_socket_path: Path, @@ -198,103 +214,42 @@ class StartCommand: ) -class NixStartScript(StartCommand): - """A start script from nixos/modules/virtualiation/qemu-vm.nix. - These Nix commands have the particular characteristic that the - machine name can be extracted out of them via a regex match. - (Admittedly a _very_ implicit contract, evtl. TODO fix) - """ - - def __init__(self, script: str): - self._cmd = script - - @property - def machine_name(self) -> str: - match = re.search("run-(.+)-vm$", self._cmd) - name = "machine" - if match: - name = match.group(1) - return name - - -class Machine: - """A handle to the machine with this name, that also knows how to manage - the machine lifecycle with the help of a start script / command.""" - +class BaseMachine(ABC): name: str - out_dir: Path + callbacks: list[Callable] tmp_dir: Path - shared_dir: Path - state_dir: Path - monitor_path: Path - qmp_path: Path - shell_path: Path - start_command: StartCommand keep_vm_state: bool - process: subprocess.Popen | None - pid: int | None - monitor: socket.socket | None - qmp_client: QMPSession | None - shell: socket.socket | None - serial_thread: threading.Thread | None - - booted: bool - connected: bool - # Store last serial console lines for use - # of wait_for_console_text - last_lines: Queue = Queue() - # Store all console output for full log retrieval - full_console_log: list[str] - callbacks: list[Callable] - def __repr__(self) -> str: - return f"" + return f"<{self.__class__.__name__} '{self.name}'>" def __init__( self, out_dir: Path, - tmp_dir: Path, - start_command: StartCommand, + name: str, logger: AbstractLogger, - name: str = "machine", - keep_vm_state: bool = False, - callbacks: list[Callable] | None = None, + tmp_dir: Path, + callbacks: list[Callable] | None, + keep_vm_state: bool, ) -> None: self.out_dir = out_dir - self.tmp_dir = tmp_dir - self.keep_vm_state = keep_vm_state self.name = name - self.start_command = start_command - self.callbacks = callbacks if callbacks is not None else [] self.logger = logger - self.full_console_log = [] + self.callbacks = callbacks if callbacks is not None else [] + self.tmp_dir = tmp_dir - # set up directories - self.shared_dir = self.tmp_dir / "shared-xchg" - self.shared_dir.mkdir(mode=0o700, exist_ok=True) + # Note: "vm" is a bit of a misnomer here. + # Consider renaming to something more generic ("machine"?) + self.keep_vm_state = keep_vm_state self.state_dir = self.tmp_dir / f"vm-state-{self.name}" - self.monitor_path = self.state_dir / "monitor" - self.qmp_path = self.state_dir / "qmp" - self.shell_path = self.state_dir / "shell" if (not self.keep_vm_state) and self.state_dir.exists(): self.cleanup_statedir() self.state_dir.mkdir(mode=0o700, exist_ok=True) - self.process = None - self.pid = None - self.monitor = None - self.qmp_client = None - self.shell = None - self.serial_thread = None - - self.booted = False - self.connected = False - - def is_up(self) -> bool: - return self.booted and self.connected + self.shared_dir = self.tmp_dir / "shared-xchg" + self.shared_dir.mkdir(mode=0o700, exist_ok=True) def log(self, msg: str) -> None: """ @@ -313,28 +268,50 @@ class Machine: my_attrs.update(attrs) return self.logger.nested(msg, my_attrs) - def wait_for_monitor_prompt(self) -> str: - assert self.monitor is not None - answer = "" - while True: - undecoded_answer = self.monitor.recv(1024) - if not undecoded_answer: - break - answer += undecoded_answer.decode() - if answer.endswith("(qemu) "): - break - return answer + @abstractmethod + def is_up(self) -> bool: + """ + Check whether the machine is running. + """ + pass - def send_monitor_command(self, command: str) -> str: + @abstractmethod + def start(self) -> None: """ - Send a command to the QEMU monitor. This allows attaching - virtual USB disks to a running machine, among other things. + Start the machine. """ - self.run_callbacks() - message = f"{command}\n".encode() - assert self.monitor is not None - self.monitor.send(message) - return self.wait_for_monitor_prompt() + pass + + @abstractmethod + def wait_for_shutdown(self) -> None: + """ + Wait for the machine to power off. This does *not* initiate a shutdown; + that's usually done via `shutdown()`. + """ + pass + + def systemctl(self, q: str, user: str | None = None) -> tuple[int, str]: + """ + Runs `systemctl` commands with optional support for + `systemctl --user` + + ```py + # run `systemctl list-jobs --no-pager` + machine.systemctl("list-jobs --no-pager") + + # spawn a shell for `any-user` and run + # `systemctl --user list-jobs --no-pager` + machine.systemctl("list-jobs --no-pager", "any-user") + ``` + """ + if user is not None: + q = q.replace("'", "\\'") + return self.execute( + f"su -l {user} --shell /bin/sh -c " + "$'XDG_RUNTIME_DIR=/run/user/`id -u` " + f"systemctl --user {q}'" + ) + return self.execute(f"systemctl {q}") def wait_for_unit( self, unit: str, user: str | None = None, timeout: int = 900 @@ -424,29 +401,6 @@ class Machine: assert match[1] == property, invalid_output_message return match[2] - def systemctl(self, q: str, user: str | None = None) -> tuple[int, str]: - """ - Runs `systemctl` commands with optional support for - `systemctl --user` - - ```py - # run `systemctl list-jobs --no-pager` - machine.systemctl("list-jobs --no-pager") - - # spawn a shell for `any-user` and run - # `systemctl --user list-jobs --no-pager` - machine.systemctl("list-jobs --no-pager", "any-user") - ``` - """ - if user is not None: - q = q.replace("'", "\\'") - return self.execute( - f"su -l {user} --shell /bin/sh -c " - "$'XDG_RUNTIME_DIR=/run/user/`id -u` " - f"systemctl --user {q}'" - ) - return self.execute(f"systemctl {q}") - def require_unit_state(self, unit: str, require_state: str = "active") -> None: """ Assert that the current state of a unit has a specific value. The default state is "active". @@ -462,6 +416,386 @@ class Machine: f"'{require_state}' but it is in state '{state}'" ) + def succeed(self, *commands: str, timeout: int | None = None) -> str: + """ + Execute a shell command, raising an exception if the exit status is + not zero, otherwise returning the standard output. Similar to `execute`, + except that the timeout is `None` by default. See `execute` for details on + command execution. + """ + output = "" + for command in commands: + with self.nested(f"must succeed: {command}"): + (status, out) = self.execute(command, timeout=timeout) + if status != 0: + self.log(f"output: {out}") + raise RequestedAssertionFailed( + f"command `{command}` failed (exit code {status})" + ) + output += out + return output + + def fail(self, *commands: str, timeout: int | None = None) -> str: + """ + Like `succeed`, but raising an exception if the command returns a zero + status. + """ + output = "" + for command in commands: + with self.nested(f"must fail: {command}"): + (status, out) = self.execute(command, timeout=timeout) + if status == 0: + raise RequestedAssertionFailed( + f"command `{command}` unexpectedly succeeded" + ) + output += out + return output + + def wait_until_succeeds(self, command: str, timeout: int = 900) -> str: + """ + Repeat a shell command with 1-second intervals until it succeeds. + Has a default timeout of 900 seconds which can be modified, e.g. + `wait_until_succeeds(cmd, timeout=10)`. See `execute` for details on + command execution. + Throws an exception on timeout. + """ + output = "" + + def check_success(_last_try: bool) -> bool: + nonlocal output + status, output = self.execute(command, timeout=timeout) + return status == 0 + + with self.nested(f"waiting for success: {command}"): + retry(check_success, timeout) + return output + + def wait_until_fails(self, command: str, timeout: int = 900) -> str: + """ + Like `wait_until_succeeds`, but repeating the command until it fails. + """ + output = "" + + def check_failure(_last_try: bool) -> bool: + nonlocal output + status, output = self.execute(command, timeout=timeout) + return status != 0 + + with self.nested(f"waiting for failure: {command}"): + retry(check_failure, timeout) + return output + + def sleep(self, secs: int) -> None: + # We want to sleep in *guest* time, not *host* time. + self.succeed(f"sleep {secs}") + + def wait_for_file(self, filename: str, timeout: int = 900) -> None: + """ + Waits until the file exists in the machine's file system. + """ + + def check_file(_last_try: bool) -> bool: + status, _ = self.execute(f"test -e {filename}") + return status == 0 + + with self.nested(f"waiting for file '{filename}'"): + retry(check_file, timeout) + + def wait_for_open_port( + self, port: int, addr: str = "localhost", timeout: int = 900 + ) -> None: + """ + Wait until a process is listening on the given TCP port and IP address + (default `localhost`). + """ + + def port_is_open(_last_try: bool) -> bool: + status, _ = self.execute(f"nc -z {addr} {port}") + return status == 0 + + with self.nested(f"waiting for TCP port {port} on {addr}"): + retry(port_is_open, timeout) + + def wait_for_open_unix_socket( + self, addr: str, is_datagram: bool = False, timeout: int = 900 + ) -> None: + """ + Wait until a process is listening on the given UNIX-domain socket + (default to a UNIX-domain stream socket). + """ + + nc_flags = [ + "-z", + "-uU" if is_datagram else "-U", + ] + + def socket_is_open(_last_try: bool) -> bool: + status, _ = self.execute(f"nc {' '.join(nc_flags)} {addr}") + return status == 0 + + with self.nested( + f"waiting for UNIX-domain {'datagram' if is_datagram else 'stream'} on '{addr}'" + ): + retry(socket_is_open, timeout) + + def wait_for_closed_port( + self, port: int, addr: str = "localhost", timeout: int = 900 + ) -> None: + """ + Wait until nobody is listening on the given TCP port and IP address + (default `localhost`). + """ + + def port_is_closed(_last_try: bool) -> bool: + status, _ = self.execute(f"nc -z {addr} {port}") + return status != 0 + + with self.nested(f"waiting for TCP port {port} on {addr} to be closed"): + retry(port_is_closed, timeout) + + def start_job(self, jobname: str, user: str | None = None) -> tuple[int, str]: + """ + Start systemd service. + """ + return self.systemctl(f"start {jobname}", user) + + def stop_job(self, jobname: str, user: str | None = None) -> tuple[int, str]: + """ + Stop systemd service. + """ + return self.systemctl(f"stop {jobname}", user) + + def wait_for_job(self, jobname: str) -> None: + self.wait_for_unit(jobname) + + def get_tty_text(self, tty: str) -> str: + """ + Get the output printed to a given TTY. + """ + status, output = self.execute( + f"fold -w$(stty -F /dev/tty{tty} size | awk '{{print $2}}') /dev/vcs{tty}" + ) + return output + + def wait_until_tty_matches(self, tty: str, regexp: str, timeout: int = 900) -> None: + """Wait until the visible output on the chosen TTY matches regular + expression. Throws an exception on timeout. + """ + matcher = re.compile(regexp) + + def tty_matches(last_try: bool) -> bool: + text = self.get_tty_text(tty) + if last_try: + self.log( + f"Last chance to match /{regexp}/ on TTY{tty}, " + f"which currently contains: {text}" + ) + return len(matcher.findall(text)) > 0 + + with self.nested(f"waiting for {regexp} to appear on tty {tty}"): + retry(tty_matches, timeout) + + def dump_tty_contents(self, tty: str) -> None: + """Debugging: Dump the contents of the TTY""" + self.execute(f"fold -w 80 /dev/vcs{tty} | systemd-cat") + + def execute( + self, + command: str, + check_return: bool = True, + check_output: bool = True, + timeout: int | None = 900, + ) -> tuple[int, str]: + self.run_callbacks() + return self._execute( + command=command, + check_return=check_return, + check_output=check_output, + timeout=timeout, + ) + + @abstractmethod + def _execute( + self, + command: str, + check_return: bool = True, + check_output: bool = True, + timeout: int | None = 900, + ) -> tuple[int, str]: ... + + def run_callbacks(self) -> None: + for callback in self.callbacks: + callback() + + def cleanup_statedir(self) -> None: + shutil.rmtree(self.state_dir) + self.logger.log(f"deleting VM state directory {self.state_dir}") + self.logger.log("if you want to keep the VM state, pass --keep-vm-state") + + def copy_from_vm(self, source: str, target_dir: str = "") -> None: + """Copy a file from the VM (specified by an in-VM source path) to a path + relative to `$out`. The file is copied via the `shared_dir` shared among + all the VMs (using a temporary directory). + """ + # Compute the source, target, and intermediate shared file names + vm_src = Path(source) + with tempfile.TemporaryDirectory(dir=self.shared_dir) as shared_td: + shared_temp = Path(shared_td) + vm_shared_temp = Path("/tmp/shared") / shared_temp.name + vm_intermediate = vm_shared_temp / vm_src.name + intermediate = shared_temp / vm_src.name + # 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 = 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(): + shutil.copytree(intermediate, abs_target) + else: + shutil.copy(intermediate, abs_target) + + def copy_from_host_via_shell(self, source: str, target: str) -> None: + """Copy a file from the host into the guest by piping it over the + shell into the destination file. Works without host-guest shared folder. + Prefer copy_from_host for whenever possible. + """ + with open(source, "rb") as fh: + content_b64 = base64.b64encode(fh.read()).decode() + self.succeed( + f"mkdir -p $(dirname {target})", + f"echo -n {content_b64} | base64 -d > {target}", + ) + + def copy_from_host(self, source: str, target: str) -> None: + """ + Copies a file from host to machine, e.g., + `copy_from_host("myfile", "/etc/my/important/file")`. + + The first argument is the file on the host. Note that the "host" refers + to the environment in which the test driver runs, which is typically the + Nix build sandbox. + + The second argument is the location of the file on the machine that will + be written to. + + The file is copied via the `shared_dir` directory which is shared among + all the VMs (using a temporary directory). + The access rights bits will mimic the ones from the host file and + user:group will be root:root. + """ + host_src = Path(source) + vm_target = Path(target) + with tempfile.TemporaryDirectory(dir=self.shared_dir) as shared_td: + shared_temp = Path(shared_td) + host_intermediate = shared_temp / host_src.name + vm_shared_temp = Path("/tmp/shared") / shared_temp.name + vm_intermediate = vm_shared_temp / host_src.name + + self.succeed(make_command(["mkdir", "-p", vm_shared_temp])) + if host_src.is_dir(): + shutil.copytree(host_src, host_intermediate) + else: + shutil.copy(host_src, host_intermediate) + self.succeed(make_command(["mkdir", "-p", vm_target.parent])) + self.succeed(make_command(["cp", "-r", vm_intermediate, vm_target])) + + +class QemuMachine(BaseMachine): + """A handle to the machine with this name, that also knows how to manage + the machine lifecycle with the help of a start script / command.""" + + name: str + out_dir: Path + shared_dir: Path + state_dir: Path + monitor_path: Path + qmp_path: Path + shell_path: Path + + start_command: QemuStartCommand + + process: subprocess.Popen | None + pid: int | None + monitor: socket.socket | None + qmp_client: QMPSession | None + shell: socket.socket | None + serial_thread: threading.Thread | None + + booted: bool + connected: bool + # Store last serial console lines for use + # of wait_for_console_text + last_lines: Queue = Queue() + # Store all console output for full log retrieval + full_console_log: list[str] + + def __init__( + self, + out_dir: Path, + tmp_dir: Path, + start_command: str, + logger: AbstractLogger, + name: str | None = None, + keep_vm_state: bool = False, + callbacks: list[Callable] | None = None, + ) -> None: + self.start_command = QemuStartCommand(start_command) + super().__init__( + out_dir=out_dir, + name=name or self.start_command.machine_name, + logger=logger, + callbacks=callbacks, + tmp_dir=tmp_dir, + keep_vm_state=keep_vm_state, + ) + + self.full_console_log = [] + + # set up directories + self.monitor_path = self.state_dir / "monitor" + self.qmp_path = self.state_dir / "qmp" + self.shell_path = self.state_dir / "shell" + + self.process = None + self.pid = None + self.monitor = None + self.qmp_client = None + self.shell = None + self.serial_thread = None + + self.booted = False + self.connected = False + + def ssh_backdoor_command(self, index: int) -> str: + return f"ssh -o User=root vsock/{index}" + + def is_up(self) -> bool: + return self.booted and self.connected + + def wait_for_monitor_prompt(self) -> str: + assert self.monitor is not None + answer = "" + while True: + undecoded_answer = self.monitor.recv(1024) + if not undecoded_answer: + break + answer += undecoded_answer.decode() + if answer.endswith("(qemu) "): + break + return answer + + def send_monitor_command(self, command: str) -> str: + """ + Send a command to the QEMU monitor. This allows attaching + virtual USB disks to a running machine, among other things. + """ + self.run_callbacks() + message = f"{command}\n".encode() + assert self.monitor is not None + self.monitor.send(message) + return self.wait_for_monitor_prompt() + def _next_newline_closed_block_from_shell(self) -> str: assert self.shell output_buffer = [] @@ -478,7 +812,7 @@ class Machine: break return "".join(output_buffer) - def execute( + def _execute( self, command: str, check_return: bool = True, @@ -517,7 +851,6 @@ class Machine: `timeout` parameter, e.g., `execute(cmd, timeout=10)` or `execute(cmd, timeout=None)`. The default is 900 seconds. """ - self.run_callbacks() self.connect() # Always run command with shell opts @@ -598,75 +931,6 @@ class Machine: break self.send_console(char.decode()) - def succeed(self, *commands: str, timeout: int | None = None) -> str: - """ - Execute a shell command, raising an exception if the exit status is - not zero, otherwise returning the standard output. Similar to `execute`, - except that the timeout is `None` by default. See `execute` for details on - command execution. - """ - output = "" - for command in commands: - with self.nested(f"must succeed: {command}"): - (status, out) = self.execute(command, timeout=timeout) - if status != 0: - self.log(f"output: {out}") - raise RequestedAssertionFailed( - f"command `{command}` failed (exit code {status})" - ) - output += out - return output - - def fail(self, *commands: str, timeout: int | None = None) -> str: - """ - Like `succeed`, but raising an exception if the command returns a zero - status. - """ - output = "" - for command in commands: - with self.nested(f"must fail: {command}"): - (status, out) = self.execute(command, timeout=timeout) - if status == 0: - raise RequestedAssertionFailed( - f"command `{command}` unexpectedly succeeded" - ) - output += out - return output - - def wait_until_succeeds(self, command: str, timeout: int = 900) -> str: - """ - Repeat a shell command with 1-second intervals until it succeeds. - Has a default timeout of 900 seconds which can be modified, e.g. - `wait_until_succeeds(cmd, timeout=10)`. See `execute` for details on - command execution. - Throws an exception on timeout. - """ - output = "" - - def check_success(_last_try: bool) -> bool: - nonlocal output - status, output = self.execute(command, timeout=timeout) - return status == 0 - - with self.nested(f"waiting for success: {command}"): - retry(check_success, timeout) - return output - - def wait_until_fails(self, command: str, timeout: int = 900) -> str: - """ - Like `wait_until_succeeds`, but repeating the command until it fails. - """ - output = "" - - def check_failure(_last_try: bool) -> bool: - nonlocal output - status, output = self.execute(command, timeout=timeout) - return status != 0 - - with self.nested(f"waiting for failure: {command}"): - retry(check_failure, timeout) - return output - def wait_for_shutdown(self) -> None: """ Wait for the VM to power off. This does *not* initiate a shutdown; @@ -710,33 +974,6 @@ class Machine: if elapsed >= timeout: raise TimeoutError - def get_tty_text(self, tty: str) -> str: - """ - Get the output printed to a given TTY. - """ - status, output = self.execute( - f"fold -w$(stty -F /dev/tty{tty} size | awk '{{print $2}}') /dev/vcs{tty}" - ) - return output - - def wait_until_tty_matches(self, tty: str, regexp: str, timeout: int = 900) -> None: - """Wait until the visible output on the chosen TTY matches regular - expression. Throws an exception on timeout. - """ - matcher = re.compile(regexp) - - def tty_matches(last_try: bool) -> bool: - text = self.get_tty_text(tty) - if last_try: - self.log( - f"Last chance to match /{regexp}/ on TTY{tty}, " - f"which currently contains: {text}" - ) - return len(matcher.findall(text)) > 0 - - with self.nested(f"waiting for {regexp} to appear on tty {tty}"): - retry(tty_matches, timeout) - def send_chars(self, chars: str, delay: float | None = 0.01) -> None: r""" Simulate typing a sequence of characters on the virtual keyboard, @@ -759,70 +996,6 @@ class Machine: with self.nested(f"waiting for file '{filename}'"): retry(check_file, timeout) - def wait_for_open_port( - self, port: int, addr: str = "localhost", timeout: int = 900 - ) -> None: - """ - Wait until a process is listening on the given TCP port and IP address - (default `localhost`). - """ - - def port_is_open(_last_try: bool) -> bool: - status, _ = self.execute(f"nc -z {addr} {port}") - return status == 0 - - with self.nested(f"waiting for TCP port {port} on {addr}"): - retry(port_is_open, timeout) - - def wait_for_open_unix_socket( - self, addr: str, is_datagram: bool = False, timeout: int = 900 - ) -> None: - """ - Wait until a process is listening on the given UNIX-domain socket - (default to a UNIX-domain stream socket). - """ - - nc_flags = [ - "-z", - "-uU" if is_datagram else "-U", - ] - - def socket_is_open(_last_try: bool) -> bool: - status, _ = self.execute(f"nc {' '.join(nc_flags)} {addr}") - return status == 0 - - with self.nested( - f"waiting for UNIX-domain {'datagram' if is_datagram else 'stream'} on '{addr}'" - ): - retry(socket_is_open, timeout) - - def wait_for_closed_port( - self, port: int, addr: str = "localhost", timeout: int = 900 - ) -> None: - """ - Wait until nobody is listening on the given TCP port and IP address - (default `localhost`). - """ - - def port_is_closed(_last_try: bool) -> bool: - status, _ = self.execute(f"nc -z {addr} {port}") - return status != 0 - - with self.nested(f"waiting for TCP port {port} on {addr} to be closed"): - retry(port_is_closed, timeout) - - def start_job(self, jobname: str, user: str | None = None) -> tuple[int, str]: - """ - Start systemd service. - """ - return self.systemctl(f"start {jobname}", user) - - def stop_job(self, jobname: str, user: str | None = None) -> tuple[int, str]: - """ - Stop systemd service. - """ - return self.systemctl(f"stop {jobname}", user) - def connect(self) -> None: """ Wait for a connection to the guest root shell @@ -902,78 +1075,6 @@ class Machine: f"Cannot convert screenshot (pnmtopng returned code {ret.returncode})" ) - def copy_from_host_via_shell(self, source: str, target: str) -> None: - """Copy a file from the host into the guest by piping it over the - shell into the destination file. Works without host-guest shared folder. - Prefer copy_from_host for whenever possible. - """ - with open(source, "rb") as fh: - content_b64 = base64.b64encode(fh.read()).decode() - self.succeed( - f"mkdir -p $(dirname {target})", - f"echo -n {content_b64} | base64 -d > {target}", - ) - - def copy_from_host(self, source: str, target: str) -> None: - """ - Copies a file from host to machine, e.g., - `copy_from_host("myfile", "/etc/my/important/file")`. - - The first argument is the file on the host. Note that the "host" refers - to the environment in which the test driver runs, which is typically the - Nix build sandbox. - - The second argument is the location of the file on the machine that will - be written to. - - The file is copied via the `shared_dir` directory which is shared among - all the VMs (using a temporary directory). - The access rights bits will mimic the ones from the host file and - user:group will be root:root. - """ - host_src = Path(source) - vm_target = Path(target) - with tempfile.TemporaryDirectory(dir=self.shared_dir) as shared_td: - shared_temp = Path(shared_td) - host_intermediate = shared_temp / host_src.name - vm_shared_temp = Path("/tmp/shared") / shared_temp.name - vm_intermediate = vm_shared_temp / host_src.name - - self.succeed(make_command(["mkdir", "-p", vm_shared_temp])) - if host_src.is_dir(): - shutil.copytree(host_src, host_intermediate) - else: - shutil.copy(host_src, host_intermediate) - self.succeed(make_command(["mkdir", "-p", vm_target.parent])) - self.succeed(make_command(["cp", "-r", vm_intermediate, vm_target])) - - def copy_from_vm(self, source: str, target_dir: str = "") -> None: - """Copy a file from the VM (specified by an in-VM source path) to a path - relative to `$out`. The file is copied via the `shared_dir` shared among - all the VMs (using a temporary directory). - """ - # Compute the source, target, and intermediate shared file names - vm_src = Path(source) - with tempfile.TemporaryDirectory(dir=self.shared_dir) as shared_td: - shared_temp = Path(shared_td) - vm_shared_temp = Path("/tmp/shared") / shared_temp.name - vm_intermediate = vm_shared_temp / vm_src.name - intermediate = shared_temp / vm_src.name - # 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 = 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(): - shutil.copytree(intermediate, abs_target) - else: - shutil.copy(intermediate, abs_target) - - def dump_tty_contents(self, tty: str) -> None: - """Debugging: Dump the contents of the TTY""" - self.execute(f"fold -w 80 /dev/vcs{tty} | systemd-cat") - def get_screen_text_variants(self) -> list[str]: """ Return a list of different interpretations of what is currently @@ -1154,11 +1255,6 @@ class Machine: self.log(f"QEMU running (pid {self.pid})") - def cleanup_statedir(self) -> None: - shutil.rmtree(self.state_dir) - self.logger.log(f"deleting VM state directory {self.state_dir}") - self.logger.log("if you want to keep the VM state, pass --keep-vm-state") - def shutdown(self) -> None: """ Shut down the machine, waiting for the VM to exit. @@ -1234,10 +1330,6 @@ class Machine: with self.nested("waiting for a window to appear"): retry(window_is_visible, timeout) - def sleep(self, secs: int) -> None: - # We want to sleep in *guest* time, not *host* time. - self.succeed(f"sleep {secs}") - def forward_port(self, host_port: int = 8080, guest_port: int = 80) -> None: """ Forward a TCP port on the host to a TCP port on the guest. @@ -1264,7 +1356,7 @@ class Machine: def release(self) -> None: if self.pid is None: return - self.logger.info(f"kill machine (pid {self.pid})") + self.logger.info(f"kill QemuMachine (pid {self.pid})") assert self.process assert self.shell assert self.monitor @@ -1278,10 +1370,6 @@ class Machine: if self.qmp_client: self.qmp_client.close() - def run_callbacks(self) -> None: - for callback in self.callbacks: - callback() - def switch_root(self) -> None: """ Transition from stage 1 to stage 2. This requires the diff --git a/nixos/lib/test-script-prepend.py b/nixos/lib/test-script-prepend.py index 6be20270c6cc..84a689f93d16 100644 --- a/nixos/lib/test-script-prepend.py +++ b/nixos/lib/test-script-prepend.py @@ -4,7 +4,7 @@ from test_driver.debug import DebugAbstract from test_driver.driver import Driver from test_driver.vlan import VLan -from test_driver.machine import Machine +from test_driver.machine import BaseMachine, QemuMachine from test_driver.logger import AbstractLogger from typing import Callable, Iterator, ContextManager, Optional, List, Dict, Any, Union from typing_extensions import Protocol @@ -35,7 +35,7 @@ class CreateMachineProtocol(Protocol): *, name: Optional[str] = None, keep_vm_state: bool = False, - ) -> Machine: + ) -> BaseMachine: raise Exception("This is just type information for the Nix test driver") @@ -43,7 +43,7 @@ start_all: Callable[[], None] subtest: Callable[[str], ContextManager[None]] retry: RetryProtocol test_script: Callable[[], None] -machines: List[Machine] +machines: List[BaseMachine] vlans: List[VLan] driver: Driver log: AbstractLogger diff --git a/nixos/lib/testing/driver.nix b/nixos/lib/testing/driver.nix index 5845ebe2695a..5623c63c4dc5 100644 --- a/nixos/lib/testing/driver.nix +++ b/nixos/lib/testing/driver.nix @@ -21,12 +21,6 @@ let ) (lib.attrValues config.nodes); vms = map (m: m.system.build.vm) (lib.attrValues config.nodes); - nodeHostNames = - let - nodesList = map (c: c.system.name) (lib.attrValues config.nodes); - in - nodesList ++ lib.optional (lib.length nodesList == 1 && !lib.elem "machine" nodesList) "machine"; - pythonizeName = name: let @@ -38,8 +32,17 @@ let uniqueVlans = lib.unique (builtins.concatLists vlans); vlanNames = map (i: "vlan${toString i}: VLan;") uniqueVlans; - pythonizedNames = map pythonizeName nodeHostNames; - machineNames = map (name: "${name}: Machine;") pythonizedNames; + + vmMachineNames = map (c: c.system.name) (lib.attrValues config.nodes); + + theOnlyMachine = + let + exactlyOneMachine = lib.length (lib.attrValues config.nodes) == 1; + in + lib.optional (exactlyOneMachine && !lib.elem "machine" vmMachineNames) "machine"; + + pythonizedVmNames = map pythonizeName (vmMachineNames ++ theOnlyMachine); + vmMachineTypeHints = map (name: "${name}: QemuMachine;") pythonizedVmNames; withChecks = lib.warnIf config.skipLint "Linting is disabled"; @@ -62,12 +65,13 @@ let '' mkdir -p $out/bin - vmStartScripts=($(for i in ${toString vms}; do echo $i/bin/run-*-vm; done)) + vmNames=(${lib.escapeShellArgs vmMachineNames}) + vmStartScripts=(${lib.escapeShellArgs (map lib.getExe vms)}) ${lib.optionalString (!config.skipTypeCheck) '' # prepend type hints so the test script can be type checked with mypy cat "${../test-script-prepend.py}" >> testScriptWithTypes - echo "${toString machineNames}" >> testScriptWithTypes + echo "${toString vmMachineTypeHints}" >> testScriptWithTypes echo "${toString vlanNames}" >> testScriptWithTypes echo -n "$testScript" >> testScriptWithTypes @@ -90,7 +94,7 @@ let echo "See https://nixos.org/manual/nixos/stable/#test-opt-skipLint" PYFLAKES_BUILTINS="$( - echo -n ${lib.escapeShellArg (lib.concatStringsSep "," pythonizedNames)}, + echo -n ${lib.escapeShellArg (lib.concatStringsSep "," pythonizedVmNames)}, cat ${lib.escapeShellArg "driver-symbols"} )" ${hostPkgs.python3Packages.pyflakes}/bin/pyflakes $out/test-script ''} @@ -98,7 +102,8 @@ let # set defaults through environment # see: ./test-driver/test-driver.py argparse implementation wrapProgram $out/bin/nixos-test-driver \ - --set startScripts "''${vmStartScripts[*]}" \ + --set vmStartScripts "''${vmStartScripts[*]}" \ + --set vmNames "''${vmNames[*]}" \ --set testScript "$out/test-script" \ --set globalTimeout "${toString config.globalTimeout}" \ --set vlans '${toString vlans}' \