Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 16 additions & 2 deletions packages/jumpstarter-cli/jumpstarter_cli/create.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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)
4 changes: 2 additions & 2 deletions packages/jumpstarter/jumpstarter/client/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
22 changes: 12 additions & 10 deletions packages/jumpstarter/jumpstarter/client/decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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)

Expand All @@ -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)

Expand Down
15 changes: 5 additions & 10 deletions packages/jumpstarter/jumpstarter/client/grpc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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]
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down
94 changes: 27 additions & 67 deletions packages/jumpstarter/jumpstarter/client/grpc_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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()
Expand All @@ -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()
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -262,26 +236,22 @@ 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"""
lease = self.create_test_lease(client="full-test-client", status="Active")

# 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
)

Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
)

Expand All @@ -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
Expand All @@ -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)


5 changes: 2 additions & 3 deletions packages/jumpstarter/jumpstarter/client/lease.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down
Loading