Unify ps implementation to fix the PM integ test

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."
     )

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}[/]")

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)

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