Restructured `do_task`

Fixed up `do_task_validate_only`

crc/scripts/set_study_status.py

@@ -7,37 +7,63 @@ from crc.scripts.script import Script
 class MyScript(Script):
 
     def get_description(self):
-        return """This is my description"""
+        return """Set the status of the current study. 
+        Status can be one of `in_progress`, `hold`, `open_for_enrollment`, or `abandoned`."""
 
     def do_task_validate_only(self, task, study_id, workflow_id, *args, **kwargs):
 
-        if 'status' in kwargs.keys():
-            new_status = kwargs['status']
-        elif len(args) > 0:
-            new_status = args[0]
+        new_status = None
+        if 'status' in kwargs.keys() or len(args) > 0:
+            if 'status' in kwargs.keys():
+                new_status = kwargs['status']
+            elif len(args) > 0:
+                new_status = args[0]
+            try:
+                study_status = getattr(StudyStatus, new_status)
+
+            except AttributeError as ae:
+                raise ApiError.from_task(code='invalid_argument',
+                                         message=f"We could not find a status matching `{new_status}`. Original message: {ae}",
+                                         task=task)
+            return study_status.value
         else:
             raise ApiError.from_task(code='missing_argument',
                                      message='You must include the new status when calling `set_study_status` script. '
-                                             'The new status can be `in_progress`, `hold`, `open_for_enrollment`, or `abandoned`.',
+                                             'The new status must be one of `in_progress`, `hold`, `open_for_enrollment`, or `abandoned`.',
                                      task=task)
-        if new_status in ['in_progress', 'hold', 'open_for_enrollment', 'abandoned']:
-            return getattr(StudyStatus, new_status).value
 
     def do_task(self, task, study_id, workflow_id, *args, **kwargs):
-        if 'status' in kwargs.keys():
-            new_status = kwargs['status']
-        elif len(args) > 0:
-            new_status = args[0]
+
+        # Get new status
+        new_status = None
+        if 'status' in kwargs.keys() or len(args) > 0:
+            if 'new_status' in kwargs.keys():
+                new_status = kwargs['new_status']
+            elif len(args) > 0:
+                new_status = args[0]
+
+            # Get StudyStatus object for new_status
+            try:
+                study_status = getattr(StudyStatus, new_status)
+
+            # Invalid argument
+            except AttributeError as ae:
+                raise ApiError.from_task(code='invalid_argument',
+                                         message=f"We could not find a status matching `{new_status}`. Original message: {ae}"
+                                         'The new status must be one of `in_progress`, `hold`, `open_for_enrollment`, or `abandoned`.',
+                                         task=task)
+
+            # Set new status
+            study_model = session.query(StudyModel).filter(StudyModel.id == study_id).first()
+            study_model.status = study_status
+            session.commit()
+
+            return study_model.status.value
+
+        # Missing argument
         else:
             raise ApiError.from_task(code='missing_argument',
                                      message='You must include the new status when calling `set_study_status` script. '
-                                             'The new status can be `in_progress`, `hold`, `open_for_enrollment`, or `abandoned`.',
+                                             'The new status must be one of `in_progress`, `hold`, `open_for_enrollment`, or `abandoned`.',
                                      task=task)
-        study_model = session.query(StudyModel).filter(StudyModel.id == study_id).first()
-        study_status = getattr(StudyStatus, new_status)
-        study_model.status = study_status
-        session.commit()
 
-        print('set_study_status')
-
-        return study_model.status.value