nixosTests.kexec: reconnect properly after 2nd kexec

By some miracle, before, it was possible to reconnect to the `node1` without
doing any relevant dance.

But now we are direct booting (¿), it seems like we need to do the right things.

This introduces a `check_output` flag for `execute` because we do not want to steal the
messages from the backdoor service as we might execute the kexec too fast compared
to when we will reconnect.

Therefore, we will let the message in the pipe if needed.

2 files changed, 14 insertions, 2 deletions

diff --git a/nixos/lib/test-driver/test_driver/machine.py b/nixos/lib/test-driver/test_driver/machine.py
index 81d3e19084ed..7ef0ae89910b 100644
--- a/nixos/lib/test-driver/test_driver/machine.py
+++ b/nixos/lib/test-driver/test_driver/machine.py
@@ -514,7 +514,11 @@ class Machine:
         return "".join(output_buffer)
 
     def execute(
-        self, command: str, check_return: bool = True, timeout: Optional[int] = 900
+        self,
+        command: str,
+        check_return: bool = True,
+        check_output: bool = True,
+        timeout: Optional[int] = 900,
     ) -> Tuple[int, str]:
         self.run_callbacks()
         self.connect()
@@ -535,6 +539,9 @@ class Machine:
         assert self.shell
         self.shell.send(out_command.encode())
 
+        if not check_output:
+            return (-2, "")
+
         # Get the output
         output = base64.b64decode(self._next_newline_closed_block_from_shell())
 
diff --git a/nixos/tests/kexec.nix b/nixos/tests/kexec.nix
index c828215bf539..5c0f0b60c853 100644
--- a/nixos/tests/kexec.nix
+++ b/nixos/tests/kexec.nix
@@ -15,6 +15,8 @@ import ./make-test-python.nix ({ pkgs, lib, ... }: {
       environment.systemPackages = [ pkgs.hello ];
       imports = [
         "${modulesPath}/installer/netboot/netboot-minimal.nix"
+        "${modulesPath}/testing/test-instrumentation.nix"
+        "${modulesPath}/profiles/qemu-guest.nix"
       ];
     };
   };
@@ -35,7 +37,10 @@ import ./make-test-python.nix ({ pkgs, lib, ... }: {
     # Kexec node1 to the toplevel of node2 via the kexec-boot script
     node1.succeed('touch /run/foo')
     node1.fail('hello')
-    node1.execute('${nodes.node2.config.system.build.kexecTree}/kexec-boot', check_return=False)
+    node1.execute('${nodes.node2.config.system.build.kexecTree}/kexec-boot', check_output=False)
+    node1.connected = False
+    node1.connect()
+    node1.wait_for_unit("multi-user.target")
     node1.succeed('! test -e /run/foo')
     node1.succeed('hello')
     node1.succeed('[ "$(hostname)" = "node2" ]')