Make use of standalone uboot driver

Author: NickCao

packages/jumpstarter-driver-flashers/jumpstarter_driver_flashers/client.py

@@ -20,7 +20,6 @@
 
 from jumpstarter_driver_flashers.bundle import FlasherBundleManifestV1Alpha1
 
-from .uboot import UbootConsole
 from jumpstarter.common.exceptions import ArgumentError
 
 debug_console_option = click.option("--console-debug", is_flag=True, help="Enable console debug mode")
@@ -42,6 +41,7 @@ def __post_init__(self):
     def set_console_debug(self, debug: bool):
         """Set console debug mode"""
         self._console_debug = debug
+        # TODO: also set console debug on uboot client
 
     @contextmanager
     def busybox_shell(self):
@@ -61,12 +61,8 @@ def bootloader_shell(self):
         self.logger.info("Setting up flasher bundle files in exporter")
         self.call("setup_flasher_bundle")
         with self._services_up():
-            with self.serial.pexpect() as console:
-                if self._console_debug:
-                    console.logfile_read = sys.stdout.buffer
-                uboot = UbootConsole(console=console, power=self.power, logger=self.logger)
-                uboot.reboot_to_console()
-                console.sendline("")
+            with self.uboot.reboot_to_console():
+                pass
             yield self.serial
 
     def flash(
@@ -377,36 +373,37 @@ def _busybox(self):
 
         This is a helper context manager that boots the device into uboot and returns a console object.
         """
-        with self.serial.pexpect() as console:
-            if self._console_debug:
-                console.logfile_read = sys.stdout.buffer
-            uboot = UbootConsole(console=console, power=self.power, logger=self.logger)
-            # make sure that the device is booted into the uboot console
-            uboot.reboot_to_console()
+
+        # make sure that the device is booted into the uboot console
+        with self.uboot.reboot_to_console():
             # run dhcp discovery and gather details useful for later
-            self._dhcp_details = uboot.setup_dhcp()
+            self._dhcp_details = self.uboot.setup_dhcp()
             self.logger.info(f"discovered dhcp details: {self._dhcp_details}")
 
             # configure the environment necessary
             env = self._generate_uboot_env()
-            uboot.set_env_dict(env)
+            self.uboot.set_env_dict(env)
 
             # load any necessary files to RAM from the tftp storage
             manifest = self.manifest
             kernel_filename = Path(manifest.get_kernel_file()).name
             kernel_address = manifest.get_kernel_address()
 
-            uboot.run_command(f"tftpboot {kernel_address} {kernel_filename}", timeout=120)
+            self.uboot.run_command(f"tftpboot {kernel_address} {kernel_filename}", timeout=120)
 
             if manifest.get_initram_file():
                 initram_filename = Path(manifest.get_initram_file()).name
                 initram_address = manifest.get_initram_address()
-                uboot.run_command(f"tftpboot {initram_address} {initram_filename}", timeout=120)
+                self.uboot.run_command(f"tftpboot {initram_address} {initram_filename}", timeout=120)
 
             if manifest.get_dtb_file():
                 dtb_filename = Path(manifest.get_dtb_file()).name
                 dtb_address = manifest.get_dtb_address()
-                uboot.run_command(f"tftpboot {dtb_address} {dtb_filename}", timeout=120)
+                self.uboot.run_command(f"tftpboot {dtb_address} {dtb_filename}", timeout=120)
+
+        with self.serial.pexpect() as console:
+            if self._console_debug:
+                console.logfile_read = sys.stdout.buffer
 
             self.logger.info(f"Running boot command: {manifest.spec.bootcmd}")
             console.send(manifest.spec.bootcmd + "\n")

packages/jumpstarter-driver-flashers/jumpstarter_driver_flashers/driver.py

@@ -4,6 +4,7 @@
 import anyio.to_thread
 from jumpstarter_driver_http.driver import HttpServer
 from jumpstarter_driver_tftp.driver import Tftp
+from jumpstarter_driver_uboot.driver import UbootConsole
 from oras.provider import Registry
 
 from .bundle import FlasherBundleManifestV1Alpha1
@@ -45,6 +46,14 @@ def __post_init__(self):
                 "'power' instance is required for BaseFlasher either via a ref ir a direct child instance"
             )
 
+        if "uboot" not in self.children:
+            self.children["uboot"] = UbootConsole(
+                children={
+                    "power": self.children["power"],
+                    "serial": self.children["serial"],
+                }
+            )
+
         # bundles that have already been downloaded in the current session
         self._downloaded = {}
         self._use_dtb = None  # use default dtb unless set by client

packages/jumpstarter-driver-flashers/jumpstarter_driver_flashers/uboot.py

@@ -18,6 +18,7 @@ dependencies = [
     "jumpstarter-driver-http",
     "jumpstarter-driver-tftp",
     "jumpstarter-driver-power",
+    "jumpstarter-driver-uboot",
 ]
 
 [tool.hatch.version]
@@ -33,6 +34,9 @@ addopts = "--cov --cov-report=html --cov-report=xml"
 log_cli = true
 log_cli_level = "INFO"
 testpaths = ["jumpstarter_driver_flashers"]
+
+[tool.uv.sources]
+jumpstarter-driver-uboot = { workspace = true }
 #asyncio_mode = "auto"
 
 [build-system]

packages/jumpstarter-driver-flashers/pyproject.toml

@@ -44,23 +44,20 @@ async def aclose(self):
 def fs_operator_for_path(path: PathBuf) -> tuple[PathBuf, Operator]:
     return Path(path).resolve(), Operator("fs", root="/")
 
+
 def operator_for_path(path: PathBuf) -> tuple[PathBuf, Operator, str]:
-    """ Create an operator for the given path
+    """Create an operator for the given path
     Return a tuple of:
         - the path
         - the operator for the given path
         - the scheme of the operator.
     """
-    if type(path) is str and path.startswith(('http://', 'https://')):
-            parsed_url = urlparse(path)
-            operator = Operator(
-                'http',
-                root='/',
-                endpoint=f"{parsed_url.scheme}://{parsed_url.netloc}"
-            )
-            return Path(parsed_url.path), operator, 'http'
+    if type(path) is str and path.startswith(("http://", "https://")):
+        parsed_url = urlparse(path)
+        operator = Operator("http", root="/", endpoint=f"{parsed_url.scheme}://{parsed_url.netloc}")
+        return Path(parsed_url.path), operator, "http"
     else:
-        return *fs_operator_for_path(path), 'fs'
+        return *fs_operator_for_path(path), "fs"
 
 
 @dataclass(kw_only=True)
@@ -529,12 +526,7 @@ def flash(
         ...
 
     @abstractmethod
-    def dump(
-        self,
-        path: PathBuf,
-        *,
-        partition: str | None = None,
-        operator: Operator | None = None):
+    def dump(self, path: PathBuf, *, partition: str | None = None, operator: Operator | None = None):
         """Dump image from DUT"""
         ...
 

packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/client.py

@@ -195,6 +195,7 @@ async def copy_exporter_file(self, /, source: Path, target: str):
                         break
                     await dst.write(bs=data)
 
+
 class FlasherInterface(metaclass=ABCMeta):
     @classmethod
     def client(cls) -> str: