Merge pull request #142747 from helsinki-systems/fix/test-runner-execute

nixos/test-runner: Fix execute() flakiness

6 files changed, 45 insertions, 17 deletions

diff --git a/nixos/doc/manual/development/writing-nixos-tests.section.md b/nixos/doc/manual/development/writing-nixos-tests.section.md
index 8471e7608af9..bf80099f1a27 100644
--- a/nixos/doc/manual/development/writing-nixos-tests.section.md
+++ b/nixos/doc/manual/development/writing-nixos-tests.section.md
@@ -159,6 +159,11 @@ The following methods are available on machine objects:
 `execute`
 
 :   Execute a shell command, returning a list `(status, stdout)`.
+    Takes an optional parameter `check_return` that defaults to `True`.
+    Setting this parameter to `False` will not check for the return code
+    and return -1 instead. This can be used for commands that shut down
+    the VM and would therefore break the pipe that would be used for
+    retrieving the return code.
 
 `succeed`
 
diff --git a/nixos/doc/manual/from_md/development/writing-nixos-tests.section.xml b/nixos/doc/manual/from_md/development/writing-nixos-tests.section.xml
index 83a96d5bb224..a3b63422433e 100644
--- a/nixos/doc/manual/from_md/development/writing-nixos-tests.section.xml
+++ b/nixos/doc/manual/from_md/development/writing-nixos-tests.section.xml
@@ -266,7 +266,13 @@ start_all()
       <listitem>
         <para>
           Execute a shell command, returning a list
-          <literal>(status, stdout)</literal>.
+          <literal>(status, stdout)</literal>. Takes an optional
+          parameter <literal>check_return</literal> that defaults to
+          <literal>True</literal>. Setting this parameter to
+          <literal>False</literal> will not check for the return code
+          and return -1 instead. This can be used for commands that shut
+          down the VM and would therefore break the pipe that would be
+          used for retrieving the return code.
         </para>
       </listitem>
     </varlistentry>
diff --git a/nixos/lib/test-driver/test-driver.py b/nixos/lib/test-driver/test-driver.py
index e4d93418a22f..a7c0484060f2 100755
--- a/nixos/lib/test-driver/test-driver.py
+++ b/nixos/lib/test-driver/test-driver.py
@@ -581,24 +581,40 @@ class Machine:
                     + "'{}' but it is in state ‘{}’".format(require_state, state)
                 )
 
-    def execute(self, command: str) -> Tuple[int, str]:
+    def _next_newline_closed_block_from_shell(self) -> str:
+        assert self.shell
+        output_buffer = []
+        while True:
+            # This receives up to 4096 bytes from the socket
+            chunk = self.shell.recv(4096)
+            if not chunk:
+                # Probably a broken pipe, return the output we have
+                break
+
+            decoded = chunk.decode()
+            output_buffer += [decoded]
+            if decoded[-1] == "\n":
+                break
+        return "".join(output_buffer)
+
+    def execute(self, command: str, check_return: bool = True) -> Tuple[int, str]:
         self.connect()
 
-        out_command = "( set -euo pipefail; {} ); echo '|!=EOF' $?\n".format(command)
+        out_command = f"( set -euo pipefail; {command} ) | (base64 --wrap 0; echo)\n"
         assert self.shell
         self.shell.send(out_command.encode())
 
-        output = ""
-        status_code_pattern = re.compile(r"(.*)\|\!=EOF\s+(\d+)")
+        # Get the output
+        output = base64.b64decode(self._next_newline_closed_block_from_shell())
 
-        while True:
-            chunk = self.shell.recv(4096).decode(errors="ignore")
-            match = status_code_pattern.match(chunk)
-            if match:
-                output += match[1]
-                status_code = int(match[2])
-                return (status_code, output)
-            output += chunk
+        if not check_return:
+            return (-1, output.decode())
+
+        # Get the return code
+        self.shell.send("echo ${PIPESTATUS[0]}\n".encode())
+        rc = int(self._next_newline_closed_block_from_shell().strip())
+
+        return (rc, output.decode())
 
     def shell_interact(self) -> None:
         """Allows you to interact with the guest shell
diff --git a/nixos/tests/hibernate.nix b/nixos/tests/hibernate.nix
index f0d5da4a95ed..508e7aa64c0d 100644
--- a/nixos/tests/hibernate.nix
+++ b/nixos/tests/hibernate.nix
@@ -95,7 +95,7 @@ in makeTest {
           "mkswap /dev/vda1 -L swap",
           # Install onto /mnt
           "nix-store --load-db < ${pkgs.closureInfo {rootPaths = [installedSystem];}}/registration",
-          "nixos-install --root /mnt --system ${installedSystem} --no-root-passwd",
+          "nixos-install --root /mnt --system ${installedSystem} --no-root-passwd --no-channel-copy >&2",
       )
       machine.shutdown()
 
@@ -110,7 +110,7 @@ in makeTest {
       )
 
       # Hibernate machine
-      hibernate.succeed("systemctl hibernate &")
+      hibernate.execute("systemctl hibernate &", check_return=False)
       hibernate.wait_for_shutdown()
 
       # Restore machine from hibernation, validate our ramfs file is there.
diff --git a/nixos/tests/kexec.nix b/nixos/tests/kexec.nix
index ec0cd9796b0e..036b9cab04f2 100644
--- a/nixos/tests/kexec.nix
+++ b/nixos/tests/kexec.nix
@@ -18,7 +18,7 @@ import ./make-test-python.nix ({ pkgs, lib, ...} : {
   testScript =
     ''
       machine.wait_for_unit("multi-user.target")
-      machine.execute("systemctl kexec &")
+      machine.execute("systemctl kexec &", check_return=False)
       machine.connected = False
       machine.wait_for_unit("multi-user.target")
     '';
diff --git a/nixos/tests/switch-test.nix b/nixos/tests/switch-test.nix
index 4caa7d98f47f..7ea07a390b80 100644
--- a/nixos/tests/switch-test.nix
+++ b/nixos/tests/switch-test.nix
@@ -392,7 +392,8 @@ import ./make-test-python.nix ({ pkgs, ...} : {
         machine.succeed("touch /testpath")
         machine.wait_until_succeeds("test -f /testpath-modified")
 
-        machine.succeed("rm /testpath /testpath-modified")
+        machine.succeed("rm /testpath")
+        machine.succeed("rm /testpath-modified")
         switch_to_specialisation("with-path-modified")
 
         machine.succeed("touch /testpath")