diff --git a/aiida/cmdline/commands/cmd_daemon.py b/aiida/cmdline/commands/cmd_daemon.py index e0853f0bce..bc568d56ba 100644 --- a/aiida/cmdline/commands/cmd_daemon.py +++ b/aiida/cmdline/commands/cmd_daemon.py @@ -17,6 +17,7 @@ import click from aiida.cmdline.commands.cmd_verdi import verdi +from aiida.cmdline.params import options from aiida.cmdline.utils import decorators, echo @@ -72,9 +73,10 @@ def execute_client_command(command: str, daemon_not_running_ok: bool = False, ** @verdi_daemon.command() @click.option('--foreground', is_flag=True, help='Run in foreground.') @click.argument('number', required=False, type=int, callback=validate_daemon_workers) +@options.TIMEOUT(default=None, required=False, type=int) @decorators.with_dbenv() @decorators.check_circus_zmq_version -def start(foreground, number): +def start(foreground, number, timeout): """Start the daemon with NUMBER workers. If the NUMBER of desired workers is not specified, the default is used, which is determined by the configuration @@ -83,14 +85,15 @@ def start(foreground, number): Returns exit code 0 if the daemon is OK, non-zero if there was an error. """ echo.echo(f'Starting the daemon with {number} workers... ', nl=False) - execute_client_command('start_daemon', number_workers=number, foreground=foreground) + execute_client_command('start_daemon', number_workers=number, foreground=foreground, timeout=timeout) @verdi_daemon.command() @click.option('--all', 'all_profiles', is_flag=True, help='Show status of all daemons.') +@options.TIMEOUT(default=None, required=False, type=int) @click.pass_context @decorators.requires_loaded_profile() -def status(ctx, all_profiles): +def status(ctx, all_profiles, timeout): """Print the status of the current daemon or all daemons. Returns exit code 0 if all requested daemons are running, else exit code 3. @@ -113,7 +116,7 @@ def status(ctx, all_profiles): echo.echo(f'{profile.name}', bold=True) try: - client.get_status() + client.get_status(timeout=timeout) except DaemonException as exception: echo.echo_error(str(exception)) daemons_running.append(False) @@ -148,26 +151,28 @@ def status(ctx, all_profiles): @verdi_daemon.command() @click.argument('number', default=1, type=int) +@options.TIMEOUT(default=None, required=False, type=int) @decorators.only_if_daemon_running() -def incr(number): +def incr(number, timeout): """Add NUMBER [default=1] workers to the running daemon. Returns exit code 0 if the daemon is OK, non-zero if there was an error. """ echo.echo(f'Starting {number} daemon workers... ', nl=False) - execute_client_command('increase_workers', number=number) + execute_client_command('increase_workers', number=number, timeout=timeout) @verdi_daemon.command() @click.argument('number', default=1, type=int) +@options.TIMEOUT(default=None, required=False, type=int) @decorators.only_if_daemon_running() -def decr(number): +def decr(number, timeout): """Remove NUMBER [default=1] workers from the running daemon. Returns exit code 0 if the daemon is OK, non-zero if there was an error. """ echo.echo(f'Stopping {number} daemon workers... ', nl=False) - execute_client_command('decrease_workers', number=number) + execute_client_command('decrease_workers', number=number, timeout=timeout) @verdi_daemon.command() @@ -184,8 +189,9 @@ def logshow(): @verdi_daemon.command() @click.option('--no-wait', is_flag=True, help='Do not wait for confirmation.') @click.option('--all', 'all_profiles', is_flag=True, help='Stop all daemons.') +@options.TIMEOUT(default=None, required=False, type=int) @click.pass_context -def stop(ctx, no_wait, all_profiles): +def stop(ctx, no_wait, all_profiles, timeout): """Stop the daemon. Returns exit code 0 if the daemon was shut down successfully (or was not running), non-zero if there was an error. @@ -199,16 +205,17 @@ def stop(ctx, no_wait, all_profiles): echo.echo('Profile: ', fg=echo.COLORS['report'], bold=True, nl=False) echo.echo(f'{profile.name}', bold=True) echo.echo('Stopping the daemon... ', nl=False) - execute_client_command('stop_daemon', daemon_not_running_ok=True, wait=not no_wait) + execute_client_command('stop_daemon', daemon_not_running_ok=True, wait=not no_wait, timeout=timeout) @verdi_daemon.command() @click.option('--reset', is_flag=True, help='Completely reset the daemon.') @click.option('--no-wait', is_flag=True, help='Do not wait for confirmation.') +@options.TIMEOUT(default=None, required=False, type=int) @click.pass_context @decorators.with_dbenv() @decorators.only_if_daemon_running() -def restart(ctx, reset, no_wait): +def restart(ctx, reset, no_wait, timeout): """Restart the daemon. By default will only reset the workers of the running daemon. After the restart the same amount of workers will be @@ -228,7 +235,7 @@ def restart(ctx, reset, no_wait): return echo.echo('Restarting the daemon... ', nl=False) - execute_client_command('restart_daemon', wait=not no_wait) + execute_client_command('restart_daemon', wait=not no_wait, timeout=timeout) @verdi_daemon.command(hidden=True) diff --git a/tests/cmdline/commands/test_daemon.py b/tests/cmdline/commands/test_daemon.py index cfc03d740d..f071501a78 100644 --- a/tests/cmdline/commands/test_daemon.py +++ b/tests/cmdline/commands/test_daemon.py @@ -116,6 +116,20 @@ def test_daemon_status_no_profile(run_cli_command): assert 'No profile loaded: make sure at least one profile is configured and a default is set.' in result.output +@pytest.mark.usefixtures('started_daemon_client', 'isolated_config') +def test_daemon_status_timeout(run_cli_command): + """Test ``verdi daemon status`` with the ``--timeout`` option. + + This will just test the timeout is accepted and doesn't cause an exception. It is not testing whether the command + will actually timeout if provided a small number, but that is difficult to make reproducible in a test environment. + """ + result = run_cli_command(cmd_daemon.status, ['--timeout', '5']) + last_line = result.output_lines[-1] + + assert f'Profile: {get_profile().name}' in result.output + assert last_line == 'Use `verdi daemon [incr | decr] [num]` to increase / decrease the number of workers' + + def get_daemon_info(_): """Mock replacement of :meth:`aiida.engine.daemon.client.DaemonClient.get_daemon_info`.""" return { @@ -165,7 +179,7 @@ def get_worker_info_broken(_): } -@patch.object(DaemonClient, 'get_status', lambda _: {'status': 'running'}) +@patch.object(DaemonClient, 'get_status', lambda *_, **__: {'status': 'running'}) @patch.object(DaemonClient, 'get_daemon_info', get_daemon_info) @patch.object(DaemonClient, 'get_worker_info', get_worker_info) @patch('aiida.cmdline.utils.common.format_local_time', format_local_time) @@ -184,7 +198,7 @@ def test_daemon_status_worker_info(run_cli_command): assert literal in result.output -@patch.object(DaemonClient, 'get_status', lambda _: {'status': 'running'}) +@patch.object(DaemonClient, 'get_status', lambda *_, **__: {'status': 'running'}) @patch.object(DaemonClient, 'get_daemon_info', get_daemon_info) @patch.object(DaemonClient, 'get_worker_info', get_worker_info_broken) @patch('aiida.cmdline.utils.common.format_local_time', format_local_time)