Merge pull request #125979 from blaggacao/nixos-test-ref/03-normalse-the-python-entrypoint

nixos/test-driver: normalize the python entrypoint

nixos/doc/manual/development/running-nixos-tests-interactively.section.md

···
 
 ```ShellSession
 $ nix-build nixos/tests/login.nix -A driverInteractive
-$ ./result/bin/nixos-test-driver
+$ ./result/bin/nixos-test-driver --interactive
 starting VDE switch for network 1
 >
 ```
···
 you to inspect the state of the VMs after the test (e.g. to debug the
 test script).
 
-To just start and experiment with the VMs, run:
-
-```ShellSession
-$ nix-build nixos/tests/login.nix -A driverInteractive
-$ ./result/bin/nixos-run-vms
-```
-
-The script `nixos-run-vms` starts the virtual machines defined by test.
-
 You can re-use the VM states coming from a previous run by setting the
 `--keep-vm-state` flag.
 
 ```ShellSession
-$ ./result/bin/nixos-run-vms --keep-vm-state
+$ ./result/bin/nixos-test-driver --interactive --keep-vm-state
 ```
 
 The machine state is stored in the `$TMPDIR/vm-state-machinename`

nixos/doc/manual/from_md/development/running-nixos-tests-interactively.section.xml

···
   </para>
   <programlisting>
 $ nix-build nixos/tests/login.nix -A driverInteractive
-$ ./result/bin/nixos-test-driver
+$ ./result/bin/nixos-test-driver --interactive
 starting VDE switch for network 1
 &gt;
 </programlisting>
···
     the test (e.g. to debug the test script).
   </para>
   <para>
-    To just start and experiment with the VMs, run:
-  </para>
-  <programlisting>
-$ nix-build nixos/tests/login.nix -A driverInteractive
-$ ./result/bin/nixos-run-vms
-</programlisting>
-  <para>
-    The script <literal>nixos-run-vms</literal> starts the virtual
-    machines defined by test.
-  </para>
-  <para>
     You can re-use the VM states coming from a previous run by setting
     the <literal>--keep-vm-state</literal> flag.
   </para>
   <programlisting>
-$ ./result/bin/nixos-run-vms --keep-vm-state
+$ ./result/bin/nixos-test-driver --interactive --keep-vm-state
 </programlisting>
   <para>
     The machine state is stored in the

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

···
 import telnetlib
 import tempfile
 import time
-import traceback
 import unicodedata
 
 CHAR_TO_KEY = {
···
             machine.wait_for_shutdown()
 
 
-def test_script() -> None:
-    exec(os.environ["testScript"])
-
-
-def run_tests() -> None:
+def run_tests(interactive: bool = False) -> None:
     global machines
-    tests = os.environ.get("tests", None)
-    if tests is not None:
-        with log.nested("running the VM test script"):
-            try:
-                exec(tests, globals())
-            except Exception as e:
-                eprint("error: ")
-                traceback.print_exc()
-                sys.exit(1)
+    if interactive:
+        ptpython.repl.embed(globals(), locals())
     else:
-        ptpython.repl.embed(locals(), globals())
-
-    # TODO: Collect coverage data
-
-    for machine in machines:
-        if machine.is_up():
-            machine.execute("sync")
+        test_script()
+        # TODO: Collect coverage data
+        for machine in machines:
+            if machine.is_up():
+                machine.execute("sync")
 
 
 def serial_stdout_on() -> None:
···
     log._print_serial_logs = False
 
 
+class EnvDefault(argparse.Action):
+    """An argpars Action that takes values from the specified
+    environment variable as the flags default value.
+    """
+
+    def __init__(self, envvar, required=False, default=None, nargs=None, **kwargs):  # type: ignore
+        if not default and envvar:
+            if envvar in os.environ:
+                if nargs is not None and (nargs.isdigit() or nargs in ["*", "+"]):
+                    default = os.environ[envvar].split()
+                else:
+                    default = os.environ[envvar]
+                kwargs["help"] = (
+                    kwargs["help"] + f" (default from environment: {default})"
+                )
+        if required and default:
+            required = False
+        super(EnvDefault, self).__init__(
+            default=default, required=required, nargs=nargs, **kwargs
+        )
+
+    def __call__(self, parser, namespace, values, option_string=None):  # type: ignore
+        setattr(namespace, self.dest, values)
+
+
 @contextmanager
 def subtest(name: str) -> Iterator[None]:
     with log.nested(name):
···
         help="re-use a VM state coming from a previous run",
         action="store_true",
     )
-    (cli_args, vm_scripts) = arg_parser.parse_known_args()
+    arg_parser.add_argument(
+        "-I",
+        "--interactive",
+        help="drop into a python repl and run the tests interactively",
+        action="store_true",
+    )
+    arg_parser.add_argument(
+        "--start-scripts",
+        metavar="START-SCRIPT",
+        action=EnvDefault,
+        envvar="startScripts",
+        nargs="*",
+        help="start scripts for participating virtual machines",
+    )
+    arg_parser.add_argument(
+        "--vlans",
+        metavar="VLAN",
+        action=EnvDefault,
+        envvar="vlans",
+        nargs="*",
+        help="vlans to span by the driver",
+    )
+    arg_parser.add_argument(
+        "testscript",
+        action=EnvDefault,
+        envvar="testScript",
+        help="the test script to run",
+        type=pathlib.Path,
+    )
+
+    args = arg_parser.parse_args()
+    global test_script
+
+    def test_script() -> None:
+        with log.nested("running the VM test script"):
+            exec(pathlib.Path(args.testscript).read_text(), globals())
 
     log = Logger()
 
-    vlan_nrs = list(dict.fromkeys(os.environ.get("VLANS", "").split()))
-    vde_sockets = [create_vlan(v) for v in vlan_nrs]
+    vde_sockets = [create_vlan(v) for v in args.vlans]
     for nr, vde_socket, _, _ in vde_sockets:
         os.environ["QEMU_VDE_SOCKET_{}".format(nr)] = vde_socket
 
     machines = [
-        create_machine({"startCommand": s, "keepVmState": cli_args.keep_vm_state})
-        for s in vm_scripts
+        create_machine({"startCommand": s, "keepVmState": args.keep_vm_state})
+        for s in args.start_scripts
     ]
     machine_eval = [
         "{0} = machines[{1}]".format(m.name, idx) for idx, m in enumerate(machines)
···
         log.close()
 
     tic = time.time()
-    run_tests()
+    run_tests(args.interactive)
     toc = time.time()
     print("test script finished in {:.2f}s".format(toc - tic))

nixos/lib/testing-python.nix

···
         ''
           mkdir -p $out
 
-          LOGFILE=/dev/null tests='exec(os.environ["testScript"])' ${driver}/bin/nixos-test-driver
+          # effectively mute the XMLLogger
+          export LOGFILE=/dev/null
+
+          ${driver}/bin/nixos-test-driver
         '';
 
       passthru = driver.passthru // {
···
       ''
         mkdir -p $out/bin
 
+        vmStartScripts=($(for i in ${toString vms}; do echo $i/bin/run-*-vm; done))
         echo -n "$testScript" > $out/test-script
+        ln -s ${testDriver}/bin/nixos-test-driver $out/bin/nixos-test-driver
+
         ${lib.optionalString (!skipLint) ''
           PYFLAKES_BUILTINS="$(
             echo -n ${lib.escapeShellArg (lib.concatStringsSep "," nodeHostNames)},
···
           )" ${python3Packages.pyflakes}/bin/pyflakes $out/test-script
         ''}
 
-        ln -s ${testDriver}/bin/nixos-test-driver $out/bin/
-        vms=($(for i in ${toString vms}; do echo $i/bin/run-*-vm; done))
+        # set defaults through environment
+        # see: ./test-driver/test-driver.py argparse implementation
         wrapProgram $out/bin/nixos-test-driver \
-          --add-flags "''${vms[*]}" \
-          --run "export testScript=\"\$(${coreutils}/bin/cat $out/test-script)\"" \
-          --set VLANS '${toString vlans}'
-        ln -s ${testDriver}/bin/nixos-test-driver $out/bin/nixos-run-vms
-        wrapProgram $out/bin/nixos-run-vms \
-          --add-flags "''${vms[*]}" \
-          --set tests 'start_all(); join_all();' \
-          --set VLANS '${toString vlans}'
+          --set startScripts "''${vmStartScripts[*]}" \
+          --set testScript "$out/test-script" \
+          --set vlans '${toString vlans}'
       '');
 
   # Make a full-blown test