diff --git a/src/drunc/integtest/process_manager_test.py b/src/drunc/integtest/process_manager_test.py index 9b4212a09..f20ac2c1e 100644 --- a/src/drunc/integtest/process_manager_test.py +++ b/src/drunc/integtest/process_manager_test.py @@ -196,14 +196,25 @@ def test_boot(run_dunerc) -> None: """Checks that boot starts the managed processes and exposes UUIDs in ps.""" lines = strip_ansi(run_dunerc.completed_process.stdout).splitlines() - ps_pre_boot = get_ps_table_after_echo(lines, "pre_boot") - ps_post_boot = get_ps_table_after_echo(lines, "post_boot") - - assert not ps_pre_boot, ( - f"Expected ps table before boot to be empty, but found {len(ps_pre_boot)} row(s): " - + ", ".join(row["friendly_name"] for row in ps_pre_boot) + # Check if no processes running in session works + pre_boot_idx = require_line_containing( + lines, + "pre_boot", + error_message="Did not find the 'pre_boot' header line in stdout.", + ) + post_boot_idx = require_line_containing( + lines, + "post_boot", + error_message="Did not find the 'post_boot' footer line in stdout.", + ) + between = lines[pre_boot_idx + 1 : post_boot_idx] + no_boot_re = "No processes running in session" + assert any(no_boot_re in line for line in between), ( + f"Did not find '{no_boot_re}' between pre_boot and post_boot.\nBetween:\n" + + "\n".join(between) ) + ps_post_boot = get_ps_table_after_echo(lines, "post_boot") assert ps_post_boot, ( "Expected ps table after boot to contain processes, but it was empty." ) diff --git a/src/drunc/process_manager/interface/commands.py b/src/drunc/process_manager/interface/commands.py index 0d2cd3c43..4c6a14089 100644 --- a/src/drunc/process_manager/interface/commands.py +++ b/src/drunc/process_manager/interface/commands.py @@ -336,40 +336,57 @@ def restart_impl(obj: ProcessManagerContext, query: ProcessQuery) -> None: obj.get_driver("process_manager").restart(query) +def ps_decorators(f): + f = click.pass_obj(f) + f = click.option( + "-w", + "--width", + type=int, + default=None, + help="Table width. Default is automatically calculated", + )(f) + f = click.option( + "-l", + "--long-format", + is_flag=True, + type=bool, + default=False, + help="Whether to have a long output", + )(f) + + return f + + @click.command("ps") @add_query_options(at_least_one=False, all_processes_by_default=True) -@click.option( - "-l", - "--long-format", - is_flag=True, - type=bool, - default=False, - help="Whether to have a long output", -) -@click.option( - "-w", - "--width", - type=int, - default=None, - help="Table width. Default is automatically calculated", -) -@click.pass_obj -def ps( +@ps_decorators +def ps(obj, query, long_format, width): + log_pm_cmd(obj) + return ps_impl(obj, query, long_format, width) + + +def ps_impl( obj: ProcessManagerContext, query: ProcessQuery, long_format: bool, width: int | None, ) -> None: log = get_logger("process_manager.shell") - log_pm_cmd(obj) log.debug(f"Running ps with query {query}") results = obj.get_driver("process_manager").ps(query) - if not results: - return - obj.print( - tabulate_process_instance_list( - results, title="Processes running", long=long_format, width=width - ), - overflow="fold", - soft_wrap=True, - ) + + # If there are processes running, tabulate them, otherwise log that there are no + # processes running. + if results.values: + obj.print( + tabulate_process_instance_list( + results, + title=f"Processes running in session {obj.session_name}", + long=long_format, + width=width, + ), + overflow="fold", + soft_wrap=True, + ) + else: + log.info(f"No processes running in session [green]{obj.session_name}[/]") diff --git a/src/drunc/unified_shell/commands.py b/src/drunc/unified_shell/commands.py index 6285ad3bf..d9f58c4a6 100644 --- a/src/drunc/unified_shell/commands.py +++ b/src/drunc/unified_shell/commands.py @@ -15,10 +15,11 @@ kill_impl, logs_decorators, logs_impl, + ps_decorators, + ps_impl, restart_impl, ) from drunc.process_manager.interface.context import ProcessManagerContext -from drunc.process_manager.utils import tabulate_process_instance_list from drunc.unified_shell.context import UnifiedShellMode from drunc.utils.shell_utils import InterruptedCommand, log_pm_cmd from drunc.utils.utils import get_logger @@ -158,35 +159,6 @@ def terminate(ctx, obj): obj.delete_driver("controller") -@click.command("ps") -@click.pass_obj -@click.pass_context -def ps(ctx, obj): - """ - Execute the process manager terminate command, but only do this for the current - session - """ - - log = get_logger("unified_shell.ps") - log_pm_cmd(obj) - session_query = ProcessQuery(session=ctx.obj.session_name) - log.info(f"Listing session [green]{ctx.obj.session_name}[/]") - results = obj.get_driver("process_manager").ps(session_query) - - # If there are processes running, tabulate them, otherwise log that there are no - # processes running. - if results.values: - obj.print( - tabulate_process_instance_list( - results, title=f"Processes running in session {ctx.obj.session_name}" - ), - overflow="fold", - soft_wrap=True, - ) - else: - log.info(f"No processes running in session [green]{ctx.obj.session_name}[/]") - - def session_injector(f): @click.pass_context def wrapper(ctx, *args, **kwargs): @@ -196,6 +168,15 @@ def wrapper(ctx, *args, **kwargs): return update_wrapper(wrapper, f) +@click.command("ps") +@session_injector +@add_query_options_no_session(at_least_one=True) +@ps_decorators +def ps(obj, query, long_format, width): + log_pm_cmd(obj) + return ps_impl(obj, query, long_format, width) + + @click.command("logs") @session_injector @add_query_options_no_session(at_least_one=True) diff --git a/tests/process_manager/interface/test_commands.py b/tests/process_manager/interface/test_commands.py index f1a5c5e4b..9c843bdcb 100644 --- a/tests/process_manager/interface/test_commands.py +++ b/tests/process_manager/interface/test_commands.py @@ -115,6 +115,7 @@ class MockContext: def __init__(self, driver=None): self.driver = driver or MockDriver() self.output = [] + self.session_name = "mock-session" def get_driver(self, name): return self.driver