atex 0.8__py3-none-any.whl → 0.10__py3-none-any.whl

atex/connection/ssh.py

@@ -18,6 +18,7 @@ import os
 import time
 import shlex
 import tempfile
+import threading
 import subprocess
 from pathlib import Path
 
@@ -38,7 +39,7 @@ DEFAULT_OPTIONS = {
 }
 
 
-class SSHError(Exception):
+class SSHError(ConnectionError):
     pass
 
 
@@ -132,16 +133,16 @@ def _rsync_host_cmd(*args, options, password=None, sudo=None):
     )
 
 
-class StatelessSSHConn(Connection):
+class StatelessSSHConnection(Connection):
     """
     Implements the Connection API using a ssh(1) client using "standalone"
     (stateless) logic - connect() and disconnect() are no-op, .cmd() simply
     executes the ssh client and .rsync() executes 'rsync -e ssh'.
 
-    Compared to ManagedSSHConn, this may be slow for many .cmd() calls,
+    Compared to ManagedSSHConnection, this may be slow for many .cmd() calls,
     but every call is stateless, there is no persistent connection.
 
-    If you need only one .cmd(), this will be faster than ManagedSSHConn.
+    If you need only one .cmd(), this will be faster than ManagedSSHConnection.
     """
 
     def __init__(self, options, *, password=None, sudo=None):
@@ -154,7 +155,6 @@ class StatelessSSHConn(Connection):
         If 'sudo' specifies a username, call sudo(8) on the remote shell
         to run under a different user on the remote host.
         """
-        super().__init__()
         self.options = DEFAULT_OPTIONS.copy()
         self.options.update(options)
         self.password = password
@@ -162,7 +162,7 @@ class StatelessSSHConn(Connection):
         self._tmpdir = None
         self._master_proc = None
 
-    def connect(self):
+    def connect(self, block=True):
         """
         Optional, .cmd() and .rsync() work without it, but it is provided here
         for compatibility with the Connection API.
@@ -178,10 +178,10 @@ class StatelessSSHConn(Connection):
         unified_options = self.options.copy()
         if options:
             unified_options.update(options)
-        unified_options["RemoteCommand"] = _shell_cmd(command, sudo=self.sudo)
+        if command:
+            unified_options["RemoteCommand"] = _shell_cmd(command, sudo=self.sudo)
         return func(
             _options_to_ssh(unified_options, password=self.password),
-            skip_frames=1,
             **func_args,
         )
 
@@ -196,7 +196,6 @@ class StatelessSSHConn(Connection):
                 password=self.password,
                 sudo=self.sudo,
             ),
-            skip_frames=1,
             check=True,
             stdin=subprocess.DEVNULL,
             **func_args,
@@ -215,17 +214,15 @@ class StatelessSSHConn(Connection):
 # checks .assert_master() and manually signals the running clients
 # when it gets DisconnectedError from it.
 
-class ManagedSSHConn(Connection):
+class ManagedSSHConnection(Connection):
     """
     Implements the Connection API using one persistently-running ssh(1) client
     started in a 'ControlMaster' mode, with additional ssh clients using that
     session to execute remote commands. Similarly, .rsync() uses it too.
 
-    This is much faster than StatelessSSHConn when executing multiple commands,
-    but contains a complex internal state (what if ControlMaster disconnects?).
-
-    Hence why this implementation provides extra non-standard-Connection methods
-    to manage this complexity.
+    This is much faster than StatelessSSHConnection when executing multiple
+    commands, but contains a complex internal state (what if ControlMaster
+    disconnects?).
     """
 
     # TODO: thread safety and locking via self.lock ?
@@ -240,7 +237,7 @@ class ManagedSSHConn(Connection):
         If 'sudo' specifies a username, call sudo(8) on the remote shell
         to run under a different user on the remote host.
         """
-        super().__init__()
+        self.lock = threading.RLock()
         self.options = DEFAULT_OPTIONS.copy()
         self.options.update(options)
         self.password = password
@@ -267,8 +264,9 @@ class ManagedSSHConn(Connection):
         proc = self._master_proc
         if not proc:
             return
+        util.debug(f"disconnecting: {self.options}")
         proc.kill()
-        # don"t zombie forever, return EPIPE on any attempts to write to us
+        # don't zombie forever, return EPIPE on any attempts to write to us
         proc.stdout.close()
         proc.wait()
         (self._tmpdir / "control.sock").unlink(missing_ok=True)
@@ -286,19 +284,26 @@ class ManagedSSHConn(Connection):
         sock = self._tmpdir / "control.sock"
 
         if not self._master_proc:
+            util.debug(f"connecting: {self.options}")
             options = self.options.copy()
             options["SessionType"] = "none"
             options["ControlMaster"] = "yes"
             options["ControlPath"] = sock
             self._master_proc = util.subprocess_Popen(
-                _options_to_ssh(options),
+                _options_to_ssh(options, password=self.password),
                 stdin=subprocess.DEVNULL,
                 stdout=subprocess.PIPE,
                 stderr=subprocess.STDOUT,
                 cwd=str(self._tmpdir),
+                start_new_session=True,  # resist Ctrl-C
             )
             os.set_blocking(self._master_proc.stdout.fileno(), False)
 
+        # NOTE: ideally, we would .read() before checking .poll() because
+        #       if the process writes a lot, it gets stuck in the pipe
+        #       (in kernel) and the process never ends; but output-appending
+        #       code would be obscure, and ssh(1) never outputs that much ..
+
         proc = self._master_proc
         if block:
             while proc.poll() is None:
@@ -309,7 +314,6 @@ class ManagedSSHConn(Connection):
                 code = proc.poll()
                 out = proc.stdout.read()
                 self._master_proc = None
-                # TODO: ConnectError should probably be generalized for Connection
                 raise ConnectError(
                     f"SSH ControlMaster failed to start on {self._tmpdir} with {code}:\n{out}",
                 )
@@ -318,55 +322,44 @@ class ManagedSSHConn(Connection):
             if code is not None:
                 out = proc.stdout.read()
                 self._master_proc = None
-                # TODO: ConnectError should probably be generalized for Connection
                 raise ConnectError(
                     f"SSH ControlMaster failed to start on {self._tmpdir} with {code}:\n{out}",
                 )
             elif not sock.exists():
                 raise BlockingIOError("SSH ControlMaster not yet ready")
 
-    def add_local_forward(self, *spec):
+    def forward(self, forward_type, *spec, cancel=False):
         """
         Add (one or more) ssh forwarding specifications as 'spec' to an
         already-connected instance. Each specification has to follow the
-        format of ssh client's LocalForward option (see ssh_config(5)).
-        """
-        self.assert_master()
-        options = self.options.copy()
-        options["LocalForward"] = spec
-        options["ControlPath"] = self._tmpdir / "control.sock"
-        util.subprocess_run(
-            _options_to_ssh(options, extra_cli_flags=("-O", "forward")),
-            skip_frames=1,
-            check=True,
-        )
+        format of LocalForward or RemoteForward (see ssh_config(5)).
+        Ie. "1234 1.2.3.4:22" or "0.0.0.0:1234 1.2.3.4:22".
 
-    def add_remote_forward(self, *spec):
-        """
-        Add (one or more) ssh forwarding specifications as 'spec' to an
-        already-connected instance. Each specification has to follow the
-        format of ssh client's RemoteForward option (see ssh_config(5)).
+        'forward_type' must be either LocalForward or RemoteForward.
+
+        If 'cancel' is True, cancel the forwarding instead of adding it.
         """
+        assert forward_type in ("LocalForward", "RemoteForward")
         self.assert_master()
-        options = self.options.copy()
-        options["RemoteForward"] = spec
+        options = DEFAULT_OPTIONS.copy()
+        options[forward_type] = spec
         options["ControlPath"] = self._tmpdir / "control.sock"
+        action = "forward" if not cancel else "cancel"
         util.subprocess_run(
-            _options_to_ssh(options, extra_cli_flags=("-O", "forward")),
-            skip_frames=1,
+            _options_to_ssh(options, extra_cli_flags=("-O", action)),
             check=True,
         )
 
-    def cmd(self, command, options=None, func=util.subprocess_run, **func_args):
+    def cmd(self, command, *, options=None, func=util.subprocess_run, **func_args):
         self.assert_master()
         unified_options = self.options.copy()
         if options:
             unified_options.update(options)
-        unified_options["RemoteCommand"] = _shell_cmd(command, sudo=self.sudo)
+        if command:
+            unified_options["RemoteCommand"] = _shell_cmd(command, sudo=self.sudo)
         unified_options["ControlPath"] = self._tmpdir / "control.sock"
         return func(
-            _options_to_ssh(unified_options, password=self.password),
-            skip_frames=1,
+            _options_to_ssh(unified_options),
             **func_args,
         )
 
@@ -380,10 +373,8 @@ class ManagedSSHConn(Connection):
             _rsync_host_cmd(
                 *args,
                 options=unified_options,
-                password=self.password,
                 sudo=self.sudo,
             ),
-            skip_frames=1,
             check=True,
             stdin=subprocess.DEVNULL,
             **func_args,

atex/executor/executor.py

@@ -1,4 +1,6 @@
 import os
+import enum
+import time
 import select
 import threading
 import contextlib
@@ -24,13 +26,15 @@ class Executor:
     and uploaded files by those tests.
 
         tests_repo = "path/to/cloned/tests"
-        tests_data = atex.fmf.FMFTests(tests_repo, "/plans/default")
+        fmf_tests = atex.fmf.FMFTests(tests_repo, "/plans/default")
 
-        with Executor(tests_data, conn) as e:
-            e.upload_tests(tests_repo)
-            e.setup_plan()
-            e.run_test("/some/test", "results/here.json", "uploaded/files/here")
+        with Executor(fmf_tests, conn) as e:
+            e.upload_tests()
+            e.plan_prepare()
+            Path("output_here").mkdir()
+            e.run_test("/some/test", "output_here")
             e.run_test(...)
+            e.plan_finish()
 
     One Executor instance may be used to run multiple tests sequentially.
     In addition, multiple Executor instances can run in parallel on the same
@@ -38,35 +42,39 @@ class Executor:
 
         conn.cmd(["mkdir", "-p", "/shared"])
 
-        with Executor(tests_data, conn, state_dir="/shared") as e:
-            e.upload_tests(tests_repo)
-            e.setup_plan()
+        with Executor(fmf_tests, conn, state_dir="/shared") as e:
+            e.upload_tests()
+            e.plan_prepare()
 
         # in parallel (ie. threading or multiprocessing)
-        with Executor(tests_data, unique_conn, state_dir="/shared") as e:
+        with Executor(fmf_tests, unique_conn, state_dir="/shared") as e:
             e.run_test(...)
     """
 
-    def __init__(self, fmf_tests, connection, *, state_dir=None):
+    def __init__(self, fmf_tests, connection, *, env=None, state_dir=None):
         """
         'fmf_tests' is a class FMFTests instance with (discovered) tests.
 
         'connection' is a class Connection instance, already fully connected.
 
+        'env' is a dict of extra environment variables to pass to the
+        plan prepare/finish scripts and to all tests.
+
         'state_dir' is a string or Path specifying path on the remote system for
         storing additional data, such as tests, execution wrappers, temporary
         plan-exported variables, etc. If left as None, a tmpdir is used.
         """
         self.lock = threading.RLock()
-        self.conn = connection
         self.fmf_tests = fmf_tests
+        self.conn = connection
+        self.env = env or {}
         self.state_dir = state_dir
         self.work_dir = None
         self.tests_dir = None
         self.plan_env_file = None
         self.cancelled = False
 
-    def setup(self):
+    def start(self):
         with self.lock:
             state_dir = self.state_dir
 
@@ -96,7 +104,10 @@ class Executor:
                 # use the tmpdir as work_dir, avoid extra mkdir over conn
                 self.work_dir = tmp_dir
 
-    def cleanup(self):
+        # create / truncate the TMT_PLAN_ENVIRONMENT_FILE
+        self.conn.cmd(("truncate", "-s", "0", self.plan_env_file), check=True)
+
+    def stop(self):
         with self.lock:
             work_dir = self.work_dir
 
@@ -109,11 +120,15 @@ class Executor:
             self.plan_env_file = None
 
     def __enter__(self):
-        self.setup()
-        return self
+        try:
+            self.start()
+            return self
+        except Exception:
+            self.stop()
+            raise
 
     def __exit__(self, exc_type, exc_value, traceback):
-        self.cleanup()
+        self.stop()
 
     def cancel(self):
         with self.lock:
@@ -125,19 +140,37 @@ class Executor:
         __init__() inside 'fmf_tests', to the remote host.
         """
         self.conn.rsync(
-            "-rv" if util.in_debug_mode() else "-rq",
-            "--delete", "--exclude=.git/",
+            "-r", "--delete", "--exclude=.git/",
             f"{self.fmf_tests.root}/",
             f"remote:{self.tests_dir}",
+            func=util.subprocess_log,
         )
 
-    def setup_plan(self):
+    def _run_prepare_scripts(self, scripts):
+        # make envionment for 'prepare' scripts
+        env = {
+            **self.fmf_tests.plan_env,
+            **self.env,
+            "TMT_PLAN_ENVIRONMENT_FILE": self.plan_env_file,
+        }
+        env_args = (f"{k}={v}" for k, v in env.items())
+        # run the scripts
+        for script in scripts:
+            self.conn.cmd(
+                ("env", *env_args, "bash"),
+                func=util.subprocess_log,
+                stderr=subprocess.STDOUT,
+                input=script,
+                check=True,
+            )
+
+    def plan_prepare(self):
         """
         Install packages and run scripts extracted from a TMT plan by a FMFTests
         instance given during class initialization.
 
-        Also prepare additional environment for tests, ie. create and export
-        a path to TMT_PLAN_ENVIRONMENT_FILE.
+        Also run additional scripts specified under the 'prepare' step inside
+        the fmf metadata of a plan.
         """
         # install packages from the plan
         if self.fmf_tests.prepare_pkgs:
@@ -146,70 +179,79 @@ class Executor:
                     "dnf", "-y", "--setopt=install_weak_deps=False",
                     "install", *self.fmf_tests.prepare_pkgs,
                 ),
-                check=True,
-                stdout=None if util.in_debug_mode() else subprocess.DEVNULL,
+                func=util.subprocess_log,
                 stderr=subprocess.STDOUT,
+                check=True,
             )
 
-        # make envionment for 'prepare' scripts
-        self.conn.cmd(("truncate", "-s", "0", self.plan_env_file), check=True)
-        env = self.fmf_tests.plan_env.copy()
-        env["TMT_PLAN_ENVIRONMENT_FILE"] = self.plan_env_file
-        env_args = (f"{k}={v}" for k, v in env.items())
+        # run 'prepare' scripts from the plan
+        if scripts := self.fmf_tests.prepare_scripts:
+            self._run_prepare_scripts(scripts)
 
-        # run the prepare scripts
-        for script in self.fmf_tests.prepare_scripts:
-            self.conn.cmd(
-                ("env", *env_args, "bash"),
-                input=script,
-                text=True,
-                check=True,
-                stdout=None if util.in_debug_mode() else subprocess.DEVNULL,
-                stderr=subprocess.STDOUT,
-            )
+    def plan_finish(self):
+        """
+        Run any scripts specified under the 'finish' step inside
+        the fmf metadata of a plan.
+        """
+        if scripts := self.fmf_tests.finish_scripts:
+            self._run_prepare_scripts(scripts)
 
-    def run_test(self, test_name, json_file, files_dir, *, env=None):
+    class State(enum.Enum):
+        STARTING_TEST = enum.auto()
+        READING_CONTROL = enum.auto()
+        WAITING_FOR_EXIT = enum.auto()
+        RECONNECTING = enum.auto()
+
+    def run_test(self, test_name, output_dir, *, env=None):
         """
         Run one test on the remote system.
 
         'test_name' is a string with test name.
 
-        'json_file' is a destination file (string or Path) for results.
-
-        'files_dir' is a destination dir (string or Path) for uploaded files.
+        'output_dir' is a destination dir (string or Path) for results reported
+        and files uploaded by the test. Results are always stored in a line-JSON
+        format in a file named 'results', files are always uploaded to directory
+        named 'files', both inside 'output_dir'.
+        The path for 'output_dir' must already exist and be an empty directory
+        (ie. typically a tmpdir).
 
         'env' is a dict of extra environment variables to pass to the test.
 
         Returns an integer exit code of the test script.
         """
+        output_dir = Path(output_dir)
         test_data = self.fmf_tests.tests[test_name]
 
-        # start with fmf-plan-defined environment
-        env_vars = self.fmf_tests.plan_env.copy()
-        # append fmf-test-defined environment into it
-        for item in fmf.listlike(test_data, "environment"):
-            env_vars.update(item)
-        # append additional variables typically exported by tmt
-        env_vars["TMT_PLAN_ENVIRONMENT_FILE"] = self.plan_env_file
-        env_vars["TMT_TEST_NAME"] = test_name
-        env_vars["ATEX_TEST_NAME"] = test_name
-        # append variables given to this function call
-        if env:
-            env_vars.update(env)
-
         # run a setup script, preparing wrapper + test scripts
         setup_script = scripts.test_setup(
             test=scripts.Test(test_name, test_data, self.fmf_tests.test_dirs[test_name]),
             tests_dir=self.tests_dir,
             wrapper_exec=f"{self.work_dir}/wrapper.sh",
             test_exec=f"{self.work_dir}/test.sh",
+            test_yaml=f"{self.work_dir}/metadata.yaml",
         )
         self.conn.cmd(("bash",), input=setup_script, text=True, check=True)
 
+        # start with fmf-plan-defined environment
+        env_vars = {
+            **self.fmf_tests.plan_env,
+            "TMT_PLAN_ENVIRONMENT_FILE": self.plan_env_file,
+            "TMT_TEST_NAME": test_name,
+            "TMT_TEST_METADATA": f"{self.work_dir}/metadata.yaml",
+        }
+        # append fmf-test-defined environment into it
+        for item in fmf.listlike(test_data, "environment"):
+            env_vars.update(item)
+        # append the Executor-wide environment passed to __init__()
+        env_vars.update(self.env)
+        # append variables given to this function call
+        if env:
+            env_vars.update(env)
+
         with contextlib.ExitStack() as stack:
-            reporter = stack.enter_context(Reporter(json_file, files_dir))
-            testout_fd = stack.enter_context(reporter.open_tmpfile())
+            reporter = stack.enter_context(Reporter(output_dir, "results", "files"))
             duration = Duration(test_data.get("duration", "5m"))
+            control = testcontrol.TestControl(reporter=reporter, duration=duration)
 
             test_proc = None
             control_fd = None
@@ -223,23 +265,19 @@ class Executor:
                     test_proc.wait()
                 raise TestAbortedError(msg) from None
 
+            exception = None
+
             try:
-                # TODO: probably enum
-                state = "starting_test"
+                state = self.State.STARTING_TEST
                 while not duration.out_of_time():
                     with self.lock:
                         if self.cancelled:
                             abort("cancel requested")
 
-                    if state == "starting_test":
+                    if state == self.State.STARTING_TEST:
                         control_fd, pipe_w = os.pipe()
                         os.set_blocking(control_fd, False)
-                        control = testcontrol.TestControl(
-                            control_fd=control_fd,
-                            reporter=reporter,
-                            duration=duration,
-                            testout_fd=testout_fd,
-                        )
+                        control.reassign(control_fd)
                         # reconnect/reboot count (for compatibility)
                         env_vars["TMT_REBOOT_COUNT"] = str(reconnects)
                         env_vars["TMT_TEST_RESTART_COUNT"] = str(reconnects)
@@ -249,13 +287,13 @@ class Executor:
                         test_proc = self.conn.cmd(
                             ("env", *env_args, f"{self.work_dir}/wrapper.sh"),
                             stdout=pipe_w,
-                            stderr=testout_fd,
+                            stderr=reporter.testout_fobj.fileno(),
                             func=util.subprocess_Popen,
                         )
                         os.close(pipe_w)
-                        state = "reading_control"
+                        state = self.State.READING_CONTROL
 
-                    elif state == "reading_control":
+                    elif state == self.State.READING_CONTROL:
                         rlist, _, xlist = select.select((control_fd,), (), (control_fd,), 0.1)
                         if xlist:
                             abort(f"got exceptional condition on control_fd {control_fd}")
@@ -264,9 +302,9 @@ class Executor:
                             if control.eof:
                                 os.close(control_fd)
                                 control_fd = None
-                                state = "waiting_for_exit"
+                                state = self.State.WAITING_FOR_EXIT
 
-                    elif state == "waiting_for_exit":
+                    elif state == self.State.WAITING_FOR_EXIT:
                         # control stream is EOF and it has nothing for us to read,
                         # we're now just waiting for proc to cleanly terminate
                         try:
@@ -279,7 +317,7 @@ class Executor:
                                 self.conn.disconnect()
                                 # if reconnect was requested, do so, otherwise abort
                                 if control.reconnect:
-                                    state = "reconnecting"
+                                    state = self.State.RECONNECTING
                                     if control.reconnect != "always":
                                         control.reconnect = None
                                 else:
@@ -291,13 +329,20 @@ class Executor:
                         except subprocess.TimeoutExpired:
                             pass
 
-                    elif state == "reconnecting":
+                    elif state == self.State.RECONNECTING:
                         try:
                             self.conn.connect(block=False)
                             reconnects += 1
-                            state = "starting_test"
+                            state = self.State.STARTING_TEST
                         except BlockingIOError:
-                            pass
+                            # avoid 100% CPU spinning if the connection it too slow
+                            # to come up (ie. ssh ControlMaster socket file not created)
+                            time.sleep(0.5)
+                        except ConnectionError:
+                            # can happen when ie. ssh is connecting over a LocalForward port,
+                            # causing 'read: Connection reset by peer' instead of timeout
+                            # - just retry again after a short delay
+                            time.sleep(0.5)
 
                     else:
                         raise AssertionError("reached unexpected state")
@@ -305,12 +350,19 @@ class Executor:
                 else:
                     abort("test duration timeout reached")
 
-                # testing successful, do post-testing tasks
+                # testing successful
 
                 # test wrapper hasn't provided exitcode
                 if control.exit_code is None:
                     abort("exitcode not reported, wrapper bug?")
 
+                return control.exit_code
+
+            except Exception as e:
+                exception = e
+                raise
+
+            finally:
                 # partial results that were never reported
                 if control.partial_results:
                     for result in control.partial_results.values():
@@ -320,59 +372,32 @@ class Executor:
                             control.nameless_result_seen = True
                         if testout := result.get("testout"):
                             try:
-                                reporter.link_tmpfile_to(testout_fd, testout, name)
+                                reporter.link_testout(testout, name)
                             except FileExistsError:
                                 raise testcontrol.BadReportJSONError(
                                     f"file '{testout}' already exists",
                                 ) from None
                         reporter.report(result)
 
-                # test hasn't reported a result for itself, add an automatic one
-                # as specified in RESULTS.md
-                # {"status": "pass", "testout": "output.txt"}
-                if not control.nameless_result_seen:
-                    reporter.link_tmpfile_to(testout_fd, "output.txt")
+                # if an unexpected infrastructure-related exception happened
+                if exception:
+                    try:
+                        reporter.link_testout("output.txt")
+                    except FileExistsError:
+                        pass
                     reporter.report({
-                        "status": "pass" if control.exit_code == 0 else "fail",
+                        "status": "infra",
+                        "note": repr(exception),
                         "testout": "output.txt",
                     })
 
-                return control.exit_code
-
-            except Exception:
-                # if the test hasn't reported a result for itself, but still
-                # managed to break something, provide at least the default log
-                # for manual investigation - otherwise test output disappears
-                if not control.nameless_result_seen:
+                # if the test hasn't reported a result for itself
+                elif not control.nameless_result_seen:
                     try:
-                        reporter.link_tmpfile_to(testout_fd, "output.txt")
-                        reporter.report({
-                            "status": "infra",
-                            "testout": "output.txt",
-                        })
-                    # in case outout.txt exists as a directory
+                        reporter.link_testout("output.txt")
                     except FileExistsError:
                         pass
-                raise
-
-
-#__all__ = [
-#    info.name for info in _pkgutil.iter_modules(__spec__.submodule_search_locations)
-#]
-#
-#
-#import importlib as _importlib
-#import pkgutil as _pkgutil
-#
-#
-#def __dir__():
-#    return __all__
-#
-#
-## lazily import submodules
-#def __getattr__(attr):
-#    # importing a module known to exist
-#    if attr in __all__:
-#        return _importlib.import_module(f".{attr}", __name__)
-#    else:
-#        raise AttributeError(f"module '{__name__}' has no attribute '{attr}'")
+                    reporter.report({
+                        "status": "pass" if control.exit_code == 0 else "fail",
+                        "testout": "output.txt",
+                    })