Merge pull request #1370 from locustio/remove-task-arguments

Remove task arguments

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):

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):

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)