Add ability to flush also plan data before reboot

pholica · May 13, 2024 · 73ff913 · 73ff913
1 parent 09c3818
commit 73ff913
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git 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.
 
@@ -368,7 +368,10 @@ def handle_reboot(self) -> bool:
             if flush_artifacts:
                 # Pull artifacts created in the plan data directory
                 self.logger.debug("Flush artifacts requested, pull the test data directory.", level=2)
-                self.guest.pull(source=self.test_data_path)
+                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
@@ -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: