Is it possible to assign task weight in Locust from a command-line argument?

Viewed 66

I have managed to pass user-defined weights into an HttpUser class as environment variables by using command-line arguments, but I seem unable to use these weights in the task decorator. For example, my code looks like this:

from locust import HttpUser, task, between, events
from other_module.other_class import OtherClass

@events.init_command_line_parser.add_listener
def _(parser):
    parser.add_argument('--user-defined-weight', type=int,
                        env_var='LOCUST_USER_DEFINED_WEIGHT',
                        default=1,
                        help='User defined weight')

@events.test_start.add_listener
def _(environment, **kw):
    print(
        (
            '\nCustom arguments supplied:\n'
            'User defined weight: '
            f'{environment.parsed_options.user_defined_weight}'
        )
    )

class MyClass(HttpUser):

    ''' My class docstring here '''

    wait_time = between(1, 5)

    def on_start(self):
        self.obj = OtherClass()

    def SomeMethod(self, param):
        data = None
        if param == 'A':
            data = #### Some_data
        else:
            data = #### Some_other_data
        return data


    @task(<use of environment.parsed_options.user_defined_weight here instead of 1, 3, 5 etc>)
    def some_task(self):
        # Do stuff using OtherClass object
        data = self.SomeDethod('A')
        # More stuff happening here etc

When I try this, the line with the @task decorator returns NameError: name 'environment' is not defined, or NameError: name 'self' is not defined if I pass the weight as self.environment.parsed_options.user_defined_weight instead etc.

I have also tried using tasks as a dictionary instead of decorators, with callable tasks as keys and weights as values, but Locust seems to totally ignore the weights. For example, let's say I modify my code to this (using a 2-task example here)

from locust import HttpUser, task, between, events
from other_module.other_class import OtherClass

@events.init_command_line_parser.add_listener
def _(parser):
    parser.add_argument('--user-defined-weight', type=int,
                        env_var='LOCUST_USER_DEFINED_WEIGHT',
                        default=1,
                        help='User defined weight')
    parser.add_argument('--another-user-defined-weight', type=int,
                        env_var='LOCUST_ANOTHER_USER_DEFINED_WEIGHT',
                        default=1,
                        help='Another user defined weight')

@events.test_start.add_listener
def _(environment, **kw):
    print(
        (
            '\nCustom arguments supplied:\n'
            'User defined weight: '
            f'{environment.parsed_options.user_defined_weight}'
            'Another user defined weight: '
            f'{environment.parsed_options.another_user_defined_weight}'
        )
    )

def some_task(user):
    # Do stuff using OtherClass object
    data = user.SomeDethod('A')
    # More stuff happening here etc

def some_other_task(user):
    # Do stuff using OtherClass object
    other_data = user.SomeDethod('B')
    # More stuff happening here etc

class MyClass(HttpUser):

    ''' My class docstring here '''

    wait_time = between(1, 5)

    def on_start(self):
        self.obj = OtherClass()
        tasks = {
            some_task(self): \
                self.environment.parsed_options.user_defined_weight,
            some_other_task(self): \
                self.environment.parsed_options.another_user_defined_weight
        }

    def SomeMethod(self, param):
        data = None
        if param == 'A':
            data = Some_data
        else:
            data = Some_ther_data
        return data


    # HttpUser class needs to have at least one task, even if it does nothing
    @task
    def nothing(self):
        pass

However, if I now set the weight for e.g. some_other_task to zero, Locust will still execute some_other_task...

Any ideas how to pass user-defined weights to tasks, either using the @task decorator, or the tasks list ?

Many thanks!

0 Answers
Related