Merge pull request #348 from jumpstarter-dev/uboot-even-better

UbootConsole: Reuse pexpect instance

Author: NickCao

packages/jumpstarter-driver-uboot/jumpstarter_driver_uboot/client.py

@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 from functools import cached_property
 
 import pexpect
@@ -11,6 +12,7 @@ class UbootConsoleClient(CompositeClient):
     def prompt(self) -> str:
         return self.call("get_prompt")
 
+    @contextmanager
     def reboot_to_console(self) -> None:
         """
         Reboot to U-Boot console
@@ -22,6 +24,7 @@ def reboot_to_console(self) -> None:
         self.power.cycle()
 
         self.logger.info("Waiting for U-Boot prompt...")
+
         with self.serial.pexpect() as p:
             for _ in range(100):  # TODO: configurable retries
                 try:
@@ -30,21 +33,30 @@ def reboot_to_console(self) -> None:
                 except pexpect.TIMEOUT:
                     continue
 
-                return
-
-            raise RuntimeError("Failed to get U-Boot prompt")
-
-    def run_command(self, cmd: str, timeout: int = 60) -> bytes:
-        self.logger.info(f"Running command: {cmd}")
-        with self.serial.pexpect() as p:
-            p.sendline("")
-            p.expect_exact(self.prompt, timeout=timeout)
-            p.sendline(cmd)
-            p.expect_exact(self.prompt, timeout=timeout)
-            return p.before
+                break
+            else:
+                raise RuntimeError("Failed to get U-Boot prompt")
+
+            self.p = p
+            try:
+                yield
+            finally:
+                delattr(self, "p")
+
+    def run_command(self, cmd: str, timeout: int = 60, *, _internal_log=True) -> bytes:
+        if _internal_log:
+            self.logger.info(f"Running command: {cmd}")
+        if not hasattr(self, "p"):
+            raise RuntimeError("Not in a reboot_to_console context")
+        self.p.sendline("")
+        self.p.expect_exact(self.prompt, timeout=timeout)
+        self.p.sendline(cmd)
+        self.p.expect_exact(self.prompt, timeout=timeout)
+        return self.p.before
 
     def run_command_checked(self, cmd: str, timeout: int = 60, check=True) -> list[str]:
-        output = self.run_command("{}; echo $?".format(cmd))
+        self.logger.info(f"Running command checked: {cmd}")
+        output = self.run_command("{}; echo $?".format(cmd), _internal_log=False)
         parsed = output.strip().decode().splitlines()
 
         if len(parsed) < 2:

packages/jumpstarter-driver-uboot/jumpstarter_driver_uboot/driver_test.py

@@ -46,33 +46,33 @@ def test_driver_uboot_console(uboot_image):
         root.qemu.flasher.flash(uboot_image, partition="bios")
 
         uboot = root.uboot
-        uboot.reboot_to_console()
 
-        assert uboot.run_command_checked("version") == [
-            "U-Boot 2024.10 (Oct 11 2024 - 00:00:00 +0000)",
-            "",
-        ]
+        with uboot.reboot_to_console():
+            assert uboot.run_command_checked("version") == [
+                "U-Boot 2024.10 (Oct 11 2024 - 00:00:00 +0000)",
+                "",
+            ]
 
-        print(uboot.setup_dhcp())
+            print(uboot.setup_dhcp())
 
-        uboot.set_env_dict(
-            {
-                "foo": "bar",
-                "baz": "qux",
-            }
-        )
+            uboot.set_env_dict(
+                {
+                    "foo": "bar",
+                    "baz": "qux",
+                }
+            )
 
-        assert uboot.get_env("foo") == "bar"
-        assert uboot.get_env("baz") == "qux"
+            assert uboot.get_env("foo") == "bar"
+            assert uboot.get_env("baz") == "qux"
 
-        uboot.set_env_dict(
-            {
-                "foo": "qux",
-                "baz": None,
-            }
-        )
+            uboot.set_env_dict(
+                {
+                    "foo": "qux",
+                    "baz": None,
+                }
+            )
 
-        assert uboot.get_env("foo") == "qux"
-        assert uboot.get_env("baz") is None
+            assert uboot.get_env("foo") == "qux"
+            assert uboot.get_env("baz") is None
 
-        root.qemu.power.off()
+            root.qemu.power.off()