Refactor FSM instance creation

src/app/app.py

@@ -7,11 +7,18 @@ from tornado.web import Application
 from tornado.ioloop import IOLoop
 
 from config import WEB_PORT
-from handlers import ZMQWebSocketHandler
+from handlers import FSMManagingSocketHandler
+from sql_injection_fsm import SQLInjectionFSM
+
+
+class SQLInjectionFSMManagingSocketHandler(FSMManagingSocketHandler):
+    def __init__(self, application, request, **kwargs):
+        super().__init__(SQLInjectionFSM, application, request, **kwargs)
+
 
 if __name__ == '__main__':
     routes = [
-        (r'/ws', ZMQWebSocketHandler),
+        (r'/ws', SQLInjectionFSMManagingSocketHandler),
     ]
     application = Application(
         routes,

src/app/handlers/__init__.py

@@ -1 +1 @@
-from .zmq_websocket_handler import ZMQWebSocketHandler
+from .zmq_websocket_handler import FSMManagingSocketHandler

src/app/handlers/zmq_websocket_handler.py

@@ -3,8 +3,6 @@ import json
 from tornado.websocket import WebSocketHandler
 from util import parse_anchor_from_message
 
-# from buttons import fsm
-from sql_injection_fsm import fsm
 import event_handler_connector
 
 
@@ -23,12 +21,10 @@ class ZMQWebSocketHandler(WebSocketHandler):
 
     def on_message(self, message):
         logging.debug('Received on WebSocket: {}'.format(message))
-        fsm.trigger(parse_anchor_from_message(message), message=message)
+        self.send_message(*self.make_response(message))
-        anchor = 'FSMUpdate'
+
-        response = { 'anchor': anchor,
+    def make_response(self, message):
-                     'data': { 'current_state': fsm.state,
+        raise NotImplementedError
-                               'valid_transitions': [{'trigger': trigger} for trigger in fsm.machine.get_triggers()]}}
-        self.send_message(json.dumps(response), anchor)
 
     def send_message(self, message: str, anchor: str = None):
         event_handler_connector.send_message(message, anchor)
@@ -40,3 +36,17 @@ class ZMQWebSocketHandler(WebSocketHandler):
     def check_origin(self, origin):
         return True
 
+
+class FSMManagingSocketHandler(ZMQWebSocketHandler):
+    def __init__(self, FSM, application, request, **kwargs):
+        super().__init__(application, request, **kwargs)
+        self.fsm = FSM()
+
+    def make_response(self, message):
+        self.fsm.trigger(parse_anchor_from_message(message), message=message)
+        anchor = 'FSMUpdate'
+        response = {'anchor': anchor,
+                    'data': {'current_state': self.fsm.state,
+                             'valid_transitions':
+                                 [{'trigger': trigger} for trigger in self.fsm.machine.get_triggers()]}}
+        return json.dumps(response), anchor

src/app/sql_injection_fsm.py

@@ -24,6 +24,3 @@ class SQLInjectionFSM(FSMBase):
 
     def __init__(self):
         super().__init__('start')
-
-
-fsm = SQLInjectionFSM()