diff --git a/SAS/TMSS/backend/test/test_utils.py b/SAS/TMSS/backend/test/test_utils.py index fb0dbac7715f6acf83c367b7871c016a6ac09f52..d4e5a13929f0650074797f704f9436ee64b8e211 100644 --- a/SAS/TMSS/backend/test/test_utils.py +++ b/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: