Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove task arguments #1370

Merged
merged 1 commit into from
May 12, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 10 additions & 52 deletions locust/test/test_locust_class.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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):
Expand Down Expand Up @@ -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):
Expand Down
5 changes: 1 addition & 4 deletions locust/test/test_runners.py
Original file line number Diff line number Diff line change
Expand Up @@ -822,10 +822,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):
Expand Down
29 changes: 12 additions & 17 deletions locust/user/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,10 +191,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:
Expand Down Expand Up @@ -235,22 +232,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.

Expand All @@ -259,11 +255,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:
Expand Down Expand Up @@ -343,11 +338,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)