TMSS-417: Refactoring

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

@@ -29,14 +29,13 @@ 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.common.test_utils import integration_test
-from time import sleep
-from datetime import datetime, timedelta
 
 from collections import deque
 from json import loads as JSONLoads
 from threading import Thread, Event
 import websocket
 
+
 @integration_test
 class TestSubtaskSchedulingService(unittest.TestCase):
     '''
@@ -107,36 +106,29 @@ class TestSubtaskSchedulingService(unittest.TestCase):
 
             self.start_ws_client()  # Start ws client
 
+            def test_object(json_test, obj_type, action):   # Check if the correct/expected json_blobs arrive in the ws client
+                # Wait for incoming ws message
+                if not self.sync_event.wait(timeout=5):
+                    raise TimeoutError()
+                self.sync_event.clear()
+                # Assert json_blobs
+                json_blob = {'id': json_test['id'], 'object_type': obj_type, 'action': action, '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
             json_schedulingset = self.test_data_creator.post_data_and_get_response_as_json_object(self.test_data_creator.SchedulingSet(), url_postfix='/scheduling_set')
             json_schedulingunitdraft = self.test_data_creator.post_data_and_get_response_as_json_object(self.test_data_creator.SchedulingUnitDraft(scheduling_set_url=json_schedulingset['url']), url_postfix='/scheduling_unit_draft')
-            if not self.sync_event.wait(timeout=5):
-                raise TimeoutError()
-            self.sync_event.clear()
-            json_blob = {'id': json_schedulingunitdraft['id'], 'object_type': 'scheduling_unit_draft',
-                         'action': 'create'}
-            json_blob['object'] = json_schedulingunitdraft
-            recv_msg = self.msg_queue.popleft()
-            self.assertEqual(json_blob, recv_msg)
-
-            # TODO: Refactoring
-
-            json_blob = {}
-            json_taskdraft = self.test_data_creator.post_data_and_get_response_as_json_object(
-                self.test_data_creator.TaskDraft(scheduling_unit_draft_url=json_schedulingunitdraft['url']),
-                url_postfix='/task_draft')
-            if not self.sync_event.wait(timeout=5):
-                raise TimeoutError()
-            self.sync_event.clear()
-            json_blob = {'id': json_taskdraft['id'], 'object_type': 'task_draft',
-                         'action': 'create'}
-            json_blob['object'] = json_taskdraft
-            recv_msg = self.msg_queue.popleft()
-            self.assertEqual(json_blob, recv_msg)
+            test_object(json_schedulingunitdraft, 'scheduling_unit_draft', 'create')
 
+            # Test task_draft create
+            json_taskdraft = self.test_data_creator.post_data_and_get_response_as_json_object(
+                 self.test_data_creator.TaskDraft(scheduling_unit_draft_url=json_schedulingunitdraft['url']),
+                 url_postfix='/task_draft')
+            test_object(json_taskdraft, 'task_draft', 'create')
 
             # TODO: create/update/delete objects like SubTask, TaskBlueprint etc
-            # TODO: check if the correct/expected json_blobs arrive in the websocket client
 
 
 if __name__ == '__main__':