nixos/test-runner: Fix execute() flakiness

Instead of using the magic string, we now just base64-encode everything
and check for a newline.

1 files changed, 28 insertions, 12 deletions

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