···
···
def __init__(self) -> None:
···
self.xml.startElement("logfile", attrs={})
self._print_serial_logs = True
+
def _eprint(*args: object, **kwargs: Any) -> None:
+
print(*args, file=sys.stderr, **kwargs)
self.xml.endElement("logfile")
···
self.xml.characters(message)
self.xml.endElement("line")
+
def info(self, *args, **kwargs) -> None: # type: ignore
+
self.log(*args, **kwargs)
+
def warning(self, *args, **kwargs) -> None: # type: ignore
+
self.log(*args, **kwargs)
+
def error(self, *args, **kwargs) -> None: # type: ignore
+
self.log(*args, **kwargs)
def log(self, message: str, attributes: Dict[str, str] = {}) -> None:
+
self._eprint(self.maybe_prefix(message, attributes))
self.log_line(message, attributes)
def log_serial(self, message: str, machine: str) -> None:
self.enqueue({"msg": message, "machine": machine, "type": "serial"})
if self._print_serial_logs:
+
Style.DIM + "{} # {}".format(machine, message) + Style.RESET_ALL
def enqueue(self, item: Dict[str, str]) -> None:
···
def nested(self, message: str, attributes: Dict[str, str] = {}) -> Iterator[None]:
+
self._eprint(self.maybe_prefix(message, attributes))
self.xml.startElement("nest", attrs={})
self.xml.startElement("head", attributes)
···
self.xml.endElement("nest")
+
def make_command(args: list) -> str:
+
return " ".join(map(shlex.quote, (map(str, args))))
+
def retry(fn: Callable, timeout: int = 900) -> None:
+
"""Call the given function repeatedly, with 1 second intervals,
+
until it returns True or a timeout is reached.
+
for _ in range(timeout):
+
raise Exception(f"action timed out after {timeout} seconds")
def _perform_ocr_on_screenshot(
screenshot_path: str, model_ids: Iterable[int]
···
+
"""The Base Start Command knows how to append the necesary
+
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.
+
monitor_socket_path: pathlib.Path,
+
shell_socket_path: pathlib.Path,
+
allow_reboot: bool = False, # TODO: unused, legacy?
+
display_available = any(x in os.environ for x in ["DISPLAY", "WAYLAND_DISPLAY"])
+
if not display_available:
+
display_opts += " -nographic"
+
" -device virtio-serial"
+
" -device virtconsole,chardev=shell"
+
" -device virtio-rng-pci"
+
# TODO: qemu script already catpures this env variable, legacy?
+
qemu_opts += " " + os.environ.get("QEMU_OPTS", "")
+
f" -monitor unix:{monitor_socket_path}"
+
f" -chardev socket,id=shell,path={shell_socket_path}"
+
state_dir: pathlib.Path,
+
shared_dir: pathlib.Path,
+
# We make a copy to not update the current environment
+
"TMPDIR": str(state_dir),
+
"SHARED_DIR": str(shared_dir),
+
state_dir: pathlib.Path,
+
shared_dir: pathlib.Path,
+
monitor_socket_path: pathlib.Path,
+
shell_socket_path: pathlib.Path,
+
return subprocess.Popen(
+
self.cmd(monitor_socket_path, shell_socket_path),
+
stdin=subprocess.DEVNULL,
+
stdout=subprocess.PIPE,
+
stderr=subprocess.STDOUT,
+
env=self.build_environment(state_dir, shared_dir),
+
class NixStartScript(StartCommand):
+
"""A start script from nixos/modules/virtualiation/qemu-vm.nix
+
that also satisfies the requirement of the BaseStartCommand.
+
These Nix commands have the particular charactersitic 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):
+
def machine_name(self) -> str:
+
match = re.search("run-(.+)-vm$", self._cmd)
+
class LegacyStartCommand(StartCommand):
+
"""Used in some places to create an ad-hoc machine instead of
+
using nix test instrumentation + module system for that purpose.
+
netBackendArgs: Optional[str] = None,
+
netFrontendArgs: Optional[str] = None,
+
hda: Optional[Tuple[pathlib.Path, str]] = None,
+
cdrom: Optional[str] = None,
+
usb: Optional[str] = None,
+
bios: Optional[str] = None,
+
qemuFlags: Optional[str] = None,
+
self._cmd = "qemu-kvm -m 384"
+
net_backend = "-netdev user,id=net0"
+
net_frontend = "-device virtio-net-pci,netdev=net0"
+
if netBackendArgs is not None:
+
net_backend += "," + netBackendArgs
+
if netFrontendArgs is not None:
+
net_frontend += "," + netFrontendArgs
+
self._cmd += f" {net_backend} {net_frontend}"
+
hda_path = hda[0].resolve()
+
if hda_interface == "scsi":
+
f" -drive id=hda,file={hda_path},werror=report,if=none"
+
" -device scsi-hd,drive=hda"
+
hda_cmd += f" -drive file={hda_path},if={hda_interface},werror=report"
+
self._cmd += f" -cdrom {cdrom}"
# https://github.com/qemu/qemu/blob/master/docs/usb2.txt
+
f" -drive id=usbdisk,file={usb},if=none,readonly"
+
" -device usb-storage,drive=usbdisk "
+
self._cmd += f" -bios {bios}"
+
if qemuFlags is not None:
+
self._cmd += f" {qemuFlags}"
+
"""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."""
+
shared_dir: pathlib.Path
+
state_dir: pathlib.Path
+
monitor_path: pathlib.Path
+
shell_path: pathlib.Path
+
start_command: StartCommand
+
process: Optional[subprocess.Popen] = None
+
pid: Optional[int] = None
+
monitor: Optional[socket.socket] = None
+
shell: Optional[socket.socket] = None
+
connected: bool = False
+
# Store last serial console lines for use
+
# of wait_for_console_text
+
last_lines: Queue = Queue()
+
def __repr__(self) -> str:
+
return f"<Machine '{self.name}'>"
+
start_command: StartCommand,
+
keep_vm_state: bool = False,
+
allow_reboot: bool = False,
+
self.keep_vm_state = keep_vm_state
+
self.allow_reboot = allow_reboot
+
self.start_command = start_command
+
self.shared_dir = self.tmp_dir / "shared-xchg"
+
self.shared_dir.mkdir(mode=0o700, exist_ok=True)
+
self.state_dir = self.tmp_dir / f"vm-state-{self.name}"
+
self.monitor_path = self.state_dir / "monitor"
+
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)
+
def create_startcommand(args: Dict[str, str]) -> StartCommand:
+
"Using legacy create_startcommand(),"
+
"please use proper nix test vm instrumentation, instead"
+
"to generate the appropriate nixos test vm qemu startup script"
+
hda_arg: str = args.get("hda", "")
+
hda_arg_path: pathlib.Path = pathlib.Path(hda_arg)
+
hda = (hda_arg_path, args.get("hdaInterface", ""))
+
return LegacyStartCommand(
+
netBackendArgs=args.get("netBackendArgs"),
+
netFrontendArgs=args.get("netFrontendArgs"),
+
cdrom=args.get("cdrom"),
+
qemuFlags=args.get("qemuFlags"),
return self.booted and self.connected
def log(self, msg: str) -> None:
+
rootlog.log(msg, {"machine": self.name})
def log_serial(self, msg: str) -> None:
+
rootlog.log_serial(msg, self.name)
def nested(self, msg: str, attrs: Dict[str, str] = {}) -> _GeneratorContextManager:
my_attrs = {"machine": self.name}
+
return rootlog.nested(msg, my_attrs)
def wait_for_monitor_prompt(self) -> str:
assert self.monitor is not None
···
out_command = "( set -euo pipefail; {} ); echo '|!=EOF' $?\n".format(command)
self.shell.send(out_command.encode())
···
Should only be used during test development, not in the production test."""
self.log("Terminal is ready (there is no prompt):")
["socat", "READLINE", f"FD:{self.shell.fileno()}"],
pass_fds=[self.shell.fileno()],
···
with self.nested("waiting for the VM to power off"):
···
with self.nested("waiting for the VM to finish booting"):
···
+
def clear(path: pathlib.Path) -> pathlib.Path:
+
def create_socket(path: pathlib.Path) -> socket.socket:
s = socket.socket(family=socket.AF_UNIX, type=socket.SOCK_STREAM)
+
monitor_socket = create_socket(clear(self.monitor_path))
+
shell_socket = create_socket(clear(self.shell_path))
+
self.process = self.start_command.run(
+
self.monitor, _ = monitor_socket.accept()
+
self.shell, _ = 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
for _line in self.process.stdout:
# Ignore undecodable bytes that may occur in boot menus
line = _line.decode(errors="ignore").replace("\r", "").rstrip()
···
self.log("QEMU running (pid {})".format(self.pid))
def cleanup_statedir(self) -> None:
+
shutil.rmtree(self.state_dir)
+
rootlog.log(f"deleting VM state directory {self.state_dir}")
+
rootlog.log("if you want to keep the VM state, pass --keep-vm-state")
def shutdown(self) -> None:
self.shell.send("poweroff\n".encode())
···
"""Make the machine reachable."""
self.send_monitor_command("set_link virtio-net-pci.1 on")
+
def release(self) -> None:
+
rootlog.info(f"kill machine (pid {self.pid})")
+
self.process.terminate()
+
"""This class handles a VLAN that the run-vm scripts identify via its
+
number handles. The network's lifetime equals the object's lifetime.
+
socket_dir: pathlib.Path
+
process: subprocess.Popen
+
def __repr__(self) -> str:
+
return f"<Vlan Nr. {self.nr}>"
+
def __init__(self, nr: int, tmp_dir: pathlib.Path):
+
self.socket_dir = tmp_dir / f"vde{self.nr}.ctl"
+
# TODO: don't side-effect environment here
+
os.environ[f"QEMU_VDE_SOCKET_{self.nr}"] = str(self.socket_dir)
+
rootlog.info("start vlan")
+
pty_master, pty_slave = pty.openpty()
+
self.process = subprocess.Popen(
+
["vde_switch", "-s", self.socket_dir, "--dirmode", "0700"],
+
stdout=subprocess.PIPE,
+
stderr=subprocess.PIPE,
+
self.pid = self.process.pid
+
self.fd = os.fdopen(pty_master, "w")
+
self.fd.write("version\n")
+
# TODO: perl version checks if this can be read from
+
# an if not, dies. we could hang here forever. Fix it.
+
assert self.process.stdout is not None
+
self.process.stdout.readline()
+
if not (self.socket_dir / "ctl").exists():
+
rootlog.error("cannot start vde_switch")
+
rootlog.info(f"running vlan (pid {self.pid})")
+
def __del__(self) -> None:
+
rootlog.info(f"kill vlan (pid {self.pid})")
+
self.process.terminate()
+
"""A handle to the driver that sets up the environment
+
machines: List[Machine]
+
start_scripts: List[str],
+
keep_vm_state: bool = False,
+
tmp_dir = pathlib.Path(os.environ.get("TMPDIR", tempfile.gettempdir()))
+
tmp_dir.mkdir(mode=0o700, exist_ok=True)
+
with rootlog.nested("start all VLans"):
+
self.vlans = [VLan(nr, tmp_dir) for nr in vlans]
+
def cmd(scripts: List[str]) -> Iterator[NixStartScript]:
+
yield NixStartScript(s)
+
keep_vm_state=keep_vm_state,
+
for cmd in cmd(start_scripts)
+
def clean_up() -> None:
+
with rootlog.nested("clean up"):
+
for machine in self.machines:
+
def subtest(self, name: str) -> Iterator[None]:
+
"""Group logs under a given test name"""
+
with rootlog.nested(name):
+
rootlog.error(f'Test "{name}" failed with error:')
+
def test_symbols(self) -> Dict[str, Any]:
+
def subtest(name: str) -> Iterator[None]:
+
return self.subtest(name)
+
general_symbols = dict(
+
start_all=self.start_all,
+
test_script=self.test_script,
+
machines=self.machines,
+
create_machine=self.create_machine,
+
run_tests=self.run_tests,
+
join_all=self.join_all,
+
serial_stdout_off=self.serial_stdout_off,
+
serial_stdout_on=self.serial_stdout_on,
+
Machine=Machine, # for typing
+
m.name: self.machines[idx] for idx, m in enumerate(self.machines)
+
f"vlan{v.nr}": self.vlans[idx] for idx, v in enumerate(self.vlans)
+
"additionally exposed symbols:\n "
+
+ ", ".join(map(lambda m: m.name, self.machines))
+
+ ", ".join(map(lambda v: f"vlan{v.nr}", self.vlans))
+
+ ", ".join(list(general_symbols.keys()))
+
return {**general_symbols, **machine_symbols, **vlan_symbols}
+
def test_script(self) -> None:
+
"""Run the test script"""
+
with rootlog.nested("run the VM test script"):
+
symbols = self.test_symbols() # call eagerly
+
exec(self.tests, symbols, None)
+
def run_tests(self) -> None:
+
"""Run the test script (for non-interactive test runs)"""
# TODO: Collect coverage data
+
for machine in self.machines:
+
def start_all(self) -> None:
+
"""Start all machines"""
+
with rootlog.nested("start all VMs"):
+
for machine in self.machines:
+
def join_all(self) -> None:
+
"""Wait for all machines to shut down"""
+
with rootlog.nested("wait for all VMs to finish"):
+
for machine in self.machines:
+
machine.wait_for_shutdown()
+
def create_machine(self, args: Dict[str, Any]) -> Machine:
+
"Using legacy create_machine(), please instantiate the"
+
"Machine class directly, instead"
+
tmp_dir = pathlib.Path(os.environ.get("TMPDIR", tempfile.gettempdir()))
+
tmp_dir.mkdir(mode=0o700, exist_ok=True)
+
if args.get("startCommand"):
+
start_command: str = args.get("startCommand", "")
+
cmd = NixStartScript(start_command)
+
name = args.get("name", cmd.machine_name)
+
cmd = Machine.create_startcommand(args) # type: ignore
+
name = args.get("name", "machine")
+
keep_vm_state=args.get("keep_vm_state", False),
+
allow_reboot=args.get("allow_reboot", False),
+
def serial_stdout_on(self) -> None:
+
rootlog._print_serial_logs = True
+
def serial_stdout_off(self) -> None:
+
rootlog._print_serial_logs = False
class EnvDefault(argparse.Action):
···
setattr(namespace, self.dest, values)
if __name__ == "__main__":
arg_parser = argparse.ArgumentParser(prog="nixos-test-driver")
···
args = arg_parser.parse_args()
+
if not args.keep_vm_state:
+
rootlog.info("Machine state will be reset. To keep it, pass --keep-vm-state")
+
args.start_scripts, args.vlans, args.testscript.read_text(), args.keep_vm_state
+
ptpython.repl.embed(driver.test_symbols(), {})
+
rootlog.info(f"test script finished in {(toc-tic):.2f}s")
pyrox.dev