diff --git a/python/packages/jumpstarter-driver-pyserial/README.md b/python/packages/jumpstarter-driver-pyserial/README.md index 806775e0d..6ba16de61 100644 --- a/python/packages/jumpstarter-driver-pyserial/README.md +++ b/python/packages/jumpstarter-driver-pyserial/README.md @@ -46,6 +46,8 @@ export: | check_present | Check if the serial port exists during exporter initialization, disable if you are connecting to a dynamically created port (i.e. USB from your DUT) | bool | no | True | | cps | Characters per second throttling limit. When set, data transmission will be throttled to simulate slow typing. Useful for devices that can't handle fast input | float | no | None | | disable_hupcl | Disable HUPCL on POSIX systems to avoid toggling DTR/RTS on close (can prevent MCU reset on serial disconnect) | bool | no | False | +| power_control_ref | Explicit power device name from DUT tree for Ctrl-] x3 hotkey (skips auto-discovery). Only needed in multi-power setups | str | no | None (auto-discover) | +| power_control_method | Power cycle method sequence for Ctrl-] x3 hotkey. Supports method names (`cycle`, `reset`, `on`, `off`) and `sleep:N` delays. Set to `[]` or `null` to disable | list[str] | no | `["cycle"]` | ### NVDemuxSerial Driver @@ -161,7 +163,9 @@ Start an interactive serial console with direct terminal access. j serial start-console ``` -Exit the console by pressing CTRL+B three times. +**Hotkeys:** +- **CTRL+B x3**: Exit the console +- **CTRL+] x3**: Power cycle the board (if a power driver is available in the DUT tree) ### pipe diff --git a/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/client.py b/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/client.py index af7481f0f..6ca5ff425 100644 --- a/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/client.py +++ b/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/client.py @@ -1,17 +1,29 @@ +import logging import sys +import time from contextlib import contextmanager from typing import Optional import click -from anyio import BrokenResourceError, EndOfStream, create_task_group, open_file +from anyio import BrokenResourceError, EndOfStream, create_task_group, open_file, sleep, to_thread from anyio.streams.file import FileReadStream from jumpstarter_driver_network.adapters import PexpectAdapter from pexpect.fdpexpect import fdspawn -from .console import Console +from .console import Console, ConsoleStreamDrop from jumpstarter.client import DriverClient from jumpstarter.client.decorators import driver_click_group +logger = logging.getLogger(__name__) + +KNOWN_POWER_CLIENTS = frozenset({ + "jumpstarter_driver_power.client.PowerClient", + "jumpstarter_driver_power.client.VirtualPowerClient", + "jumpstarter_driver_ridesx.client.RideSXPowerClient", + "jumpstarter_driver_noyito_relay.client.NoyitoPowerClient", + "jumpstarter_driver_snmp.client.SNMPServerClient", +}) + class PySerialClient(DriverClient): """ @@ -125,6 +137,94 @@ async def _stdin_to_serial(self, stream) -> tuple[int, int]: return bytes_read, bytes_sent + def _find_power_client(self): + # Check if hotkey is disabled + method_label = self.labels.get("jumpstarter.dev/pyserial/power-control-method", "cycle") + if not method_label: + return None + + root = getattr(self, 'root', None) + if root is None: + return None + + # Explicit ref takes precedence + ref_label = self.labels.get("jumpstarter.dev/pyserial/power-control-ref") + if ref_label: + power_client = root.children.get(ref_label) + if power_client is None: + logger.warning( + "power_control_ref '%s' not found in DUT tree — power cycle hotkey disabled", + ref_label + ) + return None + return power_client + + # Auto-discovery: collect all power-capable clients + candidates = [] + self._collect_power_clients(root, candidates) + + if len(candidates) == 0: + return None + if len(candidates) == 1: + return candidates[0] + + # Multiple candidates — ambiguous + names = [c.labels.get("jumpstarter.dev/name", "unknown") for c in candidates] + logger.warning( + "Multiple power drivers found (%s) — power cycle hotkey disabled. " + "Set power_control_ref to select one explicitly.", + ", ".join(names) + ) + return None + + def _collect_power_clients(self, client, result, seen_uuids=None): + if seen_uuids is None: + seen_uuids = set() + client_uuid = getattr(client, 'uuid', None) + if client_uuid and client_uuid in seen_uuids: + return + if client_uuid: + seen_uuids.add(client_uuid) + client_class = client.labels.get("jumpstarter.dev/client") + if client_class in KNOWN_POWER_CLIENTS: + result.append(client) + for child in client.children.values(): + self._collect_power_clients(child, result, seen_uuids) + + def _make_power_cycle(self, power_client): + method_label = self.labels.get("jumpstarter.dev/pyserial/power-control-method", "cycle") + methods = [m for m in method_label.split(",") if m] + + # Pre-validate and parse all steps + steps = [] + for method in methods: + if method.startswith("sleep:"): + try: + delay = float(method.split(":", 1)[1]) + steps.append(("sleep", delay)) + except (ValueError, IndexError): + logger.warning("Invalid sleep step '%s' — power cycle hotkey disabled", method) + return None + else: + operation = getattr(power_client, method, None) + if callable(operation): + steps.append(("operation", operation)) + else: + logger.warning( + "Power client does not have callable method '%s' — power cycle hotkey disabled", + method + ) + return None + + async def _cycle(): + for kind, value in steps: + if kind == "sleep": + await sleep(value) + else: + await to_thread.run_sync(value) + + return _cycle + def cli(self): # noqa: C901 @driver_click_group(self) def base(): @@ -134,9 +234,23 @@ def base(): @base.command() def start_console(): """Start serial port console""" + power_client = self._find_power_client() + on_power_cycle = self._make_power_cycle(power_client) if power_client is not None else None click.echo("\nStarting serial port console ... exit with CTRL+B x 3 times\n") - console = Console(serial_client=self) - console.run() + if on_power_cycle is not None: + click.echo("Power cycle: CTRL+] x 3 times\n") + retries = 0 + while retries < 30: + console = Console(serial_client=self, on_power_cycle=on_power_cycle) + try: + console.run() + break + except ConsoleStreamDrop: + click.echo("\r\nSerial connection lost, reconnecting...\n", err=True) + retries += 1 + time.sleep(1) + else: + click.echo("\nSerial connection lost (reconnect attempts exhausted).\n", err=True) @base.command() @click.option( diff --git a/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/client_test.py b/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/client_test.py new file mode 100644 index 000000000..43d0bcefd --- /dev/null +++ b/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/client_test.py @@ -0,0 +1,188 @@ +import threading +from unittest.mock import MagicMock + +from .driver import PySerial +from jumpstarter.common.utils import serve + + +def test_find_power_client_no_root(): + with serve(PySerial(url="loop://", power_control_method=["cycle"])) as client: + # No root attribute set → should return None + assert client._find_power_client() is None + + +def test_find_power_client_auto_discover(): + power = MagicMock(spec=["cycle", "children", "labels"]) + power.children = {} + power.labels = {"jumpstarter.dev/client": "jumpstarter_driver_power.client.PowerClient"} + root = MagicMock(spec=["children", "labels"]) + root.children = {"power": power} + root.labels = {} + + with serve(PySerial(url="loop://")) as client: + object.__setattr__(client, "root", root) + assert client._find_power_client() is power + + +def test_make_power_cycle_calls_cycle(): + called = threading.Event() + power = MagicMock() + power.cycle = MagicMock(side_effect=lambda: called.set()) + + with serve(PySerial(url="loop://")) as client: + cycle_fn = client._make_power_cycle(power) + client.portal.call(cycle_fn) + assert called.is_set() + + +def test_find_power_client_ambiguous(): + power1 = MagicMock(spec=["children", "labels"]) + power1.children = {} + power1.labels = { + "jumpstarter.dev/client": "jumpstarter_driver_power.client.PowerClient", + "jumpstarter.dev/name": "power1", + } + power2 = MagicMock(spec=["children", "labels"]) + power2.children = {} + power2.labels = { + "jumpstarter.dev/client": "jumpstarter_driver_power.client.PowerClient", + "jumpstarter.dev/name": "power2", + } + root = MagicMock(spec=["children", "labels"]) + root.children = {"power1": power1, "power2": power2} + root.labels = {} + + with serve(PySerial(url="loop://")) as client: + object.__setattr__(client, "root", root) + assert client._find_power_client() is None + + +def test_find_power_client_explicit_ref(): + power = MagicMock(spec=["children", "labels"]) + power.children = {} + power.labels = {"jumpstarter.dev/client": "jumpstarter_driver_power.client.PowerClient"} + other = MagicMock(spec=["children", "labels"]) + other.children = {} + other.labels = {"jumpstarter.dev/client": "other.driver.Client"} + root = MagicMock(spec=["children", "labels"]) + root.children = {"power": power, "other": other} + root.labels = {} + + with serve(PySerial(url="loop://", power_control_ref="power")) as client: + object.__setattr__(client, "root", root) + assert client._find_power_client() is power + + +def test_find_power_client_explicit_ref_missing(): + root = MagicMock(spec=["children", "labels"]) + root.children = {} + root.labels = {} + + with serve(PySerial(url="loop://", power_control_ref="nonexistent")) as client: + object.__setattr__(client, "root", root) + assert client._find_power_client() is None + + +def test_find_power_client_non_power_ignored(): + gpio = MagicMock(spec=["children", "labels"]) + gpio.children = {} + gpio.labels = {"jumpstarter.dev/client": "jumpstarter_driver_gpiod.client.DigitalOutputClient"} + root = MagicMock(spec=["children", "labels"]) + root.children = {"gpio": gpio} + root.labels = {} + + with serve(PySerial(url="loop://")) as client: + object.__setattr__(client, "root", root) + assert client._find_power_client() is None + + +def test_make_power_cycle_custom_method(): + called_sequence = [] + power = MagicMock() + power.off = MagicMock(side_effect=lambda: called_sequence.append("off")) + power.on = MagicMock(side_effect=lambda: called_sequence.append("on")) + + with serve(PySerial(url="loop://", power_control_method=["off", "on"])) as client: + cycle_fn = client._make_power_cycle(power) + client.portal.call(cycle_fn) + assert called_sequence == ["off", "on"] + + +def test_find_power_client_disabled_via_empty_list(): + power = MagicMock(spec=["children", "labels"]) + power.children = {} + power.labels = {"jumpstarter.dev/client": "jumpstarter_driver_power.client.PowerClient"} + root = MagicMock(spec=["children", "labels"]) + root.children = {"power": power} + root.labels = {} + + with serve(PySerial(url="loop://", power_control_method=[])) as client: + object.__setattr__(client, "root", root) + assert client._find_power_client() is None + + +def test_make_power_cycle_missing_method(): + power = MagicMock(spec=["children", "labels"]) + power.children = {} + power.labels = {} + + with serve(PySerial(url="loop://", power_control_method=["nonexistent_method"])) as client: + result = client._make_power_cycle(power) + assert result is None + + +def test_make_power_cycle_with_sleep(): + called_sequence = [] + power = MagicMock() + power.off = MagicMock(side_effect=lambda: called_sequence.append("off")) + power.on = MagicMock(side_effect=lambda: called_sequence.append("on")) + + with serve(PySerial(url="loop://", power_control_method=["off", "sleep:0.01", "on"])) as client: + cycle_fn = client._make_power_cycle(power) + client.portal.call(cycle_fn) + assert called_sequence == ["off", "on"] + + +def test_find_power_client_disabled_via_none(): + power = MagicMock(spec=["children", "labels"]) + power.children = {} + power.labels = {"jumpstarter.dev/client": "jumpstarter_driver_power.client.PowerClient"} + root = MagicMock(spec=["children", "labels"]) + root.children = {"power": power} + root.labels = {} + + with serve(PySerial(url="loop://", power_control_method=None)) as client: + object.__setattr__(client, "root", root) + assert client._find_power_client() is None + + +def test_collect_power_clients_dedup_proxy(): + # Simulate Proxy scenario: same power driver instance appears twice in tree + # (once via proxy delegation, once via direct parent) + from uuid import uuid4 + shared_uuid = uuid4() + + power1 = MagicMock(spec=["children", "labels", "uuid"]) + power1.children = {} + power1.labels = { + "jumpstarter.dev/client": "jumpstarter_driver_power.client.PowerClient", + "jumpstarter.dev/name": "power", + } + power1.uuid = shared_uuid + + power2 = MagicMock(spec=["children", "labels", "uuid"]) + power2.children = {} + power2.labels = { + "jumpstarter.dev/client": "jumpstarter_driver_power.client.PowerClient", + "jumpstarter.dev/name": "power", + } + power2.uuid = shared_uuid # Same UUID as power1 + + root = MagicMock(spec=["children", "labels"]) + root.children = {"power1": power1, "power2": power2} + root.labels = {} + + with serve(PySerial(url="loop://")) as client: + object.__setattr__(client, "root", root) + # Should find only one power client despite two references + assert client._find_power_client() is power1 diff --git a/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/console.py b/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/console.py index b315b2f55..ce52a15c4 100644 --- a/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/console.py +++ b/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/console.py @@ -1,21 +1,31 @@ +import logging import sys import termios import tty +from collections.abc import Awaitable, Callable from contextlib import contextmanager -from anyio import create_task_group +from anyio import EndOfStream, create_task_group from anyio.streams.file import FileReadStream, FileWriteStream from jumpstarter.client import DriverClient +logger = logging.getLogger(__name__) + class ConsoleExit(Exception): pass +class ConsoleStreamDrop(Exception): + """Serial stream dropped; caller may reconnect.""" + pass + + class Console: - def __init__(self, serial_client: DriverClient): + def __init__(self, serial_client: DriverClient, on_power_cycle: Callable[[], Awaitable[None]] | None = None): self.serial_client = serial_client + self.on_power_cycle = on_power_cycle def run(self): with self.setraw(): @@ -31,32 +41,53 @@ def setraw(self): termios.tcsetattr(sys.stdin.fileno(), termios.TCSADRAIN, original) async def __run(self): - async with self.serial_client.stream_async(method="connect") as stream: - try: - async with create_task_group() as tg: - tg.start_soon(self.__serial_to_stdout, stream) - tg.start_soon(self.__stdin_to_serial, stream) - except* ConsoleExit: - pass + try: + async with self.serial_client.stream_async(method="connect") as stream: + try: + async with create_task_group() as tg: + tg.start_soon(self.__serial_to_stdout, stream) + tg.start_soon(self.__stdin_to_serial, stream) + except* ConsoleExit: + pass + except* ConsoleStreamDrop: + raise ConsoleStreamDrop() from None + except EndOfStream: + raise ConsoleStreamDrop() from None async def __serial_to_stdout(self, stream): stdout = FileWriteStream(sys.stdout.buffer) - while True: - data = await stream.receive() - await stdout.send(data) - sys.stdout.flush() + try: + while True: + data = await stream.receive() + await stdout.send(data) + sys.stdout.flush() + except EndOfStream: + raise ConsoleStreamDrop() from None async def __stdin_to_serial(self, stream): stdin = FileReadStream(sys.stdin.buffer) ctrl_b_count = 0 + ctrl_bracket_count = 0 # Ctrl-] x3 triggers power cycle while True: data = await stdin.receive(max_bytes=1) if not data: continue if data == b"\x02": # Ctrl-B ctrl_b_count += 1 + ctrl_bracket_count = 0 if ctrl_b_count == 3: raise ConsoleExit + elif data == b"\x1d": # Ctrl-] + ctrl_bracket_count += 1 + ctrl_b_count = 0 + if ctrl_bracket_count == 3: + if self.on_power_cycle is not None: + await self.on_power_cycle() + else: + logger.warning("Power cycle hotkey pressed but no power driver available") + ctrl_bracket_count = 0 + continue else: ctrl_b_count = 0 + ctrl_bracket_count = 0 await stream.send(data) diff --git a/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/console_test.py b/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/console_test.py new file mode 100644 index 000000000..4e8f691d4 --- /dev/null +++ b/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/console_test.py @@ -0,0 +1,91 @@ +import os +import threading +import time +from unittest.mock import MagicMock, patch + +from .console import Console +from .driver import PySerial +from jumpstarter.common.utils import serve + + +def _start_console(client, on_power_cycle=None): + """Run Console.run() in a thread with a PTY substituted for stdin. + + Returns (master_fd, thread, result_dict). Write keypresses to master_fd; + the result dict gets an 'exc' key if the console thread raises. + """ + master_fd, slave_fd = os.openpty() + slave_file = os.fdopen(slave_fd, "rb", buffering=0) + + mock_stdin = MagicMock() + mock_stdin.fileno.return_value = slave_fd + mock_stdin.buffer = slave_file + + result = {} + + def _run(): + with patch("sys.stdin", mock_stdin): + console = Console(serial_client=client, on_power_cycle=on_power_cycle) + try: + console.run() + except Exception as e: + result["exc"] = e + slave_file.close() + + t = threading.Thread(target=_run, daemon=True) + t.start() + return master_fd, t, result + + +def test_ctrl_b_exits(): + with serve(PySerial(url="loop://")) as client: + master_fd, t, result = _start_console(client) + try: + time.sleep(0.1) + os.write(master_fd, b"a") + os.write(master_fd, b"\x02\x02\x02") + t.join(timeout=5) + finally: + os.close(master_fd) + + assert not t.is_alive(), "console did not exit after Ctrl-B x3" + assert "exc" not in result + + +def test_ctrl_bracket_triggers_power_cycle(): + power_cycled = threading.Event() + + async def on_power_cycle(): + power_cycled.set() + + with serve(PySerial(url="loop://")) as client: + master_fd, t, result = _start_console(client, on_power_cycle=on_power_cycle) + try: + time.sleep(0.1) + os.write(master_fd, b"\x1d\x1d\x1d") + assert power_cycled.wait(timeout=5), "power cycle was not triggered" + assert t.is_alive(), "console exited after power cycle" + os.write(master_fd, b"\x02\x02\x02") + t.join(timeout=5) + finally: + os.close(master_fd) + + assert not t.is_alive() + assert "exc" not in result + + +def test_ctrl_bracket_without_power_client(): + with serve(PySerial(url="loop://")) as client: + master_fd, t, result = _start_console(client, on_power_cycle=None) + try: + time.sleep(0.1) + os.write(master_fd, b"\x1d\x1d\x1d") + time.sleep(0.1) + assert t.is_alive(), "console exited unexpectedly on Ctrl-] without power client" + os.write(master_fd, b"\x02\x02\x02") + t.join(timeout=5) + finally: + os.close(master_fd) + + assert not t.is_alive() + assert "exc" not in result diff --git a/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/driver.py b/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/driver.py index 443e04af6..40bf8fd00 100644 --- a/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/driver.py +++ b/python/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/driver.py @@ -101,6 +101,18 @@ class PySerial(Driver): check_present: bool = field(default=True) cps: Optional[float] = field(default=None) # characters per second throttling disable_hupcl: bool = field(default=False) + power_control_ref: Optional[str] = field( + default=None, + metadata={"description": "Explicit power device name from DUT tree for Ctrl-] hotkey (skips auto-discovery)"}, + ) + power_control_method: list[str] | str | None = field( + default_factory=lambda: ["cycle"], + metadata={ + "description": "Power cycle method sequence for Ctrl-] hotkey. " + "Supports method names (cycle, reset, on, off) and sleep:N. " + "Set to [] or null to disable hotkey. Default: ['cycle']" + }, + ) _transport: Any = field(default=None, init=False, repr=False) def __post_init__(self): @@ -109,10 +121,29 @@ def __post_init__(self): if self.check_present and self.url != LOOP: serial_for_url(self.url, baudrate=self.baudrate) + # Normalize power_control_method: None stays None, string becomes list + if isinstance(self.power_control_method, str): + object.__setattr__(self, "power_control_method", [self.power_control_method]) + @classmethod def client(cls) -> str: return "jumpstarter_driver_pyserial.client.PySerialClient" + def extra_labels(self) -> dict[str, str]: + labels = {} + if self.power_control_ref is not None: + labels["jumpstarter.dev/pyserial/power-control-ref"] = self.power_control_ref + if self.power_control_method is not None: + method_list = ( + self.power_control_method + if isinstance(self.power_control_method, list) + else [self.power_control_method] + ) + labels["jumpstarter.dev/pyserial/power-control-method"] = ",".join(method_list) + else: + labels["jumpstarter.dev/pyserial/power-control-method"] = "" + return labels + def _maybe_disable_hupcl(self, serial_port: Any): """Disable HUPCL to avoid MCU reset on serial port close when supported.""" if not self.disable_hupcl or os.name != "posix" or termios is None: diff --git a/python/packages/jumpstarter/jumpstarter/client/client.py b/python/packages/jumpstarter/jumpstarter/client/client.py index 9a0539a1d..9b147e1b6 100644 --- a/python/packages/jumpstarter/jumpstarter/client/client.py +++ b/python/packages/jumpstarter/jumpstarter/client/client.py @@ -112,7 +112,6 @@ async def client_from_channel( stub = MultipathExporterStub([channel]) response = await stub.GetReport(empty_pb2.Empty()) - for index, report in enumerate(response.reports): topo[index] = [] @@ -149,4 +148,14 @@ async def client_from_channel( clients[index] = client + root_client = next(reversed(clients.values())) + + def _iter_all(client): + yield client + for child in client.children.values(): + yield from _iter_all(child) + + for c in _iter_all(root_client): + object.__setattr__(c, 'root', root_client) + return clients.popitem(last=True)[1]