customizable descriptions for all drivers

michalskrivanek · michalskrivanek · commit d02013ace86f · 2025-10-06T10:31:08.000+02:00
this allows override of the default strings to show up in the jmp shell CLI.
exporter:
  driver_x:
    type: x
    description: "this command does x"
diff --git a/packages/jumpstarter-driver-composite/jumpstarter_driver_composite/client.py b/packages/jumpstarter-driver-composite/jumpstarter_driver_composite/client.py
@@ -34,7 +34,7 @@ def close(self):
                 v.close()
 
     def cli(self):
-        @click.group
+        @click.group(help=self.description or "Generic composite device")
         @click.option(
             "--log-level",
             "log_level",
@@ -44,7 +44,6 @@ def cli(self):
             callback=_opt_log_level_callback,
         )
         def base():
-            """Generic composite device"""
             pass
 
         for k, v in self.children.items():
diff --git a/packages/jumpstarter-driver-flashers/jumpstarter_driver_flashers/client.py b/packages/jumpstarter-driver-flashers/jumpstarter_driver_flashers/client.py
@@ -750,9 +750,8 @@ def _validate_bearer_token(self, token: str | None) -> str | None:
         return token
 
     def cli(self):
-        @click.group
+        @click.group(help=self.description or "Software-defined flasher interface")
         def base():
-            """Software-defined flasher interface"""
             pass
 
         @base.command()
diff --git a/packages/jumpstarter-driver-gpiod/jumpstarter_driver_gpiod/client.py b/packages/jumpstarter-driver-gpiod/jumpstarter_driver_gpiod/client.py
@@ -34,9 +34,8 @@ def read(self):
         return PinState(int(self.call("read_pin")))
 
     def cli(self):
-        @click.group()
+        @click.group(help=self.description or "GPIO power control commands.")
         def gpio():
-            """GPIO power control commands."""
             pass
 
         for cmd in super().cli().commands.values():
@@ -79,9 +78,8 @@ def read(self):
         return PinState(int(self.call("read_pin")))
 
     def cli(self):
-        @click.group()
+        @click.group(help=self.description or "GPIO input commands.")
         def gpio():
-            """GPIO input commands."""
             pass
 
         @gpio.command()
diff --git a/packages/jumpstarter-driver-network/jumpstarter_driver_network/client.py b/packages/jumpstarter-driver-network/jumpstarter_driver_network/client.py
@@ -20,9 +20,8 @@ def address(self):
         return self.call("address")
 
     def cli(self):
-        @click.group
+        @click.group(help=self.description or "Generic Network Connection")
         def base():
-            """Generic Network Connection"""
             pass
 
         @base.command()
diff --git a/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/client.py b/packages/jumpstarter-driver-opendal/jumpstarter_driver_opendal/client.py
@@ -413,9 +413,9 @@ def cli(self):  # noqa: C901
         arg_dst = click.argument("dst", type=click.Path())
         opt_expire_second = click.option("--expire-second", type=int, required=True)
 
-        @click.group
+        @click.group(help=self.description or "Opendal Storage")
         def base():
-            """Opendal Storage"""
+            pass
 
         @base.command
         @arg_path
@@ -548,9 +548,8 @@ def dump(
         ...
 
     def cli(self):
-        @click.group
+        @click.group(help=self.description or "Generic flasher interface")
         def base():
-            """Generic flasher interface"""
             pass
 
         @base.command()
diff --git a/packages/jumpstarter-driver-power/jumpstarter_driver_power/client.py b/packages/jumpstarter-driver-power/jumpstarter_driver_power/client.py
@@ -35,9 +35,8 @@ def read(self) -> Generator[PowerReading, None, None]:
             yield PowerReading.model_validate(v, strict=True)
 
     def cli(self):
-        @click.group
+        @click.group(help=self.description or "Generic power")
         def base():
-            """Generic power"""
             pass
 
         @base.command()
diff --git a/packages/jumpstarter-driver-probe-rs/jumpstarter_driver_probe_rs/client.py b/packages/jumpstarter-driver-probe-rs/jumpstarter_driver_probe_rs/client.py
@@ -60,9 +60,8 @@ def read(self, width: int, address: int, words: int) -> list[int]:
         return [int(data, 16) for data in data_strs]
 
     def cli(self):  # noqa: C901
-        @click.group
+        @click.group(help=self.description or "probe-rs client")
         def base():
-            """probe-rs client"""
             pass
 
         @base.command()
diff --git a/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/client.py b/packages/jumpstarter-driver-pyserial/jumpstarter_driver_pyserial/client.py
@@ -36,9 +36,8 @@ def pexpect(self):
             yield adapter
 
     def cli(self):
-        @click.group
+        @click.group(help=self.description or "Serial port client")
         def base():
-            """Serial port client"""
             pass
 
         @base.command()
diff --git a/packages/jumpstarter-driver-ridesx/jumpstarter_driver_ridesx/client.py b/packages/jumpstarter-driver-ridesx/jumpstarter_driver_ridesx/client.py
@@ -102,7 +102,7 @@ def flash(
     def cli(self):
         generic_cli = FlasherClient.cli(self)
 
-        @click.group()
+        @click.group(help=self.description or "RideSX storage operations")
         def storage():
             pass
 
diff --git a/packages/jumpstarter-driver-shell/jumpstarter_driver_shell/client.py b/packages/jumpstarter-driver-shell/jumpstarter_driver_shell/client.py
@@ -41,9 +41,8 @@ def execute(*args, **kwargs):
 
     def cli(self):
         """Create CLI interface for dynamically configured shell methods"""
-        @click.group
+        @click.group(help=self.description or "Shell command executor")
         def base():
-            """Shell command executor"""
             pass
 
         # Get available methods from the driver
@@ -80,12 +79,12 @@ def method_command(args, env):
         except Exception:
             description = f"Execute the {method_name} shell method"
 
-        # Decorate and register the command
-        method_command.__doc__ = description
+        # Decorate and register the command with help text
         method_command = click.argument('args', nargs=-1, type=click.UNPROCESSED)(method_command)
         method_command = click.option('--env', '-e', multiple=True,
                      help='Environment variables in KEY=VALUE format')(method_command)
         method_command = group.command(
             name=method_name,
+            help=description,
             context_settings={"ignore_unknown_options": True, "allow_interspersed_args": False},
         )(method_command)
diff --git a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py
@@ -17,9 +17,8 @@ def off(self):
         self.call("off")
 
     def cli(self):
-        @click.group()
+        @click.group(help=self.description or "SNMP power control commands")
         def snmp():
-            """SNMP power control commands"""
             pass
 
         for cmd in super().cli().commands.values():
diff --git a/packages/jumpstarter/jumpstarter/client/base.py b/packages/jumpstarter/jumpstarter/client/base.py
@@ -33,6 +33,9 @@ class DriverClient(AsyncDriverClient):
     portal: BlockingPortal
     stack: ExitStack
 
+    description: str | None = None
+    """Driver description from GetReport(), used for CLI help text"""
+
     def call(self, method, *args):
         """
         Invoke driver call
diff --git a/packages/jumpstarter/jumpstarter/client/client.py b/packages/jumpstarter/jumpstarter/client/client.py
@@ -51,13 +51,15 @@ async def client_from_channel(
         report = reports[index]
 
         client_class = import_class(report.labels["jumpstarter.dev/client"], allow, unsafe)
+
         client = client_class(
             uuid=UUID(report.uuid),
             labels=report.labels,
             stub=stub,
             portal=portal,
             stack=stack.enter_context(ExitStack()),
             children={reports[k].labels["jumpstarter.dev/name"]: clients[k] for k in topo[index]},
+            description=getattr(report, 'description', None) or None,
         )
 
         clients[index] = client
diff --git a/packages/jumpstarter/jumpstarter/driver/base.py b/packages/jumpstarter/jumpstarter/driver/base.py
@@ -72,6 +72,9 @@ class Driver(
     resources: dict[UUID, Any] = field(default_factory=dict, init=False)
     """Dict of client side resources"""
 
+    description: str | None = None
+    """Custom description for the driver (shown in CLI help)"""
+
     log_level: str = "INFO"
     logger: logging.Logger = field(init=False)
 
@@ -206,6 +209,7 @@ def report(self, *, root=None, parent=None, name=None):
             | self.extra_labels()
             | ({"jumpstarter.dev/client": self.client()})
             | ({"jumpstarter.dev/name": name} if name else {}),
+            description=self.description or None,
         )
 
     def enumerate(self, *, root=None, parent=None, name=None):
diff --git a/packages/jumpstarter/jumpstarter/exporter/session.py b/packages/jumpstarter/jumpstarter/exporter/session.py
@@ -100,6 +100,7 @@ def __getitem__(self, key: UUID):
 
     async def GetReport(self, request, context):
         logger.debug("GetReport()")
+
         return jumpstarter_pb2.GetReportResponse(
             uuid=str(self.uuid),
             labels=self.labels,
diff --git a/packages/jumpstarter/jumpstarter/exporter/session_test.py b/packages/jumpstarter/jumpstarter/exporter/session_test.py
@@ -0,0 +1,141 @@
+"""Tests for session GetReport with descriptions"""
+
+from google.protobuf import empty_pb2
+
+from jumpstarter.common.utils import serve
+from jumpstarter.driver import Driver
+
+
+class SimpleDriver(Driver):
+    """Simple test driver"""
+
+    @classmethod
+    def client(cls):
+        return "jumpstarter.client.DriverClient"
+
+
+class CompositeDriver_(Driver):
+    """Simple composite driver for testing"""
+
+    @classmethod
+    def client(cls):
+        return "jumpstarter.client.DriverClient"
+
+
+def test_get_report_includes_descriptions():
+    """Test that GetReport includes descriptions for drivers that have them"""
+    # Create drivers with and without descriptions
+    driver_with_desc = SimpleDriver(description="Custom test driver")
+    driver_without_desc = SimpleDriver()
+
+    root = CompositeDriver_(
+        children={
+            "with_desc": driver_with_desc,
+            "without_desc": driver_without_desc,
+        }
+    )
+
+    with serve(root) as _:
+        # Get the raw report response
+
+        from jumpstarter.exporter.session import Session
+
+        # Create session manually to access GetReport
+        session = Session(
+            uuid=root.uuid,
+            labels=root.labels,
+            root_device=root,
+        )
+
+        # Call GetReport
+        import asyncio
+        response = asyncio.run(session.GetReport(empty_pb2.Empty(), None))
+
+        # Build a map of uuid -> report for easy lookup
+        reports_by_uuid = {r.uuid: r for r in response.reports}
+
+        # Verify driver with description has it in its report
+        assert str(driver_with_desc.uuid) in reports_by_uuid
+        report_with_desc = reports_by_uuid[str(driver_with_desc.uuid)]
+        assert hasattr(report_with_desc, 'description')
+        assert report_with_desc.description == "Custom test driver"
+
+        # Verify driver without description doesn't have the field set
+        assert str(driver_without_desc.uuid) in reports_by_uuid
+        report_without_desc = reports_by_uuid[str(driver_without_desc.uuid)]
+        # Optional field - either not set or empty string
+        assert not getattr(report_without_desc, 'description', None)
+
+
+def test_client_receives_description():
+    """Test that client receives description from GetReport"""
+    driver = SimpleDriver(description="Test description")
+
+    with serve(driver) as client:
+        # Description is passed during init from GetReport
+        assert client.description == "Test description"
+
+
+def test_cli_uses_description_or_default():
+    """Test that CLI uses description from GetReport or falls back to default"""
+    # Test with description set
+    driver_with_desc = SimpleDriver(description="Custom CLI description")
+    with serve(driver_with_desc) as client:
+        # Simulate what cli() method would do
+        help_text = client.description or "Default help text"
+        assert help_text == "Custom CLI description"
+
+    # Test without description
+    driver_without_desc = SimpleDriver()
+    with serve(driver_without_desc) as client:
+        help_text = client.description or "Default help text"
+        assert help_text == "Default help text"
+
+
+def test_multiple_drivers_with_descriptions():
+    """Test that multiple drivers can have different descriptions"""
+    power = SimpleDriver(description="Power control")
+    serial = SimpleDriver(description="Serial communication")
+    storage = SimpleDriver(description="Storage management")
+    plain = SimpleDriver()  # No description
+
+    root = CompositeDriver_(
+        children={
+            "power": power,
+            "serial": serial,
+            "storage": storage,
+            "plain": plain,
+        }
+    )
+
+    with serve(root) as client:
+        # Each child should have its description from GetReport
+        assert client.children['power'].description == "Power control"
+        assert client.children['serial'].description == "Serial communication"
+        assert client.children['storage'].description == "Storage management"
+        assert client.children['plain'].description is None
+
+
+def test_empty_description_not_included():
+    """Test that empty strings are not included in descriptions map"""
+    driver = SimpleDriver(description="")
+
+    with serve(driver) as _:
+        from jumpstarter.exporter.session import Session
+
+        session = Session(
+            uuid=driver.uuid,
+            labels=driver.labels,
+            root_device=driver,
+        )
+
+        import asyncio
+        response = asyncio.run(session.GetReport(empty_pb2.Empty(), None))
+
+        # Empty string should not be included in the report
+        reports_by_uuid = {r.uuid: r for r in response.reports}
+        assert str(driver.uuid) in reports_by_uuid
+        report = reports_by_uuid[str(driver.uuid)]
+        # Empty description should not be set
+        assert not getattr(report, 'description', None)
+
