@@ -1,19 +1,11 @@
1−#! /somewhere/python3
2−from contextlib import contextmanager, _GeneratorContextManager
3−from queue import Queue, Empty
4−from typing import Tuple, Any, Callable, Dict, Iterator, Optional, List, Iterable
5−from xml.sax.saxutils import XMLGenerator
6−from colorama import Style
1+from contextlib import _GeneratorContextManager
72from pathlib import Path
8−import queue
9−import io
10−import threading
11−import argparse
3+from queue import Queue
4+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
125import base64
13−import codecs
6+import io
147import os
15−import ptpython.repl
16−import pty
8+import queue
179import re
1810import shlex
1911import shutil
@@ -21,8 +13,10 @@ import socket
2113import subprocess
2214import sys
2315import tempfile
16+import threading
2417import time
25−import unicodedata
18+
19+from test_driver.logger import rootlog
2620
2721CHAR_TO_KEY = {
2822 "A": "shift-a",
@@ -88,115 +82,10 @@ CHAR_TO_KEY = {
8882}
8983
9084
91−class Logger:
92− def __init__(self) -> None:
93− self.logfile = os.environ.get("LOGFILE", "/dev/null")
94− self.logfile_handle = codecs.open(self.logfile, "wb")
95− self.xml = XMLGenerator(self.logfile_handle, encoding="utf-8")
96− self.queue: "Queue[Dict[str, str]]" = Queue()
97−
98− self.xml.startDocument()
99− self.xml.startElement("logfile", attrs={})
100−
101− self._print_serial_logs = True
102−
103− @staticmethod
104− def _eprint(*args: object, **kwargs: Any) -> None:
105− print(*args, file=sys.stderr, **kwargs)
106−
107− def close(self) -> None:
108− self.xml.endElement("logfile")
109− self.xml.endDocument()
110− self.logfile_handle.close()
111−
112− def sanitise(self, message: str) -> str:
113− return "".join(ch for ch in message if unicodedata.category(ch)[0] != "C")
114−
115− def maybe_prefix(self, message: str, attributes: Dict[str, str]) -> str:
116− if "machine" in attributes:
117− return "{}: {}".format(attributes["machine"], message)
118− return message
119−
120− def log_line(self, message: str, attributes: Dict[str, str]) -> None:
121− self.xml.startElement("line", attributes)
122− self.xml.characters(message)
123− self.xml.endElement("line")
124−
125− def info(self, *args, **kwargs) -> None: # type: ignore
126− self.log(*args, **kwargs)
127−
128− def warning(self, *args, **kwargs) -> None: # type: ignore
129− self.log(*args, **kwargs)
130−
131− def error(self, *args, **kwargs) -> None: # type: ignore
132− self.log(*args, **kwargs)
133− sys.exit(1)
134−
135− def log(self, message: str, attributes: Dict[str, str] = {}) -> None:
136− self._eprint(self.maybe_prefix(message, attributes))
137− self.drain_log_queue()
138− self.log_line(message, attributes)
139−
140− def log_serial(self, message: str, machine: str) -> None:
141− self.enqueue({"msg": message, "machine": machine, "type": "serial"})
142− if self._print_serial_logs:
143− self._eprint(
144− Style.DIM + "{} # {}".format(machine, message) + Style.RESET_ALL
145− )
146−
147− def enqueue(self, item: Dict[str, str]) -> None:
148− self.queue.put(item)
149−
150− def drain_log_queue(self) -> None:
151− try:
152− while True:
153− item = self.queue.get_nowait()
154− msg = self.sanitise(item["msg"])
155− del item["msg"]
156− self.log_line(msg, item)
157− except Empty:
158− pass
159−
160− @contextmanager
161− def nested(self, message: str, attributes: Dict[str, str] = {}) -> Iterator[None]:
162− self._eprint(self.maybe_prefix(message, attributes))
163−
164− self.xml.startElement("nest", attrs={})
165− self.xml.startElement("head", attributes)
166− self.xml.characters(message)
167− self.xml.endElement("head")
168−
169− tic = time.time()
170− self.drain_log_queue()
171− yield
172− self.drain_log_queue()
173− toc = time.time()
174− self.log("(finished: {}, in {:.2f} seconds)".format(message, toc - tic))
175−
176− self.xml.endElement("nest")
177−
178−
179−rootlog = Logger()
180−
181−
18285def make_command(args: list) -> str:
18386 return " ".join(map(shlex.quote, (map(str, args))))
18487
18588
186−def retry(fn: Callable, timeout: int = 900) -> None:
187− """Call the given function repeatedly, with 1 second intervals,
188− until it returns True or a timeout is reached.
189− """
190−
191− for _ in range(timeout):
192− if fn(False):
193− return
194− time.sleep(1)
195−
196− if not fn(True):
197− raise Exception(f"action timed out after {timeout} seconds")
198−
199−
20089def _perform_ocr_on_screenshot(
20190 screenshot_path: str, model_ids: Iterable[int]
20291) -> List[str]:
@@ -228,6 +117,20 @@ def _perform_ocr_on_screenshot(
228117 return model_results
229118
230119
120+def retry(fn: Callable, timeout: int = 900) -> None:
121+ """Call the given function repeatedly, with 1 second intervals,
122+ until it returns True or a timeout is reached.
123+ """
124+
125+ for _ in range(timeout):
126+ if fn(False):
127+ return
128+ time.sleep(1)
129+
130+ if not fn(True):
131+ raise Exception(f"action timed out after {timeout} seconds")
132+
133+
231134class StartCommand:
232135 """The Base Start Command knows how to append the necesary
233136 runtime qemu options as determined by a particular test driver
@@ -1066,286 +969,3 @@ class Machine:
1066969 self.shell.close()
1067970 self.monitor.close()
1068971 self.serial_thread.join()
1069−
1070−
1071−class VLan:
1072− """This class handles a VLAN that the run-vm scripts identify via its
1073− number handles. The network's lifetime equals the object's lifetime.
1074− """
1075−
1076− nr: int
1077− socket_dir: Path
1078−
1079− process: subprocess.Popen
1080− pid: int
1081− fd: io.TextIOBase
1082−
1083− def __repr__(self) -> str:
1084− return f"<Vlan Nr. {self.nr}>"
1085−
1086− def __init__(self, nr: int, tmp_dir: Path):
1087− self.nr = nr
1088− self.socket_dir = tmp_dir / f"vde{self.nr}.ctl"
1089−
1090− # TODO: don't side-effect environment here
1091− os.environ[f"QEMU_VDE_SOCKET_{self.nr}"] = str(self.socket_dir)
1092−
1093− rootlog.info("start vlan")
1094− pty_master, pty_slave = pty.openpty()
1095−
1096− self.process = subprocess.Popen(
1097− ["vde_switch", "-s", self.socket_dir, "--dirmode", "0700"],
1098− stdin=pty_slave,
1099− stdout=subprocess.PIPE,
1100− stderr=subprocess.PIPE,
1101− shell=False,
1102− )
1103− self.pid = self.process.pid
1104− self.fd = os.fdopen(pty_master, "w")
1105− self.fd.write("version\n")
1106−
1107− # TODO: perl version checks if this can be read from
1108− # an if not, dies. we could hang here forever. Fix it.
1109− assert self.process.stdout is not None
1110− self.process.stdout.readline()
1111− if not (self.socket_dir / "ctl").exists():
1112− rootlog.error("cannot start vde_switch")
1113−
1114− rootlog.info(f"running vlan (pid {self.pid})")
1115−
1116− def __del__(self) -> None:
1117− rootlog.info(f"kill vlan (pid {self.pid})")
1118− self.fd.close()
1119− self.process.terminate()
1120−
1121−
1122−class Driver:
1123− """A handle to the driver that sets up the environment
1124− and runs the tests"""
1125−
1126− tests: str
1127− vlans: List[VLan]
1128− machines: List[Machine]
1129−
1130− def __init__(
1131− self,
1132− start_scripts: List[str],
1133− vlans: List[int],
1134− tests: str,
1135− keep_vm_state: bool = False,
1136− ):
1137− self.tests = tests
1138−
1139− tmp_dir = Path(os.environ.get("TMPDIR", tempfile.gettempdir()))
1140− tmp_dir.mkdir(mode=0o700, exist_ok=True)
1141−
1142− with rootlog.nested("start all VLans"):
1143− self.vlans = [VLan(nr, tmp_dir) for nr in vlans]
1144−
1145− def cmd(scripts: List[str]) -> Iterator[NixStartScript]:
1146− for s in scripts:
1147− yield NixStartScript(s)
1148−
1149− self.machines = [
1150− Machine(
1151− start_command=cmd,
1152− keep_vm_state=keep_vm_state,
1153− name=cmd.machine_name,
1154− tmp_dir=tmp_dir,
1155− )
1156− for cmd in cmd(start_scripts)
1157− ]
1158−
1159− def __enter__(self) -> "Driver":
1160− return self
1161−
1162− def __exit__(self, *_: Any) -> None:
1163− with rootlog.nested("cleanup"):
1164− for machine in self.machines:
1165− machine.release()
1166−
1167− def subtest(self, name: str) -> Iterator[None]:
1168− """Group logs under a given test name"""
1169− with rootlog.nested(name):
1170− try:
1171− yield
1172− return True
1173− except Exception as e:
1174− rootlog.error(f'Test "{name}" failed with error: "{e}"')
1175− raise e
1176−
1177− def test_symbols(self) -> Dict[str, Any]:
1178− @contextmanager
1179− def subtest(name: str) -> Iterator[None]:
1180− return self.subtest(name)
1181−
1182− general_symbols = dict(
1183− start_all=self.start_all,
1184− test_script=self.test_script,
1185− machines=self.machines,
1186− vlans=self.vlans,
1187− driver=self,
1188− log=rootlog,
1189− os=os,
1190− create_machine=self.create_machine,
1191− subtest=subtest,
1192− run_tests=self.run_tests,
1193− join_all=self.join_all,
1194− retry=retry,
1195− serial_stdout_off=self.serial_stdout_off,
1196− serial_stdout_on=self.serial_stdout_on,
1197− Machine=Machine, # for typing
1198− )
1199− machine_symbols = {m.name: m for m in self.machines}
1200− # If there's exactly one machine, make it available under the name
1201− # "machine", even if it's not called that.
1202− if len(self.machines) == 1:
1203− (machine_symbols["machine"],) = self.machines
1204− vlan_symbols = {
1205− f"vlan{v.nr}": self.vlans[idx] for idx, v in enumerate(self.vlans)
1206− }
1207− print(
1208− "additionally exposed symbols:\n "
1209− + ", ".join(map(lambda m: m.name, self.machines))
1210− + ",\n "
1211− + ", ".join(map(lambda v: f"vlan{v.nr}", self.vlans))
1212− + ",\n "
1213− + ", ".join(list(general_symbols.keys()))
1214− )
1215− return {**general_symbols, **machine_symbols, **vlan_symbols}
1216−
1217− def test_script(self) -> None:
1218− """Run the test script"""
1219− with rootlog.nested("run the VM test script"):
1220− symbols = self.test_symbols() # call eagerly
1221− exec(self.tests, symbols, None)
1222−
1223− def run_tests(self) -> None:
1224− """Run the test script (for non-interactive test runs)"""
1225− self.test_script()
1226− # TODO: Collect coverage data
1227− for machine in self.machines:
1228− if machine.is_up():
1229− machine.execute("sync")
1230−
1231− def start_all(self) -> None:
1232− """Start all machines"""
1233− with rootlog.nested("start all VMs"):
1234− for machine in self.machines:
1235− machine.start()
1236−
1237− def join_all(self) -> None:
1238− """Wait for all machines to shut down"""
1239− with rootlog.nested("wait for all VMs to finish"):
1240− for machine in self.machines:
1241− machine.wait_for_shutdown()
1242−
1243− def create_machine(self, args: Dict[str, Any]) -> Machine:
1244− rootlog.warning(
1245− "Using legacy create_machine(), please instantiate the"
1246− "Machine class directly, instead"
1247− )
1248− tmp_dir = Path(os.environ.get("TMPDIR", tempfile.gettempdir()))
1249− tmp_dir.mkdir(mode=0o700, exist_ok=True)
1250−
1251− if args.get("startCommand"):
1252− start_command: str = args.get("startCommand", "")
1253− cmd = NixStartScript(start_command)
1254− name = args.get("name", cmd.machine_name)
1255− else:
1256− cmd = Machine.create_startcommand(args) # type: ignore
1257− name = args.get("name", "machine")
1258−
1259− return Machine(
1260− tmp_dir=tmp_dir,
1261− start_command=cmd,
1262− name=name,
1263− keep_vm_state=args.get("keep_vm_state", False),
1264− allow_reboot=args.get("allow_reboot", False),
1265− )
1266−
1267− def serial_stdout_on(self) -> None:
1268− rootlog._print_serial_logs = True
1269−
1270− def serial_stdout_off(self) -> None:
1271− rootlog._print_serial_logs = False
1272−
1273−
1274−class EnvDefault(argparse.Action):
1275− """An argpars Action that takes values from the specified
1276− environment variable as the flags default value.
1277− """
1278−
1279− def __init__(self, envvar, required=False, default=None, nargs=None, **kwargs): # type: ignore
1280− if not default and envvar:
1281− if envvar in os.environ:
1282− if nargs is not None and (nargs.isdigit() or nargs in ["*", "+"]):
1283− default = os.environ[envvar].split()
1284− else:
1285− default = os.environ[envvar]
1286− kwargs["help"] = (
1287− kwargs["help"] + f" (default from environment: {default})"
1288− )
1289− if required and default:
1290− required = False
1291− super(EnvDefault, self).__init__(
1292− default=default, required=required, nargs=nargs, **kwargs
1293− )
1294−
1295− def __call__(self, parser, namespace, values, option_string=None): # type: ignore
1296− setattr(namespace, self.dest, values)
1297−
1298−
1299−if __name__ == "__main__":
1300− arg_parser = argparse.ArgumentParser(prog="nixos-test-driver")
1301− arg_parser.add_argument(
1302− "-K",
1303− "--keep-vm-state",
1304− help="re-use a VM state coming from a previous run",
1305− action="store_true",
1306− )
1307− arg_parser.add_argument(
1308− "-I",
1309− "--interactive",
1310− help="drop into a python repl and run the tests interactively",
1311− action="store_true",
1312− )
1313− arg_parser.add_argument(
1314− "--start-scripts",
1315− metavar="START-SCRIPT",
1316− action=EnvDefault,
1317− envvar="startScripts",
1318− nargs="*",
1319− help="start scripts for participating virtual machines",
1320− )
1321− arg_parser.add_argument(
1322− "--vlans",
1323− metavar="VLAN",
1324− action=EnvDefault,
1325− envvar="vlans",
1326− nargs="*",
1327− help="vlans to span by the driver",
1328− )
1329− arg_parser.add_argument(
1330− "testscript",
1331− action=EnvDefault,
1332− envvar="testScript",
1333− help="the test script to run",
1334− type=Path,
1335− )
1336−
1337− args = arg_parser.parse_args()
1338−
1339− if not args.keep_vm_state:
1340− rootlog.info("Machine state will be reset. To keep it, pass --keep-vm-state")
1341−
1342− with Driver(
1343− args.start_scripts, args.vlans, args.testscript.read_text(), args.keep_vm_state
1344− ) as driver:
1345− if args.interactive:
1346− ptpython.repl.embed(driver.test_symbols(), {})
1347− else:
1348− tic = time.time()
1349− driver.run_tests()
1350− toc = time.time()
1351− rootlog.info(f"test script finished in {(toc-tic):.2f}s")