nixos/lib/test-driver/test_driver/__init__.py

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