Noting batch mode for exiting integ tests

src/drunc/controller/controller.py

@@ -1232,3 +1232,52 @@ def who_is_in_charge(
             flag=ResponseFlag.EXECUTED_SUCCESSFULLY,
             children=response_children,
         )
+
+    ##########################################
+    ####### Integration test commands ########
+    ##########################################
+
+
+    # ORDER MATTERS!
+    @broadcasted  # outer most wrapper 1st step
+    @authentified_and_authorised(
+        action=ActionType.UPDATE, system=SystemType.CONTROLLER
+    )  # 2nd step
+    @in_control
+    @unpack_addressed_command_to()  # 3rd step
+    @publish_command_time
+    def to_error(
+        self,
+        addressed_commands: dict[str, AddressedCommand],
+        execute_on_self: bool,
+        token: Token
+    ) -> PlainText:
+        """
+        Transitions the stateful node to an error state. Used for testing purposes.
+        """
+        try:
+            if execute_on_self:
+                self.stateful_node.to_error()
+
+            response_children = self.propagate_addressed_command(
+                "to_error",
+                addressed_commands=addressed_commands,
+                token=token,
+            )
+
+            return Response(
+                name=self.name,
+                token=token,
+                data=None,
+                flag=ResponseFlag.EXECUTED_SUCCESSFULLY,
+                children=response_children,
+            )
+        except Exception as e:
+            self.log.exception(e)
+            return Response(
+                name=self.name,
+                token=token,
+                data=None,
+                flag=ResponseFlag.DRUNC_EXCEPTION_THROWN,
+                children=None,
+            )

src/drunc/controller/controller_driver.py

@@ -167,3 +167,15 @@ def expert_command(
             outformat=PlainText,
             timeout=timeout,
         )
+
+    @pack_empty_addressed_command
+    def to_error(
+        self, addressed_command: AddressedCommand, timeout: int | float = 60
+    ) -> DecodedResponse:
+        self.log.error(f"{addressed_command=}")
+        return self.send_command(
+            "to_error",
+            data=addressed_command,
+            outformat=OldDescription,
+            timeout=timeout,
+        )

src/drunc/controller/interface/commands.py

@@ -406,3 +406,32 @@ def print_result(result, prefix=""):
             print_result(child, prefix + "    ")
 
     print_result(result)
+
+@click.command("to_error")
+@click.option("--target", type=str, help="The target to address", default="")
+@click.option(
+    "--execute-along-path/--dont-execute-along-path",
+    is_flag=True,
+    show_default=True,
+    help="Execute the command along the path",
+    default=True,
+)
+@click.option(
+    "--execute-on-all-subsequent-children-in-path/--dont-execute-on-all-subsequent-children-in-path",
+    is_flag=True,
+    show_default=True,
+    help="Execute the command on all subsequent children in the path",
+    default=True,
+)
+@click.pass_obj
+def to_error(
+    obj: ControllerContext,
+    target: str,
+    execute_along_path: bool,
+    execute_on_all_subsequent_children_in_path: bool
+) -> None:
+    obj.get_driver("controller").to_error(
+        target=target,
+        execute_along_path=execute_along_path,
+        execute_on_all_subsequent_children_in_path=execute_on_all_subsequent_children_in_path,
+    )

src/drunc/controller/interface/shell_utils.py

@@ -1,6 +1,7 @@
 import datetime
 import logging
 import os
+import sys
 import time
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor
@@ -462,6 +463,13 @@ def run_one_fsm_command(
         f"Running transition '{transition_name}' on controller '{controller_name}', targeting: '{target if target else controller_name}'"
     )
 
+    if obj.batch_mode and obj.get_driver("controller").status().data.in_error:
+        obj.get_driver("controller").status()
+        log.error(
+            "Running in batch mode, and because error state is detected, exiting."
+        )
+        sys.exit(1)
+
     execute_along_path = False
     execute_on_all_subsequent_children_in_path = True
 

