diff --git a/locust/test/test_locust_class.py b/locust/test/test_locust_class.py index 46771054d9..29fdc42190 100644 --- a/locust/test/test_locust_class.py +++ b/locust/test/test_locust_class.py @@ -225,7 +225,7 @@ def test_schedule_task(self): def t1(l): self.t1_executed = True - def t2(l, arg): + def t2(l, arg="t2 argument"): self.t2_arg = arg class MyTasks(TaskSet): @@ -236,29 +236,9 @@ class MyTasks(TaskSet): taskset.execute_next_task() self.assertTrue(self.t1_executed) - taskset.schedule_task(t2, args=["argument to t2"]) + taskset.schedule_task(t2) taskset.execute_next_task() - self.assertEqual("argument to t2", self.t2_arg) - - def test_schedule_task_with_kwargs(self): - class MyTasks(TaskSet): - @task - def t1(self): - self.t1_executed = True - @task - def t2(self, *args, **kwargs): - self.t2_args = args - self.t2_kwargs = kwargs - loc = MyTasks(self.locust) - loc.schedule_task(loc.t2, [42], {"test_kw":"hello"}) - loc.execute_next_task() - self.assertEqual((42, ), loc.t2_args) - self.assertEqual({"test_kw":"hello"}, loc.t2_kwargs) - - loc.schedule_task(loc.t2, args=[10, 4], kwargs={"arg1":1, "arg2":2}) - loc.execute_next_task() - self.assertEqual((10, 4), loc.t2_args) - self.assertEqual({"arg1":1, "arg2":2}, loc.t2_kwargs) + self.assertEqual("t2 argument", self.t2_arg) def test_schedule_task_bound_method(self): class MyTasks(TaskSet): @@ -358,44 +338,22 @@ def a_task(self): self.assertRaises(RescheduleTaskImmediately, lambda: loc.execute_next_task()) self.assertTrue(self.locust.sub_locust_task_executed) - def test_sub_taskset_arguments(self): - class MySubTaskSet(TaskSet): - wait_time = constant(0.001) - @task() - def a_task(self): - self.user.sub_taskset_args = self.args - self.user.sub_taskset_kwargs = self.kwargs - self.interrupt() - class MyTaskSet(TaskSet): - sub_locust_args = None - sub_locust_kwargs = None - tasks = [MySubTaskSet] - - self.locust.sub_taskset_args = None - self.locust.sub_taskset_kwargs = None - - loc = MyTaskSet(self.locust) - loc.schedule_task(MySubTaskSet, args=[1,2,3], kwargs={"hello":"world"}) - self.assertRaises(RescheduleTaskImmediately, lambda: loc.execute_next_task()) - self.assertEqual((1,2,3), self.locust.sub_taskset_args) - self.assertEqual({"hello":"world"}, self.locust.sub_taskset_kwargs) - def test_on_start_interrupt(self): class SubTaskSet(TaskSet): def on_start(self): - if self.kwargs["reschedule"]: - self.interrupt(reschedule=True) - else: - self.interrupt(reschedule=False) - + self.interrupt(reschedule=False) + class RescheduleSubTaskSet(TaskSet): + def on_start(self): + self.interrupt(reschedule=True) class MyUser(User): host = "" tasks = [SubTaskSet] l = MyUser(Environment()) task_set = SubTaskSet(l) - self.assertRaises(RescheduleTaskImmediately, lambda: task_set.run(reschedule=True)) - self.assertRaises(RescheduleTask, lambda: task_set.run(reschedule=False)) + reschedule_task_set = RescheduleSubTaskSet(l) + self.assertRaises(RescheduleTaskImmediately, lambda: reschedule_task_set.run()) + self.assertRaises(RescheduleTask, lambda: task_set.run()) def test_parent_attribute(self): diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py index 87f81d8eb4..5b72425583 100644 --- a/locust/test/test_runners.py +++ b/locust/test/test_runners.py @@ -824,10 +824,7 @@ def test_exception_is_catched(self): class MyTaskSet(TaskSet): def __init__(self, *a, **kw): super(MyTaskSet, self).__init__(*a, **kw) - self._task_queue = [ - {"callable":self.will_error, "args":[], "kwargs":{}}, - {"callable":self.will_stop, "args":[], "kwargs":{}}, - ] + self._task_queue = [self.will_error, self.will_stop] @task(1) def will_error(self): diff --git a/locust/user/task.py b/locust/user/task.py index 5f6c80fffa..e474cba88a 100644 --- a/locust/user/task.py +++ b/locust/user/task.py @@ -259,10 +259,7 @@ def on_stop(self): """ pass - def run(self, *args, **kwargs): - self.args = args - self.kwargs = kwargs - + def run(self): try: self.on_start() except InterruptTaskSet as e: @@ -303,22 +300,21 @@ def run(self, *args, **kwargs): raise def execute_next_task(self): - task = self._task_queue.pop(0) - self.execute_task(task["callable"], *task["args"], **task["kwargs"]) + self.execute_task(self._task_queue.pop(0)) - def execute_task(self, task, *args, **kwargs): + def execute_task(self, task): # 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) + task() elif hasattr(task, "tasks") and issubclass(task, TaskSet): # task is another (nested) TaskSet class - task(self).run(*args, **kwargs) + task(self).run() else: # task is a function - task(self, *args, **kwargs) + task(self) - def schedule_task(self, task_callable, args=None, kwargs=None, first=False): + def schedule_task(self, task_callable, first=False): """ Add a task to the User's task execution queue. @@ -327,11 +323,10 @@ def schedule_task(self, task_callable, args=None, kwargs=None, first=False): :param kwargs: Dict of keyword arguments that will be passed to the task callable. :param 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) + self._task_queue.insert(0, task_callable) else: - self._task_queue.append(task) + self._task_queue.append(task_callable) def get_next_task(self): if not self.tasks: @@ -411,11 +406,11 @@ class DefaultTaskSet(TaskSet): def get_next_task(self): return random.choice(self.user.tasks) - def execute_task(self, task, *args, **kwargs): + def execute_task(self, task): if hasattr(task, "tasks") and issubclass(task, TaskSet): # task is (nested) TaskSet class - task(self.user).run(*args, **kwargs) + task(self.user).run() else: # task is a function - task(self.user, *args, **kwargs) + task(self.user)