nixos/test-driver: refactor Machine to BaseMachine and QemuMachine

This prepares the driver for non-QEMU backends by abstracting
the machine logic.

Co-authored-by: Jeremy Fleischman <jeremyfleischman@gmail.com>
This commit is contained in:
Kierán Meinhardt
2026-01-28 11:28:32 +01:00
co-authored by Jeremy Fleischman
parent 197287daf4
commit 799cafcc23
6 changed files with 538 additions and 421 deletions
@@ -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 = [
@@ -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()))
+28 -25
View File
@@ -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 <nixos/lib/testing/network.nix>.
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,
@@ -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"<Machine '{self.name}'>"
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<n>"""
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<n>"""
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
+3 -3
View File
@@ -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
+17 -12
View File
@@ -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}' \