src/drunc/unified_shell/commands.py

@@ -1,4 +1,5 @@
 import getpass
+import sys
 
 import click
 from druncschema.process_manager_pb2 import ProcessQuery
@@ -71,3 +72,8 @@ def boot(
         log.info("Booted successfully")
     else:
         log.error("Booted, but the top controller is in error")
+        if obj.batch_mode:
+            log.error(
+                "Unified shell: Running in batch mode, and because error state is detected, exiting."
+            )
+            sys.exit(1)

src/drunc/unified_shell/shell.py

@@ -21,6 +21,7 @@
     status,
     surrender_control,
     take_control,
+    to_error,
     wait,
     who_am_i,
     who_is_in_charge,
@@ -181,6 +182,7 @@ def unified_shell(
         f"[green]process_manager[/green] started, communicating through address [green]{process_manager_address}[/green]"
     )
     ctx.obj.reset(address_pm=process_manager_address)
+    ctx.call_on_close(lambda: on_exit(ctx, unified_shell_log))
 
     desc = None
     try:
@@ -260,6 +262,7 @@ def cleanup():
         "Adding [green]unified_shell[/green] commands to the context"
     )
     ctx.command.add_command(boot, "boot")
+    ctx.obj.dynamic_commands.add("boot")
 
     unified_shell_log.debug(
         "Adding [green]process_manager[/green] commands to the context"
@@ -270,6 +273,12 @@ def cleanup():
     ctx.command.add_command(logs, "logs")
     ctx.command.add_command(restart, "restart")
     ctx.command.add_command(ps, "ps")
+    ctx.obj.dynamic_commands.add("kill")
+    ctx.obj.dynamic_commands.add("terminate")
+    ctx.obj.dynamic_commands.add("flush")
+    ctx.obj.dynamic_commands.add("logs")
+    ctx.obj.dynamic_commands.add("restart")
+    ctx.obj.dynamic_commands.add("ps")
 
     # Not particularly proud of this...
     # We instantiate a stateful node which has the same configuration as the one from this session
@@ -339,12 +348,27 @@ def cleanup():
     ctx.command.add_command(exclude, "exclude")
     ctx.command.add_command(wait, "wait")
     ctx.command.add_command(expert_command, "expert-command")
+    ctx.command.add_command(to_error, "to-error")
+    ctx.obj.dynamic_commands.add("status")
+    ctx.obj.dynamic_commands.add("recompute_status")
+    ctx.obj.dynamic_commands.add("connect")
+    ctx.obj.dynamic_commands.add("disconnect")
+    ctx.obj.dynamic_commands.add("take_control")
+    ctx.obj.dynamic_commands.add("surrender_control")
+    ctx.obj.dynamic_commands.add("who_am_i")
+    ctx.obj.dynamic_commands.add("who_is_in_charge")
+    ctx.obj.dynamic_commands.add("include")
+    ctx.obj.dynamic_commands.add("exclude")
+    ctx.obj.dynamic_commands.add("wait")
+    ctx.obj.dynamic_commands.add("expert_command")
+    ctx.obj.dynamic_commands.add("to_error")
 
     unified_shell_log.info(
         "[green]unified_shell[/green] ready with [green]process_manager[/green] and [green]controller[/green] commands"
     )
-    ctx.call_on_close(lambda: on_exit(ctx, unified_shell_log))
 
+    if any([arg in ctx.obj.dynamic_commands for arg in sys.argv]):
+        ctx.obj.batch_mode = True
 
 def on_exit(ctx, unified_shell_log):
     """Handle exit from the shell."""

src/drunc/utils/shell_utils.py

@@ -201,6 +201,8 @@ def _reset(self, name: str, token_args: dict = {}, driver_args: dict = {}):
 
     def __init__(self, *args, **kwargs):
         log = get_logger("utils.ShellContext")
+        self.dynamic_commands = set()
+        self.batch_mode = False
         try:
             self.reset(*args, **kwargs)
         except Exception as e: