diff --git a/locust/exception.py b/locust/exception.py index c845a81a0f..d99a362bd9 100644 --- a/locust/exception.py +++ b/locust/exception.py @@ -76,3 +76,9 @@ def __init__(self, *args: object, addr=None) -> None: class RunnerAlreadyExistsError(Exception): pass + + +class NoTaskToRun(Exception): + def __init__(self, message): + super().__init__(message) + self.exit_code = 1 diff --git a/locust/test/test_locust_class.py b/locust/test/test_locust_class.py index 6fc0c0b829..dc0859744c 100644 --- a/locust/test/test_locust_class.py +++ b/locust/test/test_locust_class.py @@ -3,6 +3,7 @@ from locust.exception import ( CatchResponseError, InterruptTaskSet, + NoTaskToRun, RescheduleTask, RescheduleTaskImmediately, ResponseError, @@ -60,7 +61,7 @@ class MyUser(User): wait_time = constant(0.5) l = MyUser(self.environment) - self.assertRaisesRegex(Exception, "No tasks defined on MyUser.*", l.run) + self.assertRaisesRegex(NoTaskToRun, "No tasks defined on MyUser.*", l.run) MyUser.task = object() self.assertRaisesRegex(Exception, ".*but you have set a 'task' attribute.*", l.run) diff --git a/locust/user/task.py b/locust/user/task.py index 11089b9772..d5dbb526ac 100644 --- a/locust/user/task.py +++ b/locust/user/task.py @@ -1,6 +1,13 @@ from __future__ import annotations -from locust.exception import InterruptTaskSet, MissingWaitTimeError, RescheduleTask, RescheduleTaskImmediately, StopUser +from locust.exception import ( + InterruptTaskSet, + MissingWaitTimeError, + NoTaskToRun, + RescheduleTask, + RescheduleTaskImmediately, + StopUser, +) import logging import random @@ -361,6 +368,8 @@ def run(self): except Exception: logging.error("Uncaught exception in on_stop: \n%s", traceback.format_exc()) raise + except NoTaskToRun: + raise except Exception as e: self.user.environment.events.user_error.fire(user_instance=self, exception=e, tb=e.__traceback__) if self.user.environment.catch_exceptions: @@ -472,13 +481,12 @@ class DefaultTaskSet(TaskSet): def get_next_task(self): if not self.user.tasks: + warning_message = "Use the @task decorator or set the 'tasks' attribute of the User (or mark it as abstract = True if you only intend to subclass it)" if getattr(self.user, "task", None): extra_message = ", but you have set a 'task' attribute on your class - maybe you meant to set 'tasks'?" + raise Exception(f"No tasks defined on {self.user.__class__.__name__}{extra_message}{warning_message}") else: - extra_message = "." - raise Exception( - f"No tasks defined on {self.user.__class__.__name__}{extra_message} Use the @task decorator or set the 'tasks' attribute of the User (or mark it as abstract = True if you only intend to subclass it)" - ) + raise NoTaskToRun(f"No tasks defined on {self.user.__class__.__name__}. {warning_message}") return random.choice(self.user.tasks) def execute_task(self, task):