diff --git a/src/drunc/controller/controller.py b/src/drunc/controller/controller.py index c7ff89cf3..4f82f4721 100644 --- a/src/drunc/controller/controller.py +++ b/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, + ) diff --git a/src/drunc/controller/controller_driver.py b/src/drunc/controller/controller_driver.py index 31153af48..7213a49f8 100644 --- a/src/drunc/controller/controller_driver.py +++ b/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, + ) \ No newline at end of file diff --git a/src/drunc/controller/interface/commands.py b/src/drunc/controller/interface/commands.py index 26c4e98b9..846a56788 100644 --- a/src/drunc/controller/interface/commands.py +++ b/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, + ) \ No newline at end of file diff --git a/src/drunc/controller/interface/shell_utils.py b/src/drunc/controller/interface/shell_utils.py index d10a58658..60db392e2 100644 --- a/src/drunc/controller/interface/shell_utils.py +++ b/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 diff --git a/src/drunc/unified_shell/commands.py b/src/drunc/unified_shell/commands.py index a5b662b92..141d80e93 100644 --- a/src/drunc/unified_shell/commands.py +++ b/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) \ No newline at end of file diff --git a/src/drunc/unified_shell/shell.py b/src/drunc/unified_shell/shell.py index 259151a73..57c934bf7 100644 --- a/src/drunc/unified_shell/shell.py +++ b/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.""" diff --git a/src/drunc/utils/shell_utils.py b/src/drunc/utils/shell_utils.py index d02957cd2..9a9aadc84 100644 --- a/src/drunc/utils/shell_utils.py +++ b/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: