diff --git a/SAS/TMSS/backend/test/test_environment.py b/SAS/TMSS/backend/test/test_environment.py index 10de247fec6845c09231caa39dac28c131df7280..607a1d53e66d6f730eee2183179812eba5be3f84 100644 --- a/SAS/TMSS/backend/test/test_environment.py +++ b/SAS/TMSS/backend/test/test_environment.py @@ -631,7 +631,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, @@ -653,7 +653,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) @@ -663,6 +663,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 @@ -685,6 +686,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): @@ -855,7 +859,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)) @@ -882,6 +886,7 @@ 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) @@ -905,11 +910,12 @@ def main_scheduling_unit_blueprint_simulator(): scheduling_unit_blueprint_id = int(args[0]) - if not (options.observation or options.pipeline or options.QA or options.ingest): + if not (options.observation or options.pipeline or options.QA or options.ingest or options.cleanup): options.observation = True 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) @@ -918,7 +924,7 @@ 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), exchange=options.exchange, broker=options.broker): print("Press Ctrl-C to exit")