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