Misc cleanup & refactoring

netbox/extras/management/commands/runscript.py

@@ -33,22 +33,6 @@ class Command(BaseCommand):
         parser.add_argument('script', help="Script to run")
 
     def handle(self, *args, **options):
-        def _output_results(job):
-
-            for test_name, attrs in job.data['logs'].items():
-                self.stdout.write(
-                    "\t{}: {} success, {} info, {} warning, {} failure".format(
-                        test_name, attrs['success'], attrs['info'], attrs['warning'], attrs['failure']
-                    )
-                )
-
-        def _set_job_data(script):
-            logs = script._logs
-            job.data = {
-                'logs': logs,
-                'output': script._output,
-            }
-            return job
 
         def _run_script():
             """
@@ -64,7 +48,7 @@ class Command(BaseCommand):
                 except AbortTransaction:
                     script.log_info("Database changes have been reverted automatically.")
                     clear_events.send(request)
-                _set_job_data(script)
+                job.data = script.get_job_data()
                 job.terminate()
             except Exception as e:
                 stacktrace = traceback.format_exc()
@@ -74,10 +58,17 @@ class Command(BaseCommand):
                 script.log_info("Database changes have been reverted due to error.")
                 logger.error(f"Exception raised during script execution: {e}")
                 clear_events.send(request)
-                _set_job_data(script)
+                job.data = script.get_job_data()
                 job.terminate(status=JobStatusChoices.STATUS_ERRORED, error=repr(e))
 
-            _output_results(job)
+            # Print any test method results
+            for test_name, attrs in job.data['tests'].items():
+                self.stdout.write(
+                    "\t{}: {} success, {} info, {} warning, {} failure".format(
+                        test_name, attrs['success'], attrs['info'], attrs['warning'], attrs['failure']
+                    )
+                )
+
             logger.info(f"Script completed in {job.duration}")
 
         User = get_user_model()

netbox/extras/scripts.py

@@ -387,6 +387,16 @@ class BaseScript:
         self.run_tests()
         self.post_run()
 
+    def get_job_data(self):
+        """
+        Return a dictionary of data to attach to the script's Job.
+        """
+        return {
+            'log': self.messages,
+            'output': self.output,
+            'tests': self.tests,
+        }
+
     #
     # Form rendering
     #
@@ -610,7 +620,7 @@ def run_script(data, job, request=None, commit=True, **kwargs):
 
         return job
 
-    def _run_script():
+    def _run_script(job):
         """
         Core script execution task. We capture this within a subfunction to allow for conditionally wrapping it with
         the event_tracking context manager (which is bypassed if commit == False).
@@ -622,21 +632,17 @@ def run_script(data, job, request=None, commit=True, **kwargs):
                     if not commit:
                         raise AbortTransaction()
             except AbortTransaction:
-                msg = _("Database changes have been reverted automatically.")
+                script.log_info(message=_("Database changes have been reverted automatically."))
-                if is_report(type(script)):
-                    # script and legacy reports have different log function signatures
-                    script.log_info(message=msg)
-                else:
-                    script.log_info(msg)
-
                 if request:
                     clear_events.send(request)
-            job = set_job_data(script)
+
+            job.data = script.get_job_data()
             if script.failed:
                 logger.warning(f"Script failed")
                 job.terminate(status=JobStatusChoices.STATUS_FAILED)
             else:
                 job.terminate()
+
         except Exception as e:
             if type(e) is AbortScript:
                 msg = _("Script aborted with error: ") + str(e)
@@ -648,19 +654,13 @@ def run_script(data, job, request=None, commit=True, **kwargs):
                 logger.error(f"Script aborted with error: {e}")
             else:
                 stacktrace = traceback.format_exc()
-                msg = _("An exception occurred: : ") + f"`{type(e).__name__}: {e}`\n```\n{stacktrace}\n```"
+                script.log_failure(
-                if is_report(type(script)):
+                    message=_("An exception occurred: ") + f"`{type(e).__name__}: {e}`\n```\n{stacktrace}\n```"
-                    script.log_failure(message=msg)
+                )
-                else:
-                    script.log_failure(msg)
                 logger.error(f"Exception raised during script execution: {e}")
-            msg = _("Database changes have been reverted due to error.")
+            script.log_info(message=_("Database changes have been reverted due to error."))
-            if is_report(type(script)):
-                script.log_info(message=msg)
-            else:
-                script.log_info(msg)
 
-            job = set_job_data(script)
+            job.data = script.get_job_data()
             job.terminate(status=JobStatusChoices.STATUS_ERRORED, error=repr(e))
             if request:
                 clear_events.send(request)
@@ -671,9 +671,9 @@ def run_script(data, job, request=None, commit=True, **kwargs):
     # change logging, event rules, etc.
     if commit:
         with event_tracking(request):
-            _run_script()
+            _run_script(job)
     else:
-        _run_script()
+        _run_script(job)
 
     # Schedule the next job if an interval has been set
     if job.interval: