From 6d6c9fc3f18a641ae774d6dd2eb71777e8425775 Mon Sep 17 00:00:00 2001 From: Benny Zlotnik Date: Tue, 2 Dec 2025 13:28:03 +0200 Subject: [PATCH] cli: add support for requesting custom lease name Signed-off-by: Benny Zlotnik --- .../jumpstarter-cli/jumpstarter_cli/create.py | 18 +++- .../jumpstarter/jumpstarter/client/client.py | 4 +- .../jumpstarter/client/decorators.py | 22 +++-- .../jumpstarter/jumpstarter/client/grpc.py | 15 +-- .../jumpstarter/client/grpc_test.py | 94 ++++++------------- .../jumpstarter/jumpstarter/client/lease.py | 5 +- .../jumpstarter/client/lease_test.py | 56 ++++++----- .../jumpstarter/jumpstarter/config/client.py | 2 + 8 files changed, 97 insertions(+), 119 deletions(-) diff --git a/packages/jumpstarter-cli/jumpstarter_cli/create.py b/packages/jumpstarter-cli/jumpstarter_cli/create.py index 0dff98b53..c08598b5d 100644 --- a/packages/jumpstarter-cli/jumpstarter_cli/create.py +++ b/packages/jumpstarter-cli/jumpstarter_cli/create.py @@ -22,9 +22,17 @@ def create(): @opt_selector @opt_duration_partial(required=True) @opt_begin_time +@click.option( + "--lease-id", + type=str, + default=None, + help="Optional lease ID to request (if not provided, server will generate one)", +) @opt_output_all @handle_exceptions_with_reauthentication(relogin_client) -def create_lease(config, selector: str, duration: timedelta, begin_time: datetime | None, output: OutputType): +def create_lease( + config, selector: str, duration: timedelta, begin_time: datetime | None, lease_id: str | None, output: OutputType +): """ Create a lease @@ -48,8 +56,14 @@ def create_lease(config, selector: str, duration: timedelta, begin_time: datetim $$ exit $ jmp delete lease "${JMP_LEASE}" + You can also specify a unique custom lease ID: + + .. code-block:: bash + + $ jmp create lease -l foo=bar --duration 1d --lease-id my-custom-lease-id + """ - lease = config.create_lease(selector=selector, duration=duration, begin_time=begin_time) + lease = config.create_lease(selector=selector, duration=duration, begin_time=begin_time, lease_id=lease_id) model_print(lease, output) diff --git a/packages/jumpstarter/jumpstarter/client/client.py b/packages/jumpstarter/jumpstarter/client/client.py index f4aaef824..355bb24ca 100644 --- a/packages/jumpstarter/jumpstarter/client/client.py +++ b/packages/jumpstarter/jumpstarter/client/client.py @@ -59,8 +59,8 @@ async def client_from_channel( 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, - methods_description=getattr(report, 'methods_description', {}) or {}, + description=getattr(report, "description", None) or None, + methods_description=getattr(report, "methods_description", {}) or {}, ) clients[index] = client diff --git a/packages/jumpstarter/jumpstarter/client/decorators.py b/packages/jumpstarter/jumpstarter/client/decorators.py index 2343f8355..09daf9a5c 100644 --- a/packages/jumpstarter/jumpstarter/client/decorators.py +++ b/packages/jumpstarter/jumpstarter/client/decorators.py @@ -34,20 +34,21 @@ def on(): :param kwargs: Keyword arguments passed to DriverClickGroup :return: Decorator that creates a DriverClickGroup """ + def decorator(f: Callable) -> DriverClickGroup: # Use function docstring if no help= provided - if 'help' not in kwargs or kwargs['help'] is None: + if "help" not in kwargs or kwargs["help"] is None: if f.__doc__: - kwargs['help'] = f.__doc__.strip() + kwargs["help"] = f.__doc__.strip() # Server description overrides Click defaults - if getattr(client, 'description', None): - kwargs['help'] = client.description + if getattr(client, "description", None): + kwargs["help"] = client.description group = DriverClickGroup(client, name=f.__name__, callback=f, **kwargs) # Transfer Click parameters attached by decorators like @click.option - group.params = getattr(f, '__click_params__', []) + group.params = getattr(f, "__click_params__", []) return group @@ -74,8 +75,8 @@ def ssh(args): :return: click.command decorator """ # Server description overrides Click's defaults (help= parameter or docstring) - if getattr(client, 'description', None): - kwargs['help'] = client.description + if getattr(client, "description", None): + kwargs["help"] = client.description return click.command(**kwargs) @@ -89,13 +90,14 @@ def __init__(self, client: "DriverClient", *args: Any, **kwargs: Any) -> None: def command(self, *args: Any, **kwargs: Any) -> Callable: """Command decorator with server methods_description override support.""" + def decorator(f: Callable) -> click.Command: - name = kwargs.get('name') + name = kwargs.get("name") if not name: - name = f.__name__.lower().replace('_', '-') + name = f.__name__.lower().replace("_", "-") if name in self.client.methods_description: - kwargs['help'] = self.client.methods_description[name] + kwargs["help"] = self.client.methods_description[name] return super(DriverClickGroup, self).command(*args, **kwargs)(f) diff --git a/packages/jumpstarter/jumpstarter/client/grpc.py b/packages/jumpstarter/jumpstarter/client/grpc.py index fc6f52527..445f255c6 100644 --- a/packages/jumpstarter/jumpstarter/client/grpc.py +++ b/packages/jumpstarter/jumpstarter/client/grpc.py @@ -275,11 +275,7 @@ def model_dump_json(self, **kwargs): if not self.include_online: exclude_fields.add("online") - data = { - "exporters": [ - exporter.model_dump(mode="json", exclude=exclude_fields) for exporter in self.exporters - ] - } + data = {"exporters": [exporter.model_dump(mode="json", exclude=exclude_fields) for exporter in self.exporters]} return json.dumps(data, **json_kwargs) def model_dump(self, **kwargs): @@ -289,11 +285,8 @@ def model_dump(self, **kwargs): if not self.include_online: exclude_fields.add("online") - return { - "exporters": [ - exporter.model_dump(mode="json", exclude=exclude_fields) for exporter in self.exporters - ] - } + return {"exporters": [exporter.model_dump(mode="json", exclude=exclude_fields) for exporter in self.exporters]} + class LeaseList(BaseModel): leases: list[Lease] @@ -390,6 +383,7 @@ async def CreateLease( selector: str, duration: timedelta, begin_time: datetime | None = None, + lease_id: str | None = None, ): duration_pb = duration_pb2.Duration() duration_pb.FromTimedelta(duration) @@ -409,6 +403,7 @@ async def CreateLease( client_pb2.CreateLeaseRequest( parent="namespaces/{}".format(self.namespace), lease=lease_pb, + lease_id=lease_id or "", ) ) return Lease.from_protobuf(lease) diff --git a/packages/jumpstarter/jumpstarter/client/grpc_test.py b/packages/jumpstarter/jumpstarter/client/grpc_test.py index 8146c4933..49a1ed376 100644 --- a/packages/jumpstarter/jumpstarter/client/grpc_test.py +++ b/packages/jumpstarter/jumpstarter/client/grpc_test.py @@ -63,12 +63,7 @@ class TestAddExporterRow: def create_test_exporter(self, online=True, labels=None): if labels is None: labels = {"env": "test", "type": "device"} - return Exporter( - namespace="default", - name="test-exporter", - labels=labels, - online=online - ) + return Exporter(namespace="default", name="test-exporter", labels=labels, online=online) def test_basic_row(self): table = Table() @@ -119,11 +114,16 @@ def test_row_with_all_options(self): class TestExporterList: - def create_test_lease(self, client="test-client", status="Active", - effective_begin_time=datetime(2023, 1, 1, 10, 0, 0), - effective_duration=timedelta(hours=1), - begin_time=None, duration=timedelta(hours=1), - effective_end_time=None): + def create_test_lease( + self, + client="test-client", + status="Active", + effective_begin_time=datetime(2023, 1, 1, 10, 0, 0), + effective_duration=timedelta(hours=1), + begin_time=None, + duration=timedelta(hours=1), + effective_end_time=None, + ): lease = Mock(spec=Lease) lease.client = client lease.get_status.return_value = status @@ -135,12 +135,7 @@ def create_test_lease(self, client="test-client", status="Active", return lease def test_exporter_without_lease(self): - exporter = Exporter( - namespace="default", - name="test-exporter", - labels={"type": "device"}, - online=True - ) + exporter = Exporter(namespace="default", name="test-exporter", labels={"type": "device"}, online=True) table = Table() Exporter.rich_add_columns(table) @@ -152,11 +147,7 @@ def test_exporter_without_lease(self): def test_exporter_with_lease_no_display(self): lease = self.create_test_lease() exporter = Exporter( - namespace="default", - name="test-exporter", - labels={"type": "device"}, - online=True, - lease=lease + namespace="default", name="test-exporter", labels={"type": "device"}, online=True, lease=lease ) table = Table() @@ -170,11 +161,7 @@ def test_exporter_with_lease_no_display(self): def test_exporter_with_lease_display(self): lease = self.create_test_lease() exporter = Exporter( - namespace="default", - name="test-exporter", - labels={"type": "device"}, - online=True, - lease=lease + namespace="default", name="test-exporter", labels={"type": "device"}, online=True, lease=lease ) table = Table() @@ -198,12 +185,7 @@ def test_exporter_with_lease_display(self): assert "2023-01-01 11:00:00" in output # Expected release: begin_time (10:00:00) + duration (1h) def test_exporter_without_lease_but_show_leases(self): - exporter = Exporter( - namespace="default", - name="test-exporter", - labels={"type": "device"}, - online=True - ) + exporter = Exporter(namespace="default", name="test-exporter", labels={"type": "device"}, online=True) table = Table() options = WithOptions(show_leases=True) @@ -228,19 +210,11 @@ def test_exporter_without_lease_but_show_leases(self): def test_exporter_online_status_display(self): """Test that online status icons are correctly displayed""" # Test online exporter - exporter_online = Exporter( - namespace="default", - name="online-exporter", - labels={"type": "device"}, - online=True - ) + exporter_online = Exporter(namespace="default", name="online-exporter", labels={"type": "device"}, online=True) # Test offline exporter exporter_offline = Exporter( - namespace="default", - name="offline-exporter", - labels={"type": "device"}, - online=False + namespace="default", name="offline-exporter", labels={"type": "device"}, online=False ) # Test with online status display enabled @@ -262,7 +236,7 @@ def test_exporter_online_status_display(self): assert "online-exporter" in output assert "offline-exporter" in output assert "yes" in output # Should show "yes" for online - assert "no" in output # Should show "no" for offline + assert "no" in output # Should show "no" for offline def test_exporter_all_features_display(self): """Test all display features together: online status + lease info""" @@ -270,18 +244,14 @@ def test_exporter_all_features_display(self): # Create exporters with different combinations of online/lease status exporter_online_with_lease = Exporter( - namespace="default", - name="online-with-lease", - labels={"env": "prod"}, - online=True, - lease=lease + namespace="default", name="online-with-lease", labels={"env": "prod"}, online=True, lease=lease ) exporter_offline_no_lease = Exporter( namespace="default", name="offline-no-lease", labels={"env": "dev"}, - online=False + online=False, # No lease ) @@ -306,7 +276,7 @@ def test_exporter_all_features_display(self): assert "env=prod" in output assert "env=dev" in output assert "yes" in output # Online indicator - assert "no" in output # Offline indicator + assert "no" in output # Offline indicator assert "full-test-client" in output # Lease client assert "Active" in output # Lease status assert "Available" in output # Available status for no lease @@ -317,14 +287,10 @@ def test_exporter_lease_info_extraction(self): lease = self.create_test_lease( client="my-client", status="Expired", - effective_end_time=datetime(2023, 1, 1, 11, 0, 0) # Ended after 1 hour + effective_end_time=datetime(2023, 1, 1, 11, 0, 0), # Ended after 1 hour ) exporter = Exporter( - namespace="default", - name="test-exporter", - labels={"type": "device"}, - online=True, - lease=lease + namespace="default", name="test-exporter", labels={"type": "device"}, online=True, lease=lease ) # Manually verify the lease data that would be extracted @@ -359,7 +325,7 @@ def test_exporter_no_lease_info_extraction(self): namespace="default", name="test-exporter", labels={"type": "device"}, - online=True + online=True, # No lease attached ) @@ -380,16 +346,12 @@ def test_exporter_scheduled_lease_expected_release(self): client="my-client", status="Scheduled", effective_begin_time=None, # Not started yet - effective_duration=None, # Not started yet + effective_duration=None, # Not started yet begin_time=datetime(2023, 1, 1, 10, 0, 0), - duration=timedelta(hours=1) + duration=timedelta(hours=1), ) exporter = Exporter( - namespace="default", - name="test-exporter", - labels={"type": "device"}, - online=True, - lease=lease + namespace="default", name="test-exporter", labels={"type": "device"}, online=True, lease=lease ) # Test the table display with scheduled lease @@ -412,5 +374,3 @@ def test_exporter_scheduled_lease_expected_release(self): assert "my-client" in output assert "Scheduled" in output assert "2023-01-01 11:00:00" in output # begin_time (10:00) + duration (1h) - - diff --git a/packages/jumpstarter/jumpstarter/client/lease.py b/packages/jumpstarter/jumpstarter/client/lease.py index fc61d45f6..ac16cb945 100644 --- a/packages/jumpstarter/jumpstarter/client/lease.py +++ b/packages/jumpstarter/jumpstarter/client/lease.py @@ -69,6 +69,7 @@ async def _create(self): await self.svc.CreateLease( selector=self.selector, duration=self.duration, + lease_id=self.name, ) ).name logger.info("Acquiring lease %s for selector %s for duration %s", self.name, self.selector, self.duration) @@ -377,9 +378,7 @@ def update_status(self, message: str, force: bool = False): # Throttle updates to at most every 5 minutes unless forced now = datetime.now() should_log = ( - force - or self._last_log_time is None - or (now - self._last_log_time) >= self._log_throttle_interval + force or self._last_log_time is None or (now - self._last_log_time) >= self._log_throttle_interval ) if should_log: diff --git a/packages/jumpstarter/jumpstarter/client/lease_test.py b/packages/jumpstarter/jumpstarter/client/lease_test.py index 43bb084b9..4d7dbc854 100644 --- a/packages/jumpstarter/jumpstarter/client/lease_test.py +++ b/packages/jumpstarter/jumpstarter/client/lease_test.py @@ -32,31 +32,37 @@ def test_init_without_lease_name(self): def test_is_terminal_available_with_tty(self): """Test terminal detection when TTY is available.""" - with patch.object(sys.stdout, 'isatty', return_value=True), \ - patch.object(sys.stderr, 'isatty', return_value=True): + with ( + patch.object(sys.stdout, "isatty", return_value=True), + patch.object(sys.stderr, "isatty", return_value=True), + ): spinner = LeaseAcquisitionSpinner() assert spinner._is_terminal_available() is True def test_is_terminal_available_without_tty(self): """Test terminal detection when TTY is not available.""" - with patch.object(sys.stdout, 'isatty', return_value=False), \ - patch.object(sys.stderr, 'isatty', return_value=False): + with ( + patch.object(sys.stdout, "isatty", return_value=False), + patch.object(sys.stderr, "isatty", return_value=False), + ): spinner = LeaseAcquisitionSpinner() assert spinner._is_terminal_available() is False def test_is_terminal_available_partial_tty(self): """Test terminal detection when only one stream is TTY.""" - with patch.object(sys.stdout, 'isatty', return_value=True), \ - patch.object(sys.stderr, 'isatty', return_value=False): + with ( + patch.object(sys.stdout, "isatty", return_value=True), + patch.object(sys.stderr, "isatty", return_value=False), + ): spinner = LeaseAcquisitionSpinner() assert spinner._is_terminal_available() is False def test_context_manager_with_console(self): """Test context manager behavior when console is available.""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=True): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=True): spinner = LeaseAcquisitionSpinner("test-lease") - with patch.object(spinner.console, 'status') as mock_status: + with patch.object(spinner.console, "status") as mock_status: mock_spinner = Mock() mock_status.return_value = mock_spinner @@ -70,10 +76,10 @@ def test_context_manager_with_console(self): def test_context_manager_without_console(self): """Test context manager behavior when console is not available.""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=False): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=False): spinner = LeaseAcquisitionSpinner("test-lease") - with patch.object(spinner.console, 'status') as mock_status: + with patch.object(spinner.console, "status") as mock_status: with spinner as ctx_spinner: assert ctx_spinner is spinner assert spinner.start_time is not None @@ -81,7 +87,7 @@ def test_context_manager_without_console(self): def test_update_status_with_console(self): """Test status update when console is available.""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=True): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=True): spinner = LeaseAcquisitionSpinner("test-lease") spinner.start_time = datetime.now() @@ -98,7 +104,7 @@ def test_update_status_with_console(self): def test_update_status_without_console(self, caplog): """Test status update when console is not available (should log).""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=False): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=False): spinner = LeaseAcquisitionSpinner("test-lease") spinner.start_time = datetime.now() @@ -110,7 +116,7 @@ def test_update_status_without_console(self, caplog): def test_tick_with_console_and_message(self): """Test tick update when console is available and message exists.""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=True): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=True): spinner = LeaseAcquisitionSpinner("test-lease") spinner.start_time = datetime.now() spinner._current_message = "[blue]Test message[/blue]" @@ -127,7 +133,7 @@ def test_tick_with_console_and_message(self): def test_tick_without_console(self): """Test tick update when console is not available (should not log).""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=False): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=False): spinner = LeaseAcquisitionSpinner("test-lease") spinner.start_time = datetime.now() spinner._current_message = "[blue]Test message[/blue]" @@ -137,7 +143,7 @@ def test_tick_without_console(self): def test_tick_without_message(self): """Test tick update when no current message exists.""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=True): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=True): spinner = LeaseAcquisitionSpinner("test-lease") spinner.start_time = datetime.now() spinner._current_message = None @@ -152,7 +158,7 @@ def test_tick_without_message(self): def test_elapsed_time_formatting(self): """Test that elapsed time is formatted correctly.""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=True): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=True): spinner = LeaseAcquisitionSpinner("test-lease") spinner.start_time = datetime.now() - timedelta(seconds=65) # 1:05 spinner._current_message = "[blue]Test message[/blue]" @@ -170,10 +176,10 @@ def test_elapsed_time_formatting(self): @pytest.mark.asyncio async def test_integration_with_async_context(self): """Test integration with async context manager.""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=True): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=True): spinner = LeaseAcquisitionSpinner("test-lease") - with patch.object(spinner.console, 'status') as mock_status: + with patch.object(spinner.console, "status") as mock_status: mock_spinner = Mock() mock_status.return_value = mock_spinner @@ -195,7 +201,7 @@ async def test_async_usage(): def test_message_preservation_across_ticks(self): """Test that the base message is preserved across multiple ticks.""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=True): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=True): spinner = LeaseAcquisitionSpinner("test-lease") spinner.start_time = datetime.now() @@ -231,7 +237,7 @@ def test_start_time_initialization_in_context(self): def test_throttling_first_update_logged(self, caplog): """Test that the first update is always logged when console is not available.""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=False): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=False): spinner = LeaseAcquisitionSpinner("test-lease") spinner.start_time = datetime.now() @@ -243,7 +249,7 @@ def test_throttling_first_update_logged(self, caplog): def test_throttling_second_update_within_interval_not_logged(self, caplog): """Test that updates within 5 minutes are not logged.""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=False): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=False): spinner = LeaseAcquisitionSpinner("test-lease") spinner.start_time = datetime.now() spinner._last_log_time = datetime.now() - timedelta(minutes=2) # 2 minutes ago @@ -256,7 +262,7 @@ def test_throttling_second_update_within_interval_not_logged(self, caplog): def test_throttling_update_after_interval_logged(self, caplog): """Test that updates after 5 minutes are logged.""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=False): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=False): spinner = LeaseAcquisitionSpinner("test-lease") spinner.start_time = datetime.now() spinner._last_log_time = datetime.now() - timedelta(minutes=6) # 6 minutes ago @@ -269,7 +275,7 @@ def test_throttling_update_after_interval_logged(self, caplog): def test_throttling_forced_update_always_logged(self, caplog): """Test that forced updates are always logged regardless of throttle interval.""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=False): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=False): spinner = LeaseAcquisitionSpinner("test-lease") spinner.start_time = datetime.now() spinner._last_log_time = datetime.now() - timedelta(minutes=1) # 1 minute ago @@ -282,7 +288,7 @@ def test_throttling_forced_update_always_logged(self, caplog): def test_throttling_multiple_updates_only_logs_when_needed(self, caplog): """Test that multiple rapid updates only log at appropriate intervals.""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=False): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=False): spinner = LeaseAcquisitionSpinner("test-lease") spinner.start_time = datetime.now() @@ -311,7 +317,7 @@ def test_throttling_multiple_updates_only_logs_when_needed(self, caplog): def test_throttling_not_applied_when_console_available(self): """Test that throttling is not applied when console is available.""" - with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=True): + with patch.object(LeaseAcquisitionSpinner, "_is_terminal_available", return_value=True): spinner = LeaseAcquisitionSpinner("test-lease") spinner.start_time = datetime.now() diff --git a/packages/jumpstarter/jumpstarter/config/client.py b/packages/jumpstarter/jumpstarter/config/client.py index 1df373e09..97f92c1ec 100644 --- a/packages/jumpstarter/jumpstarter/config/client.py +++ b/packages/jumpstarter/jumpstarter/config/client.py @@ -201,12 +201,14 @@ async def create_lease( selector: str, duration: timedelta, begin_time: datetime | None = None, + lease_id: str | None = None, ): svc = ClientService(channel=await self.channel(), namespace=self.metadata.namespace) return await svc.CreateLease( selector=selector, duration=duration, begin_time=begin_time, + lease_id=lease_id, ) @_blocking_compat