TMSS-1806: added cmdline options to use the script in a more interactive way...

TMSS-1806: added cmdline options to use the script in a more interactive way against an existing TMSS(TestEnvironment) instance

SAS/TMSS/backend/test/acceptance_test/pulsar_timing_scheduling.py

@@ -29,49 +29,48 @@ from lofar.sas.tmss.test.test_environment import TMSSTestEnvironment
 
 tmss_test_env = None
 
-if __name__ == "__main__":
-    os.environ['TZ'] = 'UTC'
-    logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s', level=logging.INFO)
-    # logging.getLogger('lofar.sas.tmss.services.scheduling.constraints').level = logging.DEBUG
-
-    # Check the invocation arguments
-    from optparse import OptionParser, OptionGroup
-    parser = OptionParser('%prog [options]',
-                          description='run the tmss_slack_webhook_service which listens for TMSS event messages on the messagebus, and posts the updates on the slack webhook api.')
-
-    group = OptionGroup(parser, 'TMSS Test Environment options', description='Run this acceptance test against a fresh TMSSTestEnvironment a running one (specfied by -C).')
-    parser.add_option_group(group)
-    group.add_option('-C', '--credentials', dest='dbcredentials', type='string',
-                     default=os.environ.get('TMSS_DBCREDENTIALS', 'TMSS'),
-                     help='django dbcredentials name, default: %default')
-    group.add_option('--wipe', dest='wipe', action='store_true', help='wipe all scheduling units in the fresh/existing TMSS instance. Only available when developing.')
-
-    (options, args) = parser.parse_args()
-
-    if options.dbcredentials is not None:
-        # try to connect to existing TMSS
-        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)
-
-
-        # wipe scheduling units only in Development
-        from lofar.common import isDevelopmentEnvironment
-        if options.wipe and isDevelopmentEnvironment():
-            from lofar.common.postgres import PostgresDatabaseConnection
-            from lofar.common import dbcredentials
-            dbcreds = dbcredentials.DBCredentials().get(options.dbcredentials)
-            with PostgresDatabaseConnection(dbcreds) as db:
-                db.executeQuery('TRUNCATE tmssapp_schedulingunitdraft CASCADE;')
-                db.commit()
-
-    else:
-        # create fresh TMSSTestEnvironment
-        tmss_test_env = TMSSTestEnvironment(start_dynamic_scheduler=True, start_websocket=True,
-                                            start_subtask_scheduler=False, start_workflow_service=True)
-        tmss_test_env.start()
-        tmss_test_env.populate_schemas_and_connectors()
-
-
+os.environ['TZ'] = 'UTC'
+logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s', level=logging.INFO)
+
+# Check the invocation arguments
+from optparse import OptionParser, OptionGroup
+parser = OptionParser('%prog [options]',
+                      description='run the tmss_slack_webhook_service which listens for TMSS event messages on the messagebus, and posts the updates on the slack webhook api.')
+
+parser.add_option('-n', '--num_targets', dest='num_targets', type='int', default=-1, help='read at most this many targets (or all if -1), default: %default')
+parser.add_option('-s', '--simulate', dest='simulate', action='store_true', help='simulate a run for each unit, allowing the scheduler to schedule the next unit when one finishes')
+
+group = OptionGroup(parser, 'TMSS Test Environment options', description='Run this acceptance test against a fresh TMSSTestEnvironment a running one (specfied by -C).')
+parser.add_option_group(group)
+group.add_option('-C', '--credentials', dest='dbcredentials', type='string',
+                 default=os.environ.get('TMSS_DBCREDENTIALS', 'TMSS'),
+                 help='django dbcredentials name, default: %default')
+group.add_option('--wipe', dest='wipe', action='store_true', help='wipe all scheduling units in the fresh/existing TMSS instance. Only available when developing.')
+
+(options, args) = parser.parse_args()
+
+if options.dbcredentials is not None:
+    # try to connect to existing TMSS
+    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)
+
+
+    # wipe scheduling units only in Development
+    from lofar.common import isDevelopmentEnvironment
+    if options.wipe and isDevelopmentEnvironment():
+        from lofar.common.postgres import PostgresDatabaseConnection
+        from lofar.common import dbcredentials
+        dbcreds = dbcredentials.DBCredentials().get(options.dbcredentials)
+        with PostgresDatabaseConnection(dbcreds) as db:
+            db.executeQuery('TRUNCATE tmssapp_schedulingunitdraft CASCADE;')
+            db.commit()
+
+else:
+    # create fresh TMSSTestEnvironment
+    tmss_test_env = TMSSTestEnvironment(start_dynamic_scheduler=True, start_websocket=True,
+                                        start_subtask_scheduler=False, start_workflow_service=True)
+    tmss_test_env.start()
+    tmss_test_env.populate_schemas_and_connectors()
 
 
 from lofar.sas.tmss.tmss.tmssapp import models
