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("{0} is not a directory".format(path))
45 if not os.access(path, os.W_OK):
46 raise argparse.ArgumentTypeError(
47 "{0} is not a writeable directory".format(path)
48 )
49 return path
50
51
52def main() -> None:
53 arg_parser = argparse.ArgumentParser(prog="nixos-test-driver")
54 arg_parser.add_argument(
55 "-K",
56 "--keep-vm-state",
57 help="re-use a VM state coming from a previous run",
58 action="store_true",
59 )
60 arg_parser.add_argument(
61 "-I",
62 "--interactive",
63 help="drop into a python repl and run the tests interactively",
64 action=argparse.BooleanOptionalAction,
65 )
66 arg_parser.add_argument(
67 "--start-scripts",
68 metavar="START-SCRIPT",
69 action=EnvDefault,
70 envvar="startScripts",
71 nargs="*",
72 help="start scripts for participating virtual machines",
73 )
74 arg_parser.add_argument(
75 "--vlans",
76 metavar="VLAN",
77 action=EnvDefault,
78 envvar="vlans",
79 nargs="*",
80 help="vlans to span by the driver",
81 )
82 arg_parser.add_argument(
83 "-o",
84 "--output_directory",
85 help="""The path to the directory where outputs copied from the VM will be placed.
86 By e.g. Machine.copy_from_vm or Machine.screenshot""",
87 default=Path.cwd(),
88 type=writeable_dir,
89 )
90 arg_parser.add_argument(
91 "testscript",
92 action=EnvDefault,
93 envvar="testScript",
94 help="the test script to run",
95 type=Path,
96 )
97
98 args = arg_parser.parse_args()
99
100 if not args.keep_vm_state:
101 rootlog.info("Machine state will be reset. To keep it, pass --keep-vm-state")
102
103 with Driver(
104 args.start_scripts,
105 args.vlans,
106 args.testscript.read_text(),
107 args.output_directory.resolve(),
108 args.keep_vm_state,
109 ) as driver:
110 if args.interactive:
111 ptpython.repl.embed(driver.test_symbols(), {})
112 else:
113 tic = time.time()
114 driver.run_tests()
115 toc = time.time()
116 rootlog.info(f"test script finished in {(toc-tic):.2f}s")
117
118
119def generate_driver_symbols() -> None:
120 """
121 This generates a file with symbols of the test-driver code that can be used
122 in user's test scripts. That list is then used by pyflakes to lint those
123 scripts.
124 """
125 d = Driver([], [], "", Path())
126 test_symbols = d.test_symbols()
127 with open("driver-symbols", "w") as fp:
128 fp.write(",".join(test_symbols.keys()))