Merge branch 'main' into no-ti-j784s4-preflash

Author: NickCao

packages/jumpstarter-cli-admin/jumpstarter_cli_admin/get_test.py

@@ -219,6 +219,10 @@ async def test_get_client(_load_kube_config_mock, get_client_mock: AsyncMock):
 
 """
 
+CLIENTS_LIST_NAME = """client.jumpstarter.dev/test
+client.jumpstarter.dev/another
+"""
+
 CLIENTS_LIST_EMPTY_YAML = """apiVersion: jumpstarter.dev/v1alpha1
 items: []
 kind: ClientList
@@ -259,7 +263,7 @@ async def test_get_clients(_load_kube_config_mock, list_clients_mock: AsyncMock)
     list_clients_mock.return_value = CLIENTS_LIST
     result = await runner.invoke(get, ["clients", "--output", "name"])
     assert result.exit_code == 0
-    assert result.output == "client.jumpstarter.dev/test\n"
+    assert result.output == CLIENTS_LIST_NAME
     list_clients_mock.reset_mock()
 
     # No clients found
@@ -590,6 +594,10 @@ async def test_get_exporter_devices(_load_kube_config_mock, get_exporter_mock: A
 
 """
 
+EXPORTERS_LIST_NAME = """exporter.jumpstarter.dev/test
+exporter.jumpstarter.dev/another
+"""
+
 
 @pytest.mark.anyio
 @patch.object(ExportersV1Alpha1Api, "list_exporters")
@@ -623,7 +631,7 @@ async def test_get_exporters(_load_kube_config_mock, list_exporters_mock: AsyncM
     list_exporters_mock.return_value = EXPORTERS_LIST
     result = await runner.invoke(get, ["exporters", "--output", "name"])
     assert result.exit_code == 0
-    assert result.output == "exporter.jumpstarter.dev/test\n"
+    assert result.output == EXPORTERS_LIST_NAME
     list_exporters_mock.reset_mock()
 
     # No exporters found
@@ -754,6 +762,7 @@ async def test_get_exporters(_load_kube_config_mock, list_exporters_mock: AsyncM
 
 """
 
+EXPORTERS_DEVICES_LIST_NAME = EXPORTERS_LIST_NAME
 
 @pytest.mark.anyio
 @patch.object(ExportersV1Alpha1Api, "list_exporters")
@@ -790,7 +799,7 @@ async def test_get_exporters_devices(_load_kube_config_mock, list_exporters_mock
     list_exporters_mock.return_value = EXPORTER_DEVICES_LIST
     result = await runner.invoke(get, ["exporters", "--devices", "--output", "name"])
     assert result.exit_code == 0
-    assert result.output == "exporter.jumpstarter.dev/test\n"
+    assert result.output == EXPORTERS_DEVICES_LIST_NAME
     list_exporters_mock.reset_mock()
 
     # No exporters found
@@ -1128,6 +1137,9 @@ async def test_get_lease(_load_kube_config_mock, get_lease_mock: AsyncMock):
 
 """
 
+LEASES_LIST_NAME = """lease.jumpstarter.dev/82a8ac0d-d7ff-4009-8948-18a3c5c607b1
+lease.jumpstarter.dev/82a8ac0d-d7ff-4009-8948-18a3c5c607b2
+"""
 
 @pytest.mark.anyio
 @patch.object(LeasesV1Alpha1Api, "list_leases")
@@ -1172,7 +1184,7 @@ async def test_get_leases(_load_kube_config_mock, list_leases_mock: AsyncMock):
     list_leases_mock.return_value = V1Alpha1LeaseList(items=[IN_PROGRESS_LEASE, FINISHED_LEASE])
     result = await runner.invoke(get, ["leases", "--output", "name"])
     assert result.exit_code == 0
-    assert result.output == "lease.jumpstarter.dev/82a8ac0d-d7ff-4009-8948-18a3c5c607b1\n"
+    assert result.output == LEASES_LIST_NAME
     list_leases_mock.reset_mock()
 
     # No leases found

packages/jumpstarter-cli-admin/jumpstarter_cli_admin/print.py

@@ -40,8 +40,8 @@ def print_clients(clients: V1Alpha1List[V1Alpha1Client], namespace: str, output:
     elif output == OutputMode.YAML:
         click.echo(clients.dump_yaml())
     elif output == OutputMode.NAME:
-        if len(clients.items) > 0:
-            click.echo(f"client.jumpstarter.dev/{clients.items[0].metadata.name}")
+        for item in clients.items:
+            click.echo(f"client.jumpstarter.dev/{item.metadata.name}")
     elif len(clients.items) == 0:
         raise click.ClickException(f'No resources found in "{namespace}" namespace')
     else:
@@ -104,8 +104,8 @@ def print_exporters(exporters: V1Alpha1List[V1Alpha1Exporter], namespace: str, d
     elif output == OutputMode.YAML:
         click.echo(exporters.dump_yaml())
     elif output == OutputMode.NAME:
-        if len(exporters.items) > 0:
-            click.echo(f"exporter.jumpstarter.dev/{exporters.items[0].metadata.name}")
+        for item in exporters.items:
+            click.echo(f"exporter.jumpstarter.dev/{item.metadata.name}")
     elif len(exporters.items) == 0:
         raise click.ClickException(f'No resources found in "{namespace}" namespace')
     elif devices:
@@ -168,8 +168,8 @@ def print_leases(leases: V1Alpha1List[V1Alpha1Lease], namespace: str, output: Ou
     elif output == OutputMode.YAML:
         click.echo(leases.dump_yaml())
     elif output == OutputMode.NAME:
-        if len(leases.items) > 0:
-            click.echo(f"lease.jumpstarter.dev/{leases.items[0].metadata.name}")
+        for item in leases.items:
+            click.echo(f"lease.jumpstarter.dev/{item.metadata.name}")
     elif len(leases.items) == 0:
         raise click.ClickException(f'No resources found in "{namespace}" namespace')
     else:

packages/jumpstarter-driver-dutlink/jumpstarter_driver_dutlink/driver.py

@@ -1,28 +1,30 @@
 from __future__ import annotations
 
-import os
 import time
 from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field
 
 import pyudev
 import usb.core
 import usb.util
-from anyio import fail_after, sleep
-from anyio.streams.file import FileReadStream, FileWriteStream
+from anyio import sleep
 from jumpstarter_driver_composite.driver import CompositeInterface
 from jumpstarter_driver_opendal.driver import StorageMuxFlasherInterface
 from jumpstarter_driver_power.driver import PowerInterface, PowerReading
 from jumpstarter_driver_pyserial.driver import PySerial
 from serial.serialutil import SerialException
 
+from jumpstarter.common.storage import read_from_storage_device, write_to_storage_device
 from jumpstarter.driver import Driver, export
 
 
 @dataclass(kw_only=True)
 class DutlinkConfig:
     serial: str | None = field(default=None)
     timeout_s: int = field(default=20)  # 20 seconds, power control sequences can block USB for a long time
+    storage_timeout: int = field(default=10)
+    storage_leeway: int = field(default=6)
+    storage_fsync_timeout: int = field(default=900)
 
     dev: usb.core.Device = field(init=False)
     itf: usb.core.Interface = field(init=False)
@@ -187,44 +189,29 @@ def dut(self):
     def off(self):
         return self.control("off")
 
-    async def wait_for_storage_device(self):
-        with fail_after(20):
-            while True:
-                self.logger.debug(f"waiting for storage device {self.storage_device}")
-                if os.path.exists(self.storage_device):
-                    self.logger.debug(f"storage device {self.storage_device} is ready")
-                    # https://stackoverflow.com/a/2774125
-                    fd = os.open(self.storage_device, os.O_WRONLY)
-                    try:
-                        if os.lseek(fd, 0, os.SEEK_END) > 0:
-                            break
-                    finally:
-                        os.close(fd)
-                await sleep(1)
-
     @export
     async def write(self, src: str):
         self.host()
-        await self.wait_for_storage_device()
-        async with await FileWriteStream.from_path(self.storage_device) as stream:
-            async with self.resource(src) as res:
-                total_bytes = 0
-                next_print = 0
-                async for chunk in res:
-                    await stream.send(chunk)
-                    if total_bytes > next_print:
-                        self.logger.debug(f"{self.storage_device} written {total_bytes / (1024 * 1024)} MB")
-                        next_print += 50 * 1024 * 1024
-                    total_bytes += len(chunk)
+        async with self.resource(src) as res:
+            await write_to_storage_device(
+                self.storage_device,
+                res,
+                timeout=self.storage_timeout,
+                leeway=self.storage_leeway,
+                fsync_timeout=self.storage_fsync_timeout,
+                logger=self.logger,
+            )
 
     @export
     async def read(self, dst: str):
         self.host()
-        await self.wait_for_storage_device()
-        async with await FileReadStream.from_path(self.storage_device) as stream:
-            async with self.resource(dst) as res:
-                async for chunk in stream:
-                    await res.send(chunk)
+        async with self.resource(dst) as res:
+            await read_from_storage_device(
+                self.storage_device,
+                res,
+                timeout=self.storage_timeout,
+                logger=self.logger,
+            )
 
 
 @dataclass(kw_only=True)

packages/jumpstarter-driver-sdwire/jumpstarter_driver_sdwire/driver.py

@@ -1,15 +1,13 @@
 from __future__ import annotations
 
-import os
 from dataclasses import dataclass, field
 
 import pyudev
 import usb.core
 import usb.util
-from anyio import fail_after, sleep
-from anyio.streams.file import FileReadStream, FileWriteStream
 from jumpstarter_driver_opendal.driver import StorageMuxFlasherInterface
 
+from jumpstarter.common.storage import read_from_storage_device, write_to_storage_device
 from jumpstarter.driver import Driver, export
 
 
@@ -20,6 +18,9 @@ class SDWire(StorageMuxFlasherInterface, Driver):
     itf: usb.core.Interface = field(init=False)
 
     storage_device: str | None = field(default=None)
+    storage_timeout: int = field(default=10)
+    storage_leeway: int = field(default=6)
+    storage_fsync_timeout: int = field(default=900)
 
     def effective_storage_device(self):
         if self.storage_device is None:
@@ -101,45 +102,26 @@ def dut(self):
     def off(self):
         self.host()
 
-    async def wait_for_storage_device(self):
-        with fail_after(10):
-            storage_device = self.effective_storage_device()
-
-            while True:
-                # https://stackoverflow.com/a/2774125
-                try:
-                    fd = os.open(storage_device, os.O_WRONLY)
-                    if os.lseek(fd, 0, os.SEEK_END) > 0:
-                        break
-                except OSError as e:
-                    match e.errno:
-                        case 123:  # No medium found
-                            pass
-                        case 5:  # Input/output error
-                            pass
-                        case _:
-                            raise
-                finally:
-                    if "fd" in locals():
-                        os.close(fd)
-                await sleep(1)
-
-            return storage_device
-
     @export
     async def write(self, src: str):
         self.host()
-        storage_device = await self.wait_for_storage_device()
-        async with await FileWriteStream.from_path(storage_device) as stream:
-            async with self.resource(src) as res:
-                async for chunk in res:
-                    await stream.send(chunk)
+        async with self.resource(src) as res:
+            await write_to_storage_device(
+                self.effective_storage_device(),
+                res,
+                timeout=self.storage_timeout,
+                leeway=self.storage_leeway,
+                fsync_timeout=self.storage_fsync_timeout,
+                logger=self.logger,
+            )
 
     @export
     async def read(self, dst: str):
         self.host()
-        storage_device = await self.wait_for_storage_device()
-        async with await FileReadStream.from_path(storage_device) as stream:
-            async with self.resource(dst) as res:
-                async for chunk in stream:
-                    await res.send(chunk)
+        async with self.resource(dst) as res:
+            await read_from_storage_device(
+                self.effective_storage_device(),
+                res,
+                timeout=self.storage_timeout,
+                logger=self.logger,
+            )