@@ -81,7 +80,8 @@ from lofar.sas.tmss.tmss.tmssapp.tasks import create_scheduling_unit_draft_from_
 from lofar.sas.tmss.test.test_utils import assert_SUB_statuses, enable_dynamic_scheduling, disable_scheduling, wait_for_units_to_get_scheduled_and_simulate_run, print_scheduling_results, TestStats
 from dateutil import parser
 
-def populate_pulsar_timing_test_campaign(project: models.Project, between=None):
+
+def populate_pulsar_timing_test_campaign(project: models.Project, between=None, max_num_targets: int=-1):
     scheduling_set = models.SchedulingSet.objects.create(**SchedulingSet_test_data(project=project,
                                                                                    name="pulsar timing run %s" % (min([parser.parse(x["from"], ignoretz=True) for x in between]).date() if between else datetime.today()),))
     with open(os.path.abspath(os.path.dirname(__file__)) + '/pulsar_targets.txt', 'r') as file:
@@ -132,8 +132,7 @@ def populate_pulsar_timing_test_campaign(project: models.Project, between=None):
             scheduling_unit_blueprint = create_scheduling_unit_blueprint_and_tasks_and_subtasks_from_scheduling_unit_draft(scheduling_unit_draft)
 
 
-            # uncomment for faster debugging
-            if scheduling_set.scheduling_unit_drafts.count() >= 5:
+            if max_num_targets > -1 and scheduling_set.scheduling_unit_drafts.count() >= max_num_targets:
                 return
 
     assert_SUB_statuses(models.SchedulingUnitStatus.Choices.SCHEDULABLE.value)
@@ -156,10 +155,11 @@ def run_acceptance_test():
     stats.start_populate = datetime.utcnow()
     # populate one campaign constraint to start pretty much right away, at next midnight.
     # They have to fit in a 48 hour window
-    start_time = datetime.utcnow() #.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(hours=48)
+    start_time = datetime.utcnow().replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(hours=24)
     populate_pulsar_timing_test_campaign(project=project,
                                          between=[{"from": start_time.isoformat(),
-                                                   "to": (start_time+timedelta(hours=48)).isoformat()}])
+                                                   "to": (start_time+timedelta(hours=48)).isoformat()}],
+                                         max_num_targets=options.num_targets)
 
     # populate another campaign constraint to start a month later
     start_time = start_time + timedelta(days=3)
@@ -172,13 +172,13 @@ def run_acceptance_test():
     # start scheduling and wait for it to finish
     enable_dynamic_scheduling()
 
-    wait_for_units_to_get_scheduled_and_simulate_run(tmss_test_env)
-    stats.stop_scheduling = datetime.utcnow()
-    # assert success
-    assert_SUB_statuses(models.SchedulingUnitStatus.Choices.FINISHED.value)
+    if options.simulate:
+        wait_for_units_to_get_scheduled_and_simulate_run(tmss_test_env)
+        stats.stop_scheduling = datetime.utcnow()
+        # assert success
+        assert_SUB_statuses(models.SchedulingUnitStatus.Choices.FINISHED.value)
 
     stats.stop_test = datetime.utcnow()
-
     return stats