以编程方式分配 Locust 任务集

Question

我正在尝试根据通过 UI 传递的主机值动态分配蝗虫任务。在这个例子中，如果 host 作为“hello”传入，测试应该运行 hello 任务，否则它应该运行 world 任务。

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}

下面的示例不起作用，我收到以下错误

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)

知道如何实现这样的目标吗？我已经为示例简化了这一点，但出于所有意图和目的，我们假设 locust 实例已经在运行并且无法停止或重新启动，并且需要动态分配任务。

编辑：

试过这样做：

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

现在我看到以下错误：

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

Answer 1

创建一个任务并将逻辑放入该任务中以使其运行。

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

但是，您可以直接解决遇到的错误。即使您打算使用任务集覆盖或更改任务，您也需要在类中定义任务。有一个示例 in the documentation 但只需使用 pass 添加一个任务，这样它就不会做任何事情，那么您的覆盖应该可以工作。

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

编辑： 这应该工作，但看起来在当前版本的 Locust 中可能存在一个错误，它只在 Locust 首次启动时接受任务字典，然后只接受一个列表.在修复之前，其他答案中的示例有效。

Answer 2

import sys
from locust import HttpUser, TaskSet, task, events

class LoadTestUser(HttpUser):
    def locust_class(self, name):
        module = sys.modules[__name__]
        return getattr(module, f"{name.capitalize()}Task")

    def get_weighted_tasks(self, task_list):
        new_tasks = []
        for item in task_list:
            if "locust_task_weight" in dir(item):
                for i in range(item.locust_task_weight):
                    new_tasks.append(item)
        return new_tasks

    def get_locust_tasks(self, cls):
        tasks = []
        for maybe_task in cls.__dict__.values():
            if hasattr(maybe_task, "locust_task_weight"):
                tasks.append(maybe_task)
        return tasks

    def on_start(self):
        task_cls = self.locust_class(self.host)
        task_list = self.get_locust_tasks(task_cls)
        self.tasks = self.get_weighted_tasks(task_list)

    @task(1)
    def nothing(self):
        pass


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

    @task(100)
    def something_else(self):
        print("hello")


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

    @task(10)
    def something_else(self):
        print("world")