diff --git a/moler/helpers.py b/moler/helpers.py index 5a47c9a93..286c3bff9 100644 --- a/moler/helpers.py +++ b/moler/helpers.py @@ -574,3 +574,34 @@ def regexp_without_anchors(regexp): if regexp_str == org_regexp_str: return regexp return re.compile(regexp_str) + + +def remove_state_from_sm(source_sm: dict, state_to_remove: str) -> dict: + """ + Remove a state from a state machine dict. + :param source_sm: a dict with state machine description + :param state_to_remove: name of state to remove + :return: a new state machine dict without state_to_remove + """ + new_sm = copy.deepcopy(source_sm) + from_states = [] + for from_state in source_sm.keys(): + for to_state in source_sm[from_state].keys(): + if to_state == state_to_remove: + from_states.append(from_state) + + for to_state in source_sm[state_to_remove].keys(): + if to_state == state_to_remove: + continue + for new_from in from_states: + if new_from != to_state and new_from != state_to_remove: + break + if new_from not in new_sm: + new_sm[new_from] = {} + new_sm[new_from][to_state] = copy.deepcopy(source_sm[state_to_remove][to_state]) + + del new_sm[state_to_remove] + for from_state in from_states: + del new_sm[from_state][state_to_remove] + + return new_sm diff --git a/test/test_helpers.py b/test/test_helpers.py index 21f646948..a55fadb56 100644 --- a/test/test_helpers.py +++ b/test/test_helpers.py @@ -428,3 +428,131 @@ def test_date_parser_utc(): date_parsed = ConverterHelper.parse_date(date_str) date_expected = datetime(year=2024, month=5, day=22, hour=9, minute=11, second=48, tzinfo=tzoffset('UTC', 0)) assert date_parsed == date_expected + + +def test_remove_state_from_sm_dict(): + from moler.device.unixremote import UnixRemote + from moler.helpers import remove_state_from_sm + source_sm = { + UnixRemote.unix_local: { + UnixRemote.proxy_pc: { + "execute_command": "ssh", + "command_params": { + "target_newline": "\n" + }, + "required_command_params": [ + "host", + "login", + "password", + "expected_prompt" + ] + }, + }, + UnixRemote.proxy_pc: { # from + UnixRemote.unix_remote: { # to + "execute_command": "ssh", # using command + "command_params": { # with parameters + "target_newline": "\n" + }, + "required_command_params": [ + "host", + "login", + "password", + "expected_prompt" + ] + }, + UnixRemote.unix_local: { # to + "execute_command": "exit", # using command + "command_params": { # with parameters + "target_newline": "\n" + }, + "required_command_params": [ # with parameters + "expected_prompt" + ] + }, + }, + UnixRemote.unix_remote: { # from + UnixRemote.proxy_pc: { # to + "execute_command": "exit", # using command + "command_params": { # with parameters + "target_newline": "\n" + }, + "required_command_params": [ + "expected_prompt" + ] + }, + UnixRemote.unix_remote_root: { # to + "execute_command": "su", # using command + "command_params": { # with parameters + "password": "root_password", + "expected_prompt": r'remote_root_prompt', + "target_newline": "\n" + }, + "required_command_params": [ + ] + }, + }, + UnixRemote.unix_remote_root: { # from + UnixRemote.unix_remote: { # to + "execute_command": "exit", # using command + "command_params": { # with parameters + "target_newline": "\n", + "expected_prompt": r'remote_user_prompt' + }, + "required_command_params": [ + ] + } + } + } + + expected_sm = { + UnixRemote.unix_local: { + UnixRemote.unix_remote: { # to + "execute_command": "ssh", # using command + "command_params": { # with parameters + "target_newline": "\n" + }, + "required_command_params": [ + "host", + "login", + "password", + "expected_prompt" + ] + }, + }, + UnixRemote.unix_remote: { # from + UnixRemote.unix_local: { # to + "execute_command": "exit", # using command + "command_params": { # with parameters + "target_newline": "\n" + }, + "required_command_params": [ + "expected_prompt" + ] + }, + UnixRemote.unix_remote_root: { # to + "execute_command": "su", # using command + "command_params": { # with parameters + "password": "root_password", + "expected_prompt": r'remote_root_prompt', + "target_newline": "\n" + }, + "required_command_params": [ + ] + }, + }, + UnixRemote.unix_remote_root: { # from + UnixRemote.unix_remote: { # to + "execute_command": "exit", # using command + "command_params": { # with parameters + "target_newline": "\n", + "expected_prompt": r'remote_user_prompt' + }, + "required_command_params": [ + ] + } + } + } + + current_sm = remove_state_from_sm(source_sm, UnixRemote.proxy_pc) + assert expected_sm == current_sm