import gevent from gevent import monkey, GreenletExit import six from six.moves import xrange monkey.patch_all(thread=False) from time import time import sys import random import warnings import traceback import logging from .clients import HttpSession from . import events from .exception import LocustError, InterruptTaskSet, RescheduleTask, RescheduleTaskImmediately, StopLocust logger = logging.getLogger(__name__) def task(weight=1): """ Used as a convenience decorator to be able to declare tasks for a TaskSet inline in the class. Example:: class ForumPage(TaskSet): @task(100) def read_thread(self): pass @task(7) def create_thread(self): pass """ def decorator_func(func): func.locust_task_weight = weight return func """ Check if task was used without parentheses (not called), like this:: @task def my_task() pass """ if callable(weight): func = weight weight = 1 return decorator_func(func) else: return decorator_func class NoClientWarningRaiser(object): """ The purpose of this class is to emit a sensible error message for old test scripts that inherits from Locust, and expects there to be an HTTP client under the client attribute. """ def __getattr__(self, _): raise LocustError("No client instantiated. Did you intend to inherit from HttpLocust?") class Locust(object): """ Represents a "user" which is to be hatched and attack the system that is to be load tested. The behaviour of this user is defined by the task_set attribute, which should point to a :py:class:`TaskSet ` class. This class should usually be subclassed by a class that defines some kind of client. For example when load testing an HTTP system, you probably want to use the :py:class:`HttpLocust ` class. """ host = None """Base hostname to swarm. i.e: http://127.0.0.1:1234""" min_wait = 1000 """Minimum waiting time between the execution of locust tasks""" max_wait = 1000 """Maximum waiting time between the execution of locust tasks""" task_set = None """TaskSet class that defines the execution behaviour of this locust""" stop_timeout = None """Number of seconds after which the Locust will die. If None it won't timeout.""" weight = 10 """Probability of locust being chosen. The higher the weight, the greater is the chance of it being chosen.""" client = NoClientWarningRaiser() _catch_exceptions = True def __init__(self): super(Locust, self).__init__() def run(self): try: self.task_set(self).run() except StopLocust: pass except (RescheduleTask, RescheduleTaskImmediately) as e: six.reraise(LocustError, LocustError("A task inside a Locust class' main TaskSet (`%s.task_set` of type `%s`) seems to have called interrupt() or raised an InterruptTaskSet exception. The interrupt() function is used to hand over execution to a parent TaskSet, and should never be called in the main TaskSet which a Locust class' task_set attribute points to." % (type(self).__name__, self.task_set.__name__)), sys.exc_info()[2]) class HttpLocust(Locust): """ Represents an HTTP "user" which is to be hatched and attack the system that is to be load tested. The behaviour of this user is defined by the task_set attribute, which should point to a :py:class:`TaskSet ` class. This class creates a *client* attribute on instantiation which is an HTTP client with support for keeping a user session between requests. """ client = None """ Instance of HttpSession that is created upon instantiation of Locust. The client support cookies, and therefore keeps the session between HTTP requests. """ def __init__(self): super(HttpLocust, self).__init__() if self.host is None: raise LocustError("You must specify the base host. Either in the host attribute in the Locust class, or on the command line using the --host option.") self.client = HttpSession(base_url=self.host) class TaskSetMeta(type): """ Meta class for the main Locust class. It's used to allow Locust classes to specify task execution ratio using an {task:int} dict, or a [(task0,int), ..., (taskN,int)] list. """ def __new__(mcs, classname, bases, classDict): new_tasks = [] for base in bases: if hasattr(base, "tasks") and base.tasks: new_tasks += base.tasks if "tasks" in classDict and classDict["tasks"] is not None: tasks = classDict["tasks"] if isinstance(tasks, dict): tasks = six.iteritems(tasks) for task in tasks: if isinstance(task, tuple): task, count = task for i in xrange(0, count): new_tasks.append(task) else: new_tasks.append(task) for item in six.itervalues(classDict): if hasattr(item, "locust_task_weight"): for i in xrange(0, item.locust_task_weight): new_tasks.append(item) classDict["tasks"] = new_tasks return type.__new__(mcs, classname, bases, classDict) @six.add_metaclass(TaskSetMeta) class TaskSet(object): """ Class defining a set of tasks that a Locust user will execute. When a TaskSet starts running, it will pick a task from the *tasks* attribute, execute it, call it's wait function which will sleep a random number between *min_wait* and *max_wait* milliseconds. It will then schedule another task for execution and so on. TaskTests can be nested, which means that a TaskSet's *tasks* attribute can contain another TaskSet. If the nested TaskSet it scheduled to be executed, it will be instantiated and called from the current executing TaskSet. Execution in the currently running TaskSet will then be handed over to the nested TaskSet which will continue to run until it throws an InterruptTaskSet exception, which is done when :py:meth:`TaskSet.interrupt() ` is called. (execution will then continue in the first TaskSet). """ tasks = [] """ List with python callables that represents a locust user task. If tasks is a list, the task to be performed will be picked randomly. If tasks is a *(callable,int)* list of two-tuples, or a {callable:int} dict, the task to be performed will be picked randomly, but each task will be weighted according to it's corresponding int value. So in the following case *ThreadPage* will be fifteen times more likely to be picked than *write_post*:: class ForumPage(TaskSet): tasks = {ThreadPage:15, write_post:1} """ min_wait = None """ Minimum waiting time between the execution of locust tasks. Can be used to override the min_wait defined in the root Locust class, which will be used if not set on the TaskSet. """ max_wait = None """ Maximum waiting time between the execution of locust tasks. Can be used to override the max_wait defined in the root Locust class, which will be used if not set on the TaskSet. """ locust = None """Will refer to the root Locust class instance when the TaskSet has been instantiated""" parent = None """ Will refer to the parent TaskSet, or Locust, class instance when the TaskSet has been instantiated. Useful for nested TaskSet classes. """ def __init__(self, parent): self._task_queue = [] self._time_start = time() if isinstance(parent, TaskSet): self.locust = parent.locust elif isinstance(parent, Locust): self.locust = parent else: raise LocustError("TaskSet should be called with Locust instance or TaskSet instance as first argument") self.parent = parent # if this class doesn't have a min_wait or max_wait defined, copy it from Locust if not self.min_wait: self.min_wait = self.locust.min_wait if not self.max_wait: self.max_wait = self.locust.max_wait def run(self, *args, **kwargs): self.args = args self.kwargs = kwargs try: if hasattr(self, "on_start"): self.on_start() except InterruptTaskSet as e: if e.reschedule: six.reraise(RescheduleTaskImmediately, RescheduleTaskImmediately(e.reschedule), sys.exc_info()[2]) else: six.reraise(RescheduleTask, RescheduleTask(e.reschedule), sys.exc_info()[2]) while (True): try: if self.locust.stop_timeout is not None and time() - self._time_start > self.locust.stop_timeout: return if not self._task_queue: self.schedule_task(self.get_next_task()) try: self.execute_next_task() except RescheduleTaskImmediately: pass except RescheduleTask: self.wait() else: self.wait() except InterruptTaskSet as e: if e.reschedule: six.reraise(RescheduleTaskImmediately, RescheduleTaskImmediately(e.reschedule), sys.exc_info()[2]) else: six.reraise(RescheduleTask, RescheduleTask(e.reschedule), sys.exc_info()[2]) except StopLocust: raise except GreenletExit: raise except Exception as e: events.locust_error.fire(locust_instance=self, exception=e, tb=sys.exc_info()[2]) if self.locust._catch_exceptions: sys.stderr.write("\n" + traceback.format_exc()) self.wait() else: raise def execute_next_task(self): task = self._task_queue.pop(0) self.execute_task(task["callable"], *task["args"], **task["kwargs"]) def execute_task(self, task, *args, **kwargs): # check if the function is a method bound to the current locust, and if so, don't pass self as first argument if hasattr(task, "__self__") and task.__self__ == self: # task is a bound method on self task(*args, **kwargs) elif hasattr(task, "tasks") and issubclass(task, TaskSet): # task is another (nested) TaskSet class task(self).run(*args, **kwargs) else: # task is a function task(self, *args, **kwargs) def schedule_task(self, task_callable, args=None, kwargs=None, first=False): """ Add a task to the Locust's task execution queue. *Arguments*: * task_callable: Locust task to schedule * args: Arguments that will be passed to the task callable * kwargs: Dict of keyword arguments that will be passed to the task callable. * first: Optional keyword argument. If True, the task will be put first in the queue. """ task = {"callable":task_callable, "args":args or [], "kwargs":kwargs or {}} if first: self._task_queue.insert(0, task) else: self._task_queue.append(task) def get_next_task(self): return random.choice(self.tasks) def wait(self): millis = random.randint(self.min_wait, self.max_wait) seconds = millis / 1000.0 self._sleep(seconds) def _sleep(self, seconds): gevent.sleep(seconds) def interrupt(self, reschedule=True): """ Interrupt the TaskSet and hand over execution control back to the parent TaskSet. If *reschedule* is True (default), the parent Locust will immediately re-schedule, and execute, a new task This method should not be called by the root TaskSet (the one that is immediately, attached to the Locust class' *task_set* attribute), but rather in nested TaskSet classes further down the hierarchy. """ raise InterruptTaskSet(reschedule) @property def client(self): """ Reference to the :py:attr:`client ` attribute of the root Locust instance. """ return self.locust.client