From 2ff8bf9c6b992f2be49dfb5194fe4a032b46c41c Mon Sep 17 00:00:00 2001
From: Pavel Holica <pholica@redhat.com>
Date: Fri, 26 Apr 2024 14:03:23 +0200
Subject: [PATCH] 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             | 13 ++++++++++++-
 tmt/steps/execute/internal.py             |  5 ++++-
 tmt/steps/execute/scripts/tmt-reboot      |  6 ++++--
 tmt/steps/execute/scripts/tmt-reboot-core |  3 ++-
 4 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/tmt/steps/execute/__init__.py b/tmt/steps/execute/__init__.py
index 1f109dbc04..df83af05ac 100644
--- a/tmt/steps/execute/__init__.py
+++ b/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)
 
diff --git a/tmt/steps/execute/internal.py b/tmt/steps/execute/internal.py
index 367f73b230..d171d378e9 100644
--- a/tmt/steps/execute/internal.py
+++ b/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:
diff --git a/tmt/steps/execute/scripts/tmt-reboot b/tmt/steps/execute/scripts/tmt-reboot
index 259dbec0b2..715091b25b 100755
--- a/tmt/steps/execute/scripts/tmt-reboot
+++ b/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"
diff --git a/tmt/steps/execute/scripts/tmt-reboot-core b/tmt/steps/execute/scripts/tmt-reboot-core
index 96ed3b8aba..fb35041197 100755
--- a/tmt/steps/execute/scripts/tmt-reboot-core
+++ b/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"