Search code examples
pythonperformance-testinglocust

Assign Locust task sets programmatically


I am trying to do something where I assign the locust tasks dynamically based on the host value passed through the UI. In this example, if host is passed in as "hello", the test should run the hello task else it should run the world task.

from locust import HttpUser, TaskSet, task, events

class RandomTask1(TaskSet):
  @task(1)
  def soemthing(self):
    print("Hello!")

class RandomTask2(TaskSet):
  @task(1)
  def soemthing(self):
    print("World!")

class LoadTestUser(HttpUser):
  def on_start(self):
    host_config = self.host
    if host_config == "hello":
      tasks = {RandomTask1:1}
    else:
      tasks = {RandomTask2:1}

The example below does not work and I get the following error

Exception: No tasks defined on LoadTestUser. use the @task decorator or set the tasks property of the User (or mark it as abstract = True if you only intend to subclass it)

Any idea how I can achieve something like this? I have simplified this for the example, but for all intents and purposes, let's assume that the locust instance is already running and cannot be stopped or restarted and the tasks need to be assigned dynamically.enter image description here

Edit:

Tried doing this:

class LoadTestUser(HttpUser):
    def on_start(self):
        if self.host == "hello":
            self.tasks = {HelloTask: 1}
        else:
            self.tasks = {WorldTask: 1}

    @task
    def nothing(self):
        pass

class HelloTask(TaskSet):
    @task
    def something(self):
        print("Hello")

class WorldTask(TaskSet):
    @task
    def something(self):
        print("World")

Now I see the following error:

Traceback (most recent call last):
  File "/Users/user/project/venv/lib/python3.8/site-packages/locust/user/task.py", line 285, in run
    self.schedule_task(self.get_next_task())
  File "/Users/user/project/venv/lib/python3.8/site-packages/locust/user/task.py", line 420, in get_next_task
    return random.choice(self.user.tasks)
  File "/Users/user/opt/anaconda3/lib/python3.8/random.py", line 291, in choice
    return seq[i]
KeyError: 0

Solution

  • Create a single task and put the logic in that task for what you want to have it run.

    class LoadTestUser(HttpUser):
    
        def something1(self):
            print("Hello!")
    
        def something2(self):
            print("World!")
    
        @task
        def task_logic(self):
            if self.host == "hello":
                self.something1()
            else:
                self.something2()
    

    However, you can just address the error you're getting directly. You need to have a task defined in the class even if you intend to overwrite or change the tasks with your TaskSets. There's an example in the documentation but just add a task with pass so it doesn't do anything, then your overrides should work.

    class LoadTestUser(HttpUser):
    
        def on_start(self):
            host_config = self.host
            if host_config == "hello":
                self.tasks = {RandomTask1:1}
            else:
                self.tasks = {RandomTask2:1}
    
        @task
        def nothing(self):
            pass
    

    EDIT: This should work but looks like there could be a bug in the current version of Locust where it only accepts a dictionary for tasks when Locust first starts and then only accepts a list afterward. Until it's fixed, the example in the other answer works.