better example

SvenskaSpel · Jul 11, 2024 · 9f9b920 · 9f9b920
1 parent 2aa91be
commit 9f9b920
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 6 deletions.
diff --git a/examples/mqtt_custom_client_ex.py b/examples/mqtt_custom_client_ex.py
@@ -4,7 +4,6 @@
 from locust import task, TaskSet
 from locust_plugins.users.mqtt import MqttUser
 from locust_plugins.users.mqtt import MqttClient
-from paho.mqtt.client import MQTTv311
 
 
 # extend the MqttClient class with your own custom implementation

diff --git a/examples/mqtt_ex.py b/examples/mqtt_ex.py
@@ -2,12 +2,10 @@
 import ssl
 import time
 
-from paho.mqtt.client import MQTTv311
 from locust import task, TaskSet
 from locust.user.wait_time import between
 from locust_plugins.users.mqtt import MqttUser
 
-
 tls_context = ssl.SSLContext(ssl.PROTOCOL_TLS)
 tls_context.load_verify_locations(os.environ["LOCUST_MQTT_CAFILE"])
 
@@ -24,6 +22,9 @@ class MyUser(MqttUser):
     # 10-100 messages per second.
     wait_time = between(0.01, 0.1)
 
+    # Uncomment below if you need to set MQTTv5
+    # protocol = paho.mqtt.client.MQTTv5
+
     @task
     class MyTasks(TaskSet):
         # Sleep for a while to allow the client time to connect.

diff --git a/locust_plugins/users/mqtt.py b/locust_plugins/users/mqtt.py
@@ -404,13 +404,12 @@ class MqttUser(User):
     client_id = None
     username = None
     password = None
+    protocol = mqtt.MQTTv311
 
     def __init__(self, environment: Environment):
         super().__init__(environment)
         self.client: MqttClient = self.client_cls(
-            environment=self.environment,
-            transport=self.transport,
-            client_id=self.client_id,
+            environment=self.environment, transport=self.transport, client_id=self.client_id, protocol=self.protocol
         )
 
         if self.tls_context: