TMSS-417: Add enums, refactoring and code clean up

SAS/TMSS/services/websocket/test/t_websocket_service.py

@@ -27,7 +27,7 @@ logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s', level=loggin
 from lofar.sas.tmss.test.test_utils import TMSSTestEnvironment
 
 from lofar.messaging.messagebus import TemporaryExchange, BusListenerJanitor
-from lofar.sas.tmss.services.websocket_service import create_service
+from lofar.sas.tmss.services.websocket_service import create_service, TMSSEventMessageHandlerForWebsocket
 from lofar.common.test_utils import integration_test
 
 from collections import deque
@@ -46,6 +46,8 @@ class TestSubtaskSchedulingService(unittest.TestCase):
     def __init__(self, methodName: str = ...) -> None:
         self.msg_queue = deque()
         self.sync_event = Event()
+        self.ObjActions = TMSSEventMessageHandlerForWebsocket.ObjActions
+        self.ObjTypes = TMSSEventMessageHandlerForWebsocket.ObjTypes
         super().__init__(methodName)
 
     def start_ws_client(self):
@@ -53,9 +55,7 @@ class TestSubtaskSchedulingService(unittest.TestCase):
 
         def on_message(ws, message):
             logger.info('Received msg from ws')
-            # logger.info('\n\nBefore %s\n\n', self.msg_queue)
             self.msg_queue.append(JSONLoads(message))
-            # logger.info('\n\nAfter %s\n\n', self.msg_queue)
             self.sync_event.set()
 
         def on_error(ws, error):
@@ -67,7 +67,6 @@ class TestSubtaskSchedulingService(unittest.TestCase):
         def on_close(ws):
             logger.info('Closed ws')
 
-        websocket.enableTrace(True)
         def thread_ws_starter():
             self.ws = websocket.WebSocketApp("ws://127.0.0.1:5678/",
                                         on_open=on_open,
@@ -76,6 +75,7 @@ class TestSubtaskSchedulingService(unittest.TestCase):
                                         on_close=on_close)
             self.ws.run_forever()
 
+        websocket.enableTrace(True)
         self.t = Thread(target=thread_ws_starter, daemon=True)
         self.t.start()
 
@@ -115,44 +115,42 @@ class TestSubtaskSchedulingService(unittest.TestCase):
                     raise TimeoutError()
                 self.sync_event.clear()
                 # Assert json_blobs
-                json_blob = {'id': json_test['id'], 'object_type': obj_type, 'action': action}
-                if action == 'create' or action == 'update':
+                json_blob = {'id': json_test['id'], 'object_type': obj_type.value, 'action': action.value}
+                if action == self.ObjActions.CREATE or action == self.ObjActions.UPDATE:
                     json_blob['object'] = json_test
                 self.assertEqual(json_blob, self.msg_queue.popleft())
 
-            # TODO: Add enums for obj_type and action like in ws service
-
             # Test scheduling_unit_draft create
             logger.info('\n\n\nCreate su_draft\n\n')
             su_draft = self.test_data_creator.post_data_and_get_response_as_json_object(
                 self.test_data_creator.SchedulingUnitDraft(), '/scheduling_unit_draft/')
-            test_object(su_draft, 'scheduling_unit_draft', 'create')
+            test_object(su_draft, self.ObjTypes.SCHED_UNIT_DRAFT, self.ObjActions.CREATE)
 
             # Test task_draft create
             logger.info('\n\n\nCreate task_draft\n\n')
             task_draft = self.test_data_creator.post_data_and_get_response_as_json_object(
                 self.test_data_creator.TaskDraft(scheduling_unit_draft_url=su_draft['url']), '/task_draft/')
-            test_object(task_draft, 'task_draft', 'create')
+            test_object(task_draft, self.ObjTypes.TASK_DRAFT, self.ObjActions.CREATE)
 
             # Test scheduling_unit_blueprint create
             logger.info('\n\n\nCreate su_blueprint\n\n')
             su_blueprint = self.test_data_creator.post_data_and_get_response_as_json_object(
                 self.test_data_creator.SchedulingUnitBlueprint(scheduling_unit_draft_url=su_draft['url']),
                 '/scheduling_unit_blueprint/')
-            test_object(su_blueprint, 'scheduling_unit_blueprint', 'create')
+            test_object(su_blueprint, self.ObjTypes.SCHED_UNIT_BLUEPRINT, self.ObjActions.CREATE)
 
             # Test task_blueprint create
             logger.info('\n\n\nCreate task_blueprint\n\n')
             task_blueprint = self.test_data_creator.post_data_and_get_response_as_json_object(
                 self.test_data_creator.TaskBlueprint(scheduling_unit_blueprint_url=su_blueprint['url'],
                                                      draft_url=task_draft['url']), '/task_blueprint/')
-            test_object(task_blueprint, 'task_blueprint', 'create')
+            test_object(task_blueprint, self.ObjTypes.TASK_BLUEPRINT, self.ObjActions.CREATE)
 
             # Test subtask create
             logger.info('\n\n\nCreate subtask\n\n')
             subtask = self.test_data_creator.post_data_and_get_response_as_json_object(
                 self.test_data_creator.Subtask(task_blueprint_url=task_blueprint['url']), '/subtask/')
-            test_object(subtask, 'subtask', 'create')
+            test_object(subtask, self.ObjTypes.SUBTASK, self.ObjActions.CREATE)
 
             # TODO: Add task_draft and su_draft updates
             # Test updates
@@ -161,37 +159,37 @@ class TestSubtaskSchedulingService(unittest.TestCase):
                 logger.info('\n\n\nUpdate subtask\n\n')
                 client.set_subtask_status(subtask['id'], 'scheduled')
                 subtask = requests.get(subtask['url'], auth=self.test_data_creator.auth).json()
-                test_object(subtask, 'subtask', 'update')
+                test_object(subtask, self.ObjTypes.SUBTASK, self.ObjActions.UPDATE)
                 # Test task_blueprint update
                 logger.info('\n\n\nUpdate task_blueprint\n\n')
                 task_blueprint = requests.get(task_blueprint['url'], auth=self.test_data_creator.auth).json()
-                test_object(task_blueprint, 'task_blueprint', 'update')
+                test_object(task_blueprint, self.ObjTypes.TASK_BLUEPRINT, self.ObjActions.UPDATE)
                 # Test scheduling_unit_blueprint update
                 logger.info('\n\n\nUpdate su_blueprint\n\n')
                 su_blueprint = requests.get(su_blueprint['url'], auth=self.test_data_creator.auth).json()
-                test_object(su_blueprint, 'scheduling_unit_blueprint', 'update')
+                test_object(su_blueprint, self.ObjTypes.SCHED_UNIT_BLUEPRINT, self.ObjActions.UPDATE)
 
             # Test deletions
             # Test substask delete
             logger.info('\n\n\nDelete subtask\n\n')
             requests.delete(subtask['url'], auth=self.test_data_creator.auth)
-            test_object({'id': subtask['id']}, 'subtask', 'delete')
+            test_object({'id': subtask['id']}, self.ObjTypes.SUBTASK, self.ObjActions.DELETE)
             # Test task_blueprint delete
             logger.info('\n\n\nDelete task_blueprint\n\n')
             requests.delete(task_blueprint['url'], auth=self.test_data_creator.auth)
-            test_object({'id': task_blueprint['id']}, 'task_blueprint', 'delete')
+            test_object({'id': task_blueprint['id']}, self.ObjTypes.TASK_BLUEPRINT, self.ObjActions.DELETE)
             # Test scheduling_unit_blueprint delete
             logger.info('\n\n\nDelete su_blueprint\n\n')
             requests.delete(su_blueprint['url'], auth=self.test_data_creator.auth)
-            test_object({'id': su_blueprint['id']}, 'scheduling_unit_blueprint', 'delete')
+            test_object({'id': su_blueprint['id']}, self.ObjTypes.SCHED_UNIT_BLUEPRINT, self.ObjActions.DELETE)
             # Test task_draft delete
             logger.info('\n\n\nDelete task_draft\n\n')
             requests.delete(task_draft['url'], auth=self.test_data_creator.auth)
-            test_object({'id': task_draft['id']}, 'task_draft', 'delete')
+            test_object({'id': task_draft['id']}, self.ObjTypes.TASK_DRAFT, self.ObjActions.DELETE)
             # Test scheduling_unit_draft delete
             logger.info('\n\n\nDelete su_draft\n\n')
             requests.delete(su_draft['url'], auth=self.test_data_creator.auth)
-            test_object({'id': su_draft['id']}, 'scheduling_unit_draft', 'delete')
+            test_object({'id': su_draft['id']}, self.ObjTypes.SCHED_UNIT_DRAFT, self.ObjActions.DELETE)
 
 
 if __name__ == '__main__':