From af765f3abd7193ec5256cea3d5e8b7692dc99931 Mon Sep 17 00:00:00 2001
From: Cole Helbling <cole.helbling@determinate.systems>
Date: Thu, 18 Nov 2021 09:10:51 -0800
Subject: [PATCH] nixos/test-driver: give more functions nested labels

This will make it easier to trace through the test execution without
having to scroll through the entire kernel output.
---
 nixos/lib/test-driver/test-driver.py | 71 +++++++++++++++-------------
 1 file changed, 39 insertions(+), 32 deletions(-)

diff --git a/nixos/lib/test-driver/test-driver.py b/nixos/lib/test-driver/test-driver.py
index 7b8912acf503..da5a6cec4ec4 100755
--- a/nixos/lib/test-driver/test-driver.py
+++ b/nixos/lib/test-driver/test-driver.py
@@ -490,23 +490,24 @@ class Machine:
         return rootlog.nested(msg, my_attrs)
 
     def wait_for_monitor_prompt(self) -> str:
-        assert self.monitor is not None
-        answer = ""
-        while True:
-            undecoded_answer = self.monitor.recv(1024)
-            if not undecoded_answer:
-                break
-            answer += undecoded_answer.decode()
-            if answer.endswith("(qemu) "):
-                break
-        return answer
+        with self.nested("waiting for monitor prompt"):
+            assert self.monitor is not None
+            answer = ""
+            while True:
+                undecoded_answer = self.monitor.recv(1024)
+                if not undecoded_answer:
+                    break
+                answer += undecoded_answer.decode()
+                if answer.endswith("(qemu) "):
+                    break
+            return answer
 
     def send_monitor_command(self, command: str) -> str:
-        message = ("{}\n".format(command)).encode()
-        self.log("sending monitor command: {}".format(command))
-        assert self.monitor is not None
-        self.monitor.send(message)
-        return self.wait_for_monitor_prompt()
+        with self.nested("sending monitor command: {}".format(command)):
+            message = ("{}\n".format(command)).encode()
+            assert self.monitor is not None
+            self.monitor.send(message)
+            return self.wait_for_monitor_prompt()
 
     def wait_for_unit(self, unit: str, user: Optional[str] = None) -> None:
         """Wait for a systemd unit to get into "active" state.
@@ -533,7 +534,12 @@ class Machine:
 
             return state == "active"
 
-        retry(check_active)
+        with self.nested(
+            "waiting for unit {}{}".format(
+                unit, f" with user {user}" if user is not None else ""
+            )
+        ):
+            retry(check_active)
 
     def get_unit_info(self, unit: str, user: Optional[str] = None) -> Dict[str, str]:
         status, lines = self.systemctl('--no-pager show "{}"'.format(unit), user)
@@ -752,7 +758,8 @@ class Machine:
             status, _ = self.execute("nc -z localhost {}".format(port))
             return status != 0
 
-        retry(port_is_closed)
+        with self.nested("waiting for TCP port {} to be closed"):
+            retry(port_is_closed)
 
     def start_job(self, jobname: str, user: Optional[str] = None) -> Tuple[int, str]:
         return self.systemctl("start {}".format(jobname), user)
@@ -886,20 +893,20 @@ class Machine:
             retry(screen_matches)
 
     def wait_for_console_text(self, regex: str) -> None:
-        self.log("waiting for {} to appear on console".format(regex))
-        # Buffer the console output, this is needed
-        # to match multiline regexes.
-        console = io.StringIO()
-        while True:
-            try:
-                console.write(self.last_lines.get())
-            except queue.Empty:
-                self.sleep(1)
-                continue
-            console.seek(0)
-            matches = re.search(regex, console.read())
-            if matches is not None:
-                return
+        with self.nested("waiting for {} to appear on console".format(regex)):
+            # Buffer the console output, this is needed
+            # to match multiline regexes.
+            console = io.StringIO()
+            while True:
+                try:
+                    console.write(self.last_lines.get())
+                except queue.Empty:
+                    self.sleep(1)
+                    continue
+                console.seek(0)
+                matches = re.search(regex, console.read())
+                if matches is not None:
+                    return
 
     def send_key(self, key: str) -> None:
         key = CHAR_TO_KEY.get(key, key)
@@ -1014,7 +1021,7 @@ class Machine:
                 )
             return any(pattern.search(name) for name in names)
 
-        with self.nested("Waiting for a window to appear"):
+        with self.nested("waiting for a window to appear"):
             retry(window_is_visible)
 
     def sleep(self, secs: int) -> None: