TMSS-261: simulate cleanup

SAS/TMSS/backend/test/test_utils.py

@@ -657,7 +657,7 @@ def main_test_environment():
 
 def create_scheduling_unit_blueprint_simulator(scheduling_unit_blueprint_id: int, stop_event: threading.Event,
                                                handle_observations: bool = True, handle_pipelines: bool = True,
-                                               handle_QA: bool = True, handle_ingest: bool = True,
+                                               handle_QA: bool = True, handle_ingest: bool = True, handle_cleanup: bool = True,
                                                auto_grant_ingest_permission: bool = True,
                                                delay: float=1, duration: float=5,
                                                create_output_dataproducts: bool=False,
@@ -679,7 +679,7 @@ def create_scheduling_unit_blueprint_simulator(scheduling_unit_blueprint_id: int
     class SimulationEventHandler(TMSSEventMessageHandler):
         def __init__(self, scheduling_unit_blueprint_id: int, stop_event: threading.Event,
                      handle_observations: bool = True, handle_pipelines: bool = True,
-                     handle_QA: bool = True, handle_ingest: bool = True,
+                     handle_QA: bool = True, handle_ingest: bool = True, handle_cleanup: bool = True,
                      delay: float = 1, duration: float = 10,
                      create_output_dataproducts: bool=False) -> None:
             super().__init__(log_event_messages=False)
@@ -689,6 +689,7 @@ def create_scheduling_unit_blueprint_simulator(scheduling_unit_blueprint_id: int
             self.handle_pipelines = handle_pipelines
             self.handle_QA = handle_QA
             self.handle_ingest = handle_ingest
+            self.handle_cleanup = handle_cleanup
             self.auto_grant_ingest_permission = auto_grant_ingest_permission
             self.delay = delay
             self.duration = duration
@@ -711,6 +712,9 @@ def create_scheduling_unit_blueprint_simulator(scheduling_unit_blueprint_id: int
             if subtask.specifications_template.type.value == models.SubtaskType.Choices.INGEST.value and not self.handle_ingest:
                 return False
 
+            if subtask.specifications_template.type.value == models.SubtaskType.Choices.CLEANUP.value and not self.handle_cleanup:
+                return False
+
             return True
 
         def start_handling(self):
@@ -881,7 +885,7 @@ def create_scheduling_unit_blueprint_simulator(scheduling_unit_blueprint_id: int
     return BusListenerJanitor(TMSSBusListener(SimulationEventHandler, handler_kwargs={'scheduling_unit_blueprint_id': scheduling_unit_blueprint_id,
                                                                                       'stop_event': stop_event,
                                                                                       'handle_observations': handle_observations, 'handle_pipelines': handle_pipelines,
-                                                                                      'handle_QA': handle_QA, 'handle_ingest': handle_ingest,
+                                                                                      'handle_QA': handle_QA, 'handle_ingest': handle_ingest, 'handle_cleanup': handle_cleanup,
                                                                                       'create_output_dataproducts': create_output_dataproducts,
                                                                                       'delay': delay, 'duration': duration},
                                                                        exchange=exchange, broker=broker))
@@ -908,12 +912,14 @@ def main_scheduling_unit_blueprint_simulator():
     group.add_option('-p', '--pipeline', dest='pipeline', action='store_true', help='simulate events for pipeline subtasks')
     group.add_option('-Q', '--QA', dest='QA', action='store_true', help='simulate events for QA subtasks')
     group.add_option('-i', '--ingest', dest='ingest', action='store_true', help='simulate events for ingest subtasks')
+    group.add_option('-c', '--cleanup', dest='cleanup', action='store_true', help='simulate events for cleanup subtasks')
 
     group = OptionGroup(parser, 'Simulation parameters')
     parser.add_option_group(group)
     group.add_option('-e', '--event_delay', dest='event_delay', type='float', default=1.0, help='wait <event_delay> seconds between simulating events to mimic real-world behaviour, default: %default')
     group.add_option('-d', '--duration', dest='duration', type='float', default=60.0, help='wait <duration> seconds while "observing"/"processing" between started and finishing state to mimic real-world behaviour, default: %default')
     group.add_option('-g', '--grant_ingest_permission', dest='grant_ingest_permission', action='store_true', help='automatically grant ingest permission for ingest subtasks if needed')
+    group.add_option('-f', '--files', dest='files', action='store_true', help='create output dataproduct files')
 
     group = OptionGroup(parser, 'Messaging options')
     parser.add_option_group(group)
@@ -936,6 +942,7 @@ def main_scheduling_unit_blueprint_simulator():
         options.pipeline = True
         options.QA = True
         options.ingest = True
+        options.cleanup = True
 
     from lofar.sas.tmss.tmss import setup_and_check_tmss_django_database_connection_and_exit_on_error
     setup_and_check_tmss_django_database_connection_and_exit_on_error(options.dbcredentials)
@@ -944,8 +951,9 @@ def main_scheduling_unit_blueprint_simulator():
     with create_scheduling_unit_blueprint_simulator(scheduling_unit_blueprint_id, stop_event=stop_event,
                                                     delay=options.event_delay, duration=options.duration,
                                                     handle_observations=bool(options.observation), handle_pipelines=bool(options.pipeline),
-                                                    handle_QA=bool(options.QA), handle_ingest=bool(options.ingest),
+                                                    handle_QA=bool(options.QA), handle_ingest=bool(options.ingest), handle_cleanup=bool(options.cleanup),
                                                     auto_grant_ingest_permission=bool(options.grant_ingest_permission),
+                                                    create_output_dataproducts=bool(options.files),
                                                     exchange=options.exchange, broker=options.broker):
         print("Press Ctrl-C to exit")
         try: