Locust: Need a way for "Device" Locust class to pass a value to the TaskSet

[quanstorm]

Description of issue / feature request

Working on a MQTT Locust client for load test against an MQTT broker.

• Inheritance chain: Device class inherits MQTTLocust class, which inherits the Locust class
• DeviceBehavior class is a TaskSet class
• The Device class represents a device that publish messages to the broker, so each device spawned will get the clientID from the test dataset and do 2 things:
  
  
  
  • set the client ID in the MQTTLocust client, so the device can identify itself properly to the broker
  
  
  • pass the client ID to the task_set which needs it to set the topic
  
  

The first part works fine using super, but we've not been able to figure out a way for the Device class to tell the TaskSet which clientID it needs to run the publishing requests.
We've tried many things, including trying to __init__ the TaskSet with the clientID

Question: once the Device object figures out its clientID (through choosing randomly from a datafile) is there a way to pass the clientID to the TaskSet? We can't let the TaskSet pick it since it's too late (the ID is needed to initialize the MQTT client, and conceptually the Device would need to know its ID wouldn't it?)

Expected behavior

DeviceBehavior(TaskSet) class will happily accept the clientID and run the tasks with it

Actual behavior

When passed through the __init__ of DeviceBehavior:

locust().run(runner=self)
File "/Users/quan/Library/Python/2.7/lib/python/site-packages/locust/core.py", line 158, in run
task_set_instance = self.task_set(self)
TypeError: 'DeviceBehavior' object is not callable

Environment settings (for bug reports)

• OS: Linux
• Python version: 2.7
• Locust version: latest

Excerpts of the classes involved:

class MQTTLocust(Locust):
    def __init__(self, device_id=None, *args, **kwargs):
        super(Locust, self).__init__(*args, **kwargs)
        self.device_id = device_id
        self.client = MQTTClient(client_id=self.device_id, protocol=mqtt.MQTTv311)

class DeviceBehavior(TaskSet):
    def __init__(self, device_id):
        self.device_id = device_id
    @task
    def publish(self):
        topic = "devices/{}/messages/events/".format(self.device_id)
        ...
       self.client.publish(topic, payload=self.payload(), qos=0, timeout=PUBLISH_TIMEOUT)

class Device(MQTTLocust):
    def __init__(self):
        # Each device is identified by its ID 
        device_id = self.chooseRandomDevice()
        super(Device, self).__init__(device_id)
    task_set = AssetBehavior(would like to pass a variable here)

[devmonkey22]

TaskSet.__init__ sets the self.locust property to the locust, which in your case should be MQTTLocust. So within your DeviceBehavior TaskSet's __init__, you might be able to first call super().__init__ then get device_id from self.locust.device_id. I usually grab parameters like that in my TaskSet's on_start instead though. But maybe that's just me.

FYI - Locust's TaskSet.__init__ already takes self and parent as arguments, so you'll have to refactor your passing around of device_id out. See https://github.com/locustio/locust/blob/ff7fddd1ff770ead4fdcfb8af06213fdea5f63b9/locust/core.py#L301

[quanstorm]

@devmonkey22 you rocked! So simple.
Followed your advice I just used

    def on_start(self):
        self.tag = self.locust.device_id

so no need for any init like you mentioned.
This is considered answered.

[devmonkey22]

You should be able to Comment & Close below I think.

[quanstorm]

Thanks