From fc0ec01e9c56382f33b8ef11cb7f0254895c2e63 Mon Sep 17 00:00:00 2001 From: Miguel Angel Ajo Pelayo Date: Mon, 6 Oct 2025 21:35:51 +0200 Subject: [PATCH] Lease progress details while users wait (#688) (cherry picked from commit 5f3a04690d4c4d12d445dc077783a57af0e166e1) --- .../jumpstarter/jumpstarter/client/lease.py | 150 +++++++++--- .../jumpstarter/client/lease_test.py | 230 ++++++++++++++++++ 2 files changed, 349 insertions(+), 31 deletions(-) create mode 100644 packages/jumpstarter/jumpstarter/client/lease_test.py diff --git a/packages/jumpstarter/jumpstarter/client/lease.py b/packages/jumpstarter/jumpstarter/client/lease.py index c9bd4d5c3..835019887 100644 --- a/packages/jumpstarter/jumpstarter/client/lease.py +++ b/packages/jumpstarter/jumpstarter/client/lease.py @@ -1,4 +1,6 @@ import logging +import os +import sys from collections.abc import AsyncGenerator, Generator from contextlib import ( ExitStack, @@ -20,6 +22,7 @@ from anyio.from_thread import BlockingPortal from grpc.aio import AioRpcError, Channel from jumpstarter_protocol import jumpstarter_pb2, jumpstarter_pb2_grpc +from rich.console import Console from .exceptions import LeaseError from jumpstarter.client import client_from_path @@ -112,6 +115,17 @@ async def request_async(self): return await self._acquire() + def _update_spinner_status(self, spinner, result): + """Update spinner with appropriate status message based on lease conditions.""" + if condition_true(result.conditions, "Pending"): + pending_message = condition_message(result.conditions, "Pending") + if pending_message: + spinner.update_status(f"Waiting for lease: {pending_message}") + else: + spinner.update_status("Waiting for lease to be ready...") + else: + spinner.update_status("Waiting for server to provide status updates...") + async def _acquire(self): """Acquire a lease. @@ -119,37 +133,50 @@ async def _acquire(self): """ try: with fail_after(self.acquisition_timeout): - while True: - logger.debug("Polling Lease %s", self.name) - result = await self.get() - # lease ready - if condition_true(result.conditions, "Ready"): - logger.debug("Lease %s acquired", self.name) - self.exporter_name = result.exporter - return self - # lease unsatisfiable - if condition_true(result.conditions, "Unsatisfiable"): - message = condition_message(result.conditions, "Unsatisfiable") - logger.debug("Lease %s cannot be satisfied: %s", self.name, message) - raise LeaseError(f"the lease cannot be satisfied: {message}") - - # lease invalid - if condition_true(result.conditions, "Invalid"): - message = condition_message(result.conditions, "Invalid") - logger.debug("Lease %s is invalid: %s", self.name, message) - raise LeaseError(f"the lease is invalid: {message}") - - # lease not pending - if condition_false(result.conditions, "Pending"): - raise LeaseError( - f"Lease {self.name} is not in pending, but it isn't in Ready or Unsatisfiable state either" - ) - - # lease released - if condition_present_and_equal(result.conditions, "Ready", "False", "Released"): - raise LeaseError(f"lease {self.name} released") - - await sleep(5) + with LeaseAcquisitionSpinner(self.name) as spinner: + while True: + logger.debug("Polling Lease %s", self.name) + result = await self.get() + + # lease ready + if condition_true(result.conditions, "Ready"): + logger.debug("Lease %s acquired", self.name) + spinner.update_status(f"Lease {self.name} acquired successfully!") + self.exporter_name = result.exporter + break + + # lease unsatisfiable + if condition_true(result.conditions, "Unsatisfiable"): + message = condition_message(result.conditions, "Unsatisfiable") + logger.debug("Lease %s cannot be satisfied: %s", self.name, message) + raise LeaseError(f"the lease cannot be satisfied: {message}") + + # lease invalid + if condition_true(result.conditions, "Invalid"): + message = condition_message(result.conditions, "Invalid") + logger.debug("Lease %s is invalid: %s", self.name, message) + raise LeaseError(f"the lease is invalid: {message}") + + # lease not pending + if condition_false(result.conditions, "Pending"): + raise LeaseError( + f"Lease {self.name} is not in pending, but it isn't in Ready or " + f"Unsatisfiable state either" + ) + + # lease released + if condition_present_and_equal(result.conditions, "Ready", "False", "Released"): + raise LeaseError(f"lease {self.name} released") + + # Update spinner with appropriate status message + self._update_spinner_status(spinner, result) + + # Wait in 1-second increments with tick updates for better UX + for _ in range(5): + await sleep(1) + spinner.tick() + return self + except TimeoutError: logger.debug(f"Lease {self.name} acquisition timed out after {self.acquisition_timeout} seconds") raise LeaseError( @@ -269,3 +296,64 @@ def serve_unix(self): def monitor(self, threshold: timedelta = timedelta(minutes=5)): with self.portal.wrap_async_context_manager(self.monitor_async(threshold)): yield + + +class LeaseAcquisitionSpinner: + """Context manager for displaying a spinner during lease acquisition.""" + + def __init__(self, lease_name: str | None = None): + self.lease_name = lease_name + self.console = Console() + self.spinner = None + self.start_time = None + self._should_show_spinner = self._is_terminal_available() and not self._is_non_interactive() + self._current_message = None + + def _is_non_interactive(self) -> bool: + """Check if the user desires a NONINTERACTIVE environment.""" + return os.environ.get("NONINTERACTIVE", "false").lower() in ["true", "1"] + + def _is_terminal_available(self) -> bool: + """Check if we're running in a terminal/TTY.""" + return ( + hasattr(sys.stdout, 'isatty') and + sys.stdout.isatty() and + hasattr(sys.stderr, 'isatty') and + sys.stderr.isatty() + ) + + def __enter__(self): + self.start_time = datetime.now() + if self._should_show_spinner: + self.spinner = self.console.status( + f"Acquiring lease {self.lease_name or '...'}...", + spinner="dots", + spinner_style="blue" + ) + self.spinner.start() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if self.spinner: + self.spinner.stop() + + def update_status(self, message: str): + """Update the spinner status message.""" + if self.spinner and self._should_show_spinner: + self._current_message = f"[blue]{message}[/blue]" + elapsed = datetime.now() - self.start_time + elapsed_str = str(elapsed).split('.')[0] # Remove microseconds + self.spinner.update(f"{self._current_message} [dim]({elapsed_str})[/dim]") + else: + # Log info message when no console is available + elapsed = datetime.now() - self.start_time + elapsed_str = str(elapsed).split('.')[0] # Remove microseconds + logger.info(f"{message} ({elapsed_str})") + + def tick(self): + """Update the spinner with current elapsed time without changing the message.""" + if self.spinner and self._should_show_spinner and self._current_message: + elapsed = datetime.now() - self.start_time + elapsed_str = str(elapsed).split('.')[0] # Remove microseconds + # Use the stored current message and update with new elapsed time + self.spinner.update(f"{self._current_message} [dim]({elapsed_str})[/dim]") diff --git a/packages/jumpstarter/jumpstarter/client/lease_test.py b/packages/jumpstarter/jumpstarter/client/lease_test.py new file mode 100644 index 000000000..f00d81266 --- /dev/null +++ b/packages/jumpstarter/jumpstarter/client/lease_test.py @@ -0,0 +1,230 @@ +import asyncio +import logging +import sys +from datetime import datetime, timedelta +from unittest.mock import Mock, patch + +import pytest +from rich.console import Console + +from jumpstarter.client.lease import LeaseAcquisitionSpinner + + +class TestLeaseAcquisitionSpinner: + """Test cases for LeaseAcquisitionSpinner class.""" + + def test_init_with_lease_name(self): + """Test spinner initialization with lease name.""" + spinner = LeaseAcquisitionSpinner("test-lease-123") + assert spinner.lease_name == "test-lease-123" + assert spinner.console is not None + assert spinner.spinner is None + assert spinner.start_time is None + assert isinstance(spinner._should_show_spinner, bool) + + def test_init_without_lease_name(self): + """Test spinner initialization without lease name.""" + spinner = LeaseAcquisitionSpinner() + assert spinner.lease_name is None + assert spinner.console is not None + assert spinner.spinner is None + assert spinner.start_time is None + + 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): + 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): + 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): + 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): + spinner = LeaseAcquisitionSpinner("test-lease") + + with patch.object(spinner.console, 'status') as mock_status: + mock_spinner = Mock() + mock_status.return_value = mock_spinner + + with spinner as ctx_spinner: + assert ctx_spinner is spinner + assert spinner.start_time is not None + mock_status.assert_called_once() + mock_spinner.start.assert_called_once() + + mock_spinner.stop.assert_called_once() + + 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): + spinner = LeaseAcquisitionSpinner("test-lease") + + 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 + mock_status.assert_not_called() + + def test_update_status_with_console(self): + """Test status update when console is available.""" + with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=True): + spinner = LeaseAcquisitionSpinner("test-lease") + spinner.start_time = datetime.now() + + mock_spinner = Mock() + spinner.spinner = mock_spinner + + spinner.update_status("Test message") + + assert spinner._current_message == "[blue]Test message[/blue]" + mock_spinner.update.assert_called_once() + call_args = mock_spinner.update.call_args[0][0] + assert "[blue]Test message[/blue]" in call_args + assert "[dim](" in call_args + + 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): + spinner = LeaseAcquisitionSpinner("test-lease") + spinner.start_time = datetime.now() + + with caplog.at_level(logging.INFO): + spinner.update_status("Test message") + + assert "Test message" in caplog.text + assert spinner._current_message is None + + 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): + spinner = LeaseAcquisitionSpinner("test-lease") + spinner.start_time = datetime.now() + spinner._current_message = "[blue]Test message[/blue]" + + mock_spinner = Mock() + spinner.spinner = mock_spinner + + spinner.tick() + + mock_spinner.update.assert_called_once() + call_args = mock_spinner.update.call_args[0][0] + assert "[blue]Test message[/blue]" in call_args + assert "[dim](" in call_args + + 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): + spinner = LeaseAcquisitionSpinner("test-lease") + spinner.start_time = datetime.now() + spinner._current_message = "[blue]Test message[/blue]" + + # Should not raise any exceptions or log anything + spinner.tick() + + def test_tick_without_message(self): + """Test tick update when no current message exists.""" + with patch.object(LeaseAcquisitionSpinner, '_is_terminal_available', return_value=True): + spinner = LeaseAcquisitionSpinner("test-lease") + spinner.start_time = datetime.now() + spinner._current_message = None + + mock_spinner = Mock() + spinner.spinner = mock_spinner + + spinner.tick() + + # Should not call update when no message + mock_spinner.update.assert_not_called() + + def test_elapsed_time_formatting(self): + """Test that elapsed time is formatted correctly.""" + 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]" + + mock_spinner = Mock() + spinner.spinner = mock_spinner + + spinner.tick() + + call_args = mock_spinner.update.call_args[0][0] + # Should contain time in format like "0:01:05" + assert "[dim](" in call_args + assert "[/dim]" in call_args + + @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): + spinner = LeaseAcquisitionSpinner("test-lease") + + with patch.object(spinner.console, 'status') as mock_status: + mock_spinner = Mock() + mock_status.return_value = mock_spinner + + async def test_async_usage(): + with spinner as ctx_spinner: + ctx_spinner.update_status("Initial message") + await asyncio.sleep(0.1) # Small delay + ctx_spinner.tick() + ctx_spinner.update_status("Updated message") + + await test_async_usage() + + # Verify all expected calls were made + mock_status.assert_called_once() + assert mock_spinner.start.call_count == 1 + assert mock_spinner.stop.call_count == 1 + # update_status calls update() for each status update, tick() calls update() once + assert mock_spinner.update.call_count == 3 # 2 update_status calls + 1 tick call + + 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): + spinner = LeaseAcquisitionSpinner("test-lease") + spinner.start_time = datetime.now() + + # Set up mock before calling update_status + mock_spinner = Mock() + spinner.spinner = mock_spinner + + spinner.update_status("Waiting for lease: Test condition") + + # Call tick multiple times + for _ in range(3): + spinner.tick() + + # All calls should preserve the base message + assert mock_spinner.update.call_count == 4 # 1 update_status + 3 ticks + for call in mock_spinner.update.call_args_list: + call_args = call[0][0] + assert "[blue]Waiting for lease: Test condition[/blue]" in call_args + + def test_console_initialization(self): + """Test that console is properly initialized.""" + spinner = LeaseAcquisitionSpinner() + assert isinstance(spinner.console, Console) + + def test_start_time_initialization_in_context(self): + """Test that start_time is set when entering context.""" + spinner = LeaseAcquisitionSpinner("test-lease") + assert spinner.start_time is None + + with spinner: + assert spinner.start_time is not None + assert isinstance(spinner.start_time, datetime)