pyrox.dev
1from pathlib import Path
2import argparse
3import ptpython.repl
4import os
5import time
6
7from test_driver.logger import rootlog
8from test_driver.driver import Driver
9
10
11class EnvDefault(argparse.Action):
12 """An argpars Action that takes values from the specified
13 environment variable as the flags default value.
14 """
15
16 def __init__(self, envvar, required=False, default=None, nargs=None, **kwargs): # type: ignore
17 if not default and envvar:
18 if envvar in os.environ:
19 if nargs is not None and (nargs.isdigit() or nargs in ["*", "+"]):
20 default = os.environ[envvar].split()
21 else:
22 default = os.environ[envvar]
23 kwargs["help"] = (
24 kwargs["help"] + f" (default from environment: {default})"
25 )
26 if required and default:
27 required = False
28 super(EnvDefault, self).__init__(
29 default=default, required=required, nargs=nargs, **kwargs
30 )
31
32 def __call__(self, parser, namespace, values, option_string=None): # type: ignore
33 setattr(namespace, self.dest, values)
34
35
36def writeable_dir(arg: str) -> Path:
37 """Raises an ArgumentTypeError if the given argument isn't a writeable directory
38 Note: We want to fail as early as possible if a directory isn't writeable,
39 since an executed nixos-test could fail (very late) because of the test-driver
40 writing in a directory without proper permissions.
41 """
42 path = Path(arg)
43 if not path.is_dir():
44 raise argparse.ArgumentTypeError(f"{path} is not a directory")
45 if not os.access(path, os.W_OK):
46 raise argparse.ArgumentTypeError(f"{path} is not a writeable directory")
47 return path
48
49
50def main() -> None:
51 arg_parser = argparse.ArgumentParser(prog="nixos-test-driver")
52 arg_parser.add_argument(
53 "-K",
54 "--keep-vm-state",
55 help="re-use a VM state coming from a previous run",
56 action="store_true",
57 )
58 arg_parser.add_argument(
59 "-I",
60 "--interactive",
61 help="drop into a python repl and run the tests interactively",
62 action=argparse.BooleanOptionalAction,
63 )
64 arg_parser.add_argument(
65 "--start-scripts",
66 metavar="START-SCRIPT",
67 action=EnvDefault,
68 envvar="startScripts",
69 nargs="*",
70 help="start scripts for participating virtual machines",
71 )
72 arg_parser.add_argument(
73 "--vlans",
74 metavar="VLAN",
75 action=EnvDefault,
76 envvar="vlans",
77 nargs="*",
78 help="vlans to span by the driver",
79 )
80 arg_parser.add_argument(
81 "-o",
82 "--output_directory",
83 help="""The path to the directory where outputs copied from the VM will be placed.
84 By e.g. Machine.copy_from_vm or Machine.screenshot""",
85 default=Path.cwd(),
86 type=writeable_dir,
87 )
88 arg_parser.add_argument(
89 "testscript",
90 action=EnvDefault,
91 envvar="testScript",
92 help="the test script to run",
93 type=Path,
94 )
95
96 args = arg_parser.parse_args()
97
98 if not args.keep_vm_state:
99 rootlog.info("Machine state will be reset. To keep it, pass --keep-vm-state")
100
101 with Driver(
102 args.start_scripts,
103 args.vlans,
104 args.testscript.read_text(),
105 args.output_directory.resolve(),
106 args.keep_vm_state,
107 ) as driver:
108 if args.interactive:
109 ptpython.repl.embed(driver.test_symbols(), {})
110 else:
111 tic = time.time()
112 driver.run_tests()
113 toc = time.time()
114 rootlog.info(f"test script finished in {(toc-tic):.2f}s")
115
116
117def generate_driver_symbols() -> None:
118 """
119 This generates a file with symbols of the test-driver code that can be used
120 in user's test scripts. That list is then used by pyflakes to lint those
121 scripts.
122 """
123 d = Driver([], [], "", Path())
124 test_symbols = d.test_symbols()
125 with open("driver-symbols", "w") as fp:
126 fp.write(",".join(test_symbols.keys()))