diff --git a/lib/tfw/event_handler_base.py b/lib/tfw/event_handler_base.py index aaf408c..de170f0 100644 --- a/lib/tfw/event_handler_base.py +++ b/lib/tfw/event_handler_base.py @@ -106,25 +106,17 @@ class EventHandlerBase(ABC): pass -class FSMAwareEventHandler(EventHandlerBase, ABC): - # pylint: disable=abstract-method - """ - Abstract base class for EventHandlers which automatically - keep track of the state of the TFW FSM. - """ - def __init__(self, key): - super().__init__(key) - self.subscribe('fsm_update') +class FSMAware: + def __init__(self): self.fsm_state = None self.in_accepted_state = False self._auth_key = KeyManager().auth_key - def dispatch_handling(self, message): - if message['key'] == 'fsm_update': - if verify_message(self._auth_key, message): - self._handle_fsm_update(message) - return None - return super().dispatch_handling(message) + def update_fsm_data(self, message): + if message['key'] == 'fsm_update' and verify_message(self._auth_key, message): + self._handle_fsm_update(message) + return True + return False def _handle_fsm_update(self, message): try: @@ -144,6 +136,23 @@ class FSMAwareEventHandler(EventHandlerBase, ABC): pass +class FSMAwareEventHandler(EventHandlerBase, FSMAware, ABC): + # pylint: disable=abstract-method + """ + Abstract base class for EventHandlers which automatically + keep track of the state of the TFW FSM. + """ + def __init__(self, key): + EventHandlerBase.__init__(self, key) + FSMAware.__init__(self) + self.subscribe('fsm_update') + + def dispatch_handling(self, message): + if self.update_fsm_data(message): + return None + return super().dispatch_handling(message) + + class BroadcastingEventHandler(EventHandlerBase, ABC): # pylint: disable=abstract-method """