如何为Locust的schedule_task传递任务参数？

问题场景与代码 #

测试场景为获取新闻网站的文章列表并对文章进行评论，现有代码如下：

class MyTasks(TaskSet):
  def post_comment(self, article_id: int):
    self.client.post(
      f"/{article_id}/comment",
      json={"body": "some random comment"},
      name=f"comment_for_{article_id}",
    )

  @task
  def get_articles(self):
    with self.client.get(
      "/articles",
      name="get_articles",
    ) as response:
      if response.status_code == 200:
        for article_id in response.json()["article_ids"]:
          # 想要为特定文章调度评论任务，但无法传递文章ID？
          self.schedule_task(
            self.post_comment,
            first=True,
            # 如何在此处为post_comment添加参数？
          )

查看schedule_task的源码后，发现它仅接受无参的可调用对象，无法直接传递参数：

def schedule_task(self, task_callable, first=False):
    """
    Add a task to the User's task execution queue.

    :param task_callable: User task to schedule.
    :param first: Optional keyword argument. If True, the task will be put first in the queue.
    """
    if first:
        self._task_queue.appendleft(task_callable)
    else:
        self._task_queue.append(task_callable)

解决方案 #

Locust本身支持通过Python的函数式编程技巧为任务传递参数，以下是两种常用方法：

方法一：使用lambda表达式 #

通过lambda捕获当前的article_id，生成一个无参的可调用对象传入schedule_task：

@task
def get_articles(self):
    with self.client.get(
      "/articles",
      name="get_articles",
    ) as response:
      if response.status_code == 200:
        for article_id in response.json()["article_ids"]:
          self.schedule_task(
            lambda aid=article_id: self.post_comment(aid),
            first=True
          )

注意：给lambda的参数设置默认值aid=article_id，可避免循环变量引用导致的参数值异常。

方法二：使用functools.partial绑定参数 #

通过functools.partial预先绑定函数参数，生成新的可调用对象：

from functools import partial

class MyTasks(TaskSet):
  # ... 原有代码 ...

  @task
  def get_articles(self):
    with self.client.get(
      "/articles",
      name="get_articles",
    ) as response:
      if response.status_code == 200:
        for article_id in response.json()["article_ids"]:
          self.schedule_task(
            partial(self.post_comment, article_id=article_id),
            first=True
          )

这两种方法都能实现为调度的任务传递自定义参数，无需修改Locust的源码。

内容的提问来源于stack exchange，提问作者rzlvmp