Add ability to enforce pulling plan data during reboot

There are situations when getting the artifacts just before reboot
is very desired as the reboot may be risky operation which may lead
into situation when the system is not reachable, TMT timeouts
and is no longer able to fetch anything from the system.

This change adds optional -f switch to the tmt-reboot command
where the test can signal that it is entering "critical reboot" and
it's requesting to flush the artifacts (pull from TMT point of view)
produced by the test before the system reboots.

tmt/steps/execute/__init__.py

@@ -319,7 +319,7 @@ def handle_restart(self) -> bool:
 
         return True
 
-    def handle_reboot(self) -> bool:
+    def handle_reboot(self, flush_func=None) -> bool:
         """
         Reboot the guest if the test requested it.
 
@@ -344,6 +344,7 @@ def handle_reboot(self) -> bool:
 
         reboot_command: Optional[ShellScript] = None
         timeout: Optional[int] = None
+        flush_artifacts: bool = False
 
         if self.hard_reboot_requested:
             pass
@@ -362,6 +363,16 @@ def handle_reboot(self) -> bool:
             except ValueError:
                 timeout = None
 
+            flush_artifacts = reboot_data.get('flush_artifacts', flush_artifacts)
+
+            if flush_artifacts:
+                # Pull artifacts created in the plan data directory
+                self.logger.debug("Flush artifacts requested, pull the test data directory.", level=2)
+                if flush_func:
+                    flush_func()
+                else:
+                    self.guest.pull(source=self.test_data_path)
+
             os.remove(self.reboot_request_path)
             self.guest.push(self.test_data_path)
 

tmt/steps/execute/internal.py

@@ -1,4 +1,5 @@
 import dataclasses
+import functools
 import os
 import subprocess
 import textwrap
@@ -548,7 +549,9 @@ def _run_tests(
                     logger.verbose(
                         f"{duration} {test.name} [{progress}]", shift=shift)
                     try:
-                        if invocation.handle_reboot():
+                        if invocation.handle_reboot(
+                                flush_func=functools.partial(guest.pull, source=self.step.plan.data_directory)
+                        ):
                             continue
                     except tmt.utils.RebootTimeoutError:
                         for result in invocation.results:

tmt/steps/execute/scripts/tmt-reboot

@@ -13,11 +13,13 @@ PATH=/sbin:/usr/sbin:$PATH
 command=""
 timeout=""
 efi=True
-while getopts "c:t:e" flag; do
+flush_artifacts=False
+while getopts "c:t:e:f" flag; do
     case "${flag}" in
         c) command="${OPTARG}";;
         t) timeout="${OPTARG}";;
         e) efi=False;;
+        f) flush_artifacts=True;;
     esac
 done
 
@@ -37,4 +39,4 @@ if [ $efi = True ]; then
     fi
 fi
 
-flock "$TMT_TEST_PIDFILE_LOCK" tmt-reboot-core "$command" "$timeout"
+flock "$TMT_TEST_PIDFILE_LOCK" tmt-reboot-core "$command" "$timeout" "$flush_artifacts"

tmt/steps/execute/scripts/tmt-reboot-core

@@ -11,6 +11,7 @@ fi
 
 command="$1"
 timeout="$2"
+flush_artifacts="$3"
 
 # Thanks to being run while holding the pidfile lock, the file exists and should
 # no go away. It should be safe to read test PID and reboot-request filepath from
@@ -20,5 +21,5 @@ test_reboot_file="$(awk '{print $2}' < "$TMT_TEST_PIDFILE")"
 
 mkdir -p "$(dirname "$test_reboot_file")"
 
-printf '{"command": "%s", "timeout": "%s"}' "$command" "$timeout" > "$test_reboot_file"
+printf '{"command": "%s", "timeout": "%s", "flush_artifacts": "%s"}' "$command" "$timeout" "$flush_artifacts" > "$test_reboot_file"
 kill "$test_pid"