Locust decorators

[rahul-jain101]

Is there a way to me for class-decorate the task class? Just wondering because when I do, the class decorator gets ignored, but when I do
@task
@decorator
underneath each method it seems to work, but when I do:

@class_decorator
class CustomTaskSet(TaskSet):
@task
def task_1(self):
stuff()

[heyman]

Hi! From your description, I'm not able to decipher what the issue is. Could you provide a minimal reproducible code example? I've tested with the following code which seems to work fine:

from locust import HttpLocust, TaskSet, task
from locust.contrib.fasthttp import FastHttpLocust


def class_decorator(original_class):
    orig_init = original_class.__init__
    def __init__(self, *args, **kws):
        print("class_decorator")
        orig_init(self, *args, **kws)

    original_class.__init__ = __init__
    return original_class

@class_decorator
class UserTasks(TaskSet):
    @task
    def index(self):
        self.client.get("/")
    
    @task
    def stats(self):
        self.client.get("/stats/requests")


class WebsiteUser(FastHttpLocust):
    """
    Locust user class that does requests to the locust web server running on localhost
    """
    host = "http://127.0.0.1:8089"
    min_wait = 2000
    max_wait = 5000
    task_set = UserTasks