Merge pull request #3216 from ionutab/add_mqtt

Adding mqtt capabilities to main locust repo

Author: cyberw

docs/api.rst

@@ -28,6 +28,11 @@ FastHttpUser class
     :members: wait_time, tasks, client, abstract, rest
     :noindex:
 
+MqttUser class
+==================
+.. autoclass:: locust.contrib.mqtt.MqttUser
+    :members: __init__, host, port, transport, ws_path, tls_context, client_cls, client_id, username, password, protocol 
+
 .. _socketio:
 
 SocketIOUser class

docs/testing-other-systems.rst

@@ -109,6 +109,22 @@ Performance/load testing AI services is a little different. While you could call
 
     OpenAIUser is experimental and may change without notice.
 
+MQTT
+====
+
+Locust uses to `paho-mqtt https://github.com/eclipse-paho/paho.mqtt.python`_ to provide Mqtt connection capabilities.
+
+.. literalinclude:: ../examples/mqtt/locustfile.py
+
+Alternatively, if you need more control over the Mqtt client you can use a custom implementation.
+
+.. literalinclude:: ../examples/mqtt/locustfile_custom_mqtt_client.py
+
+.. note::
+
+    MqttUser is experimental and may change without notice.
+
+
 Other examples
 ==============
 

examples/mqtt/README.md

@@ -0,0 +1,31 @@
+# MQTT Load testing with Locust
+
+## Prerequisites
+
+Have access to a running mqtt broker.
+
+```bash
+# start mosquitto locally
+# the configuration file is required to start 
+# WINDOWS
+docker run -d -p 1883:1883 --name mqtt-broker -v .\\examples\\mqtt\\mosquitto_config\\mosquitto.conf:/mosquitto/config/mosquitto.conf eclipse-mosquitto:latest
+# UNIX
+docker run -d -p 1883:1883 --name mqtt-broker -v ./examples/mqtt/mosquitto_config/mosquitto.conf:/mosquitto/config/mosquitto.conf eclipse-mosquitto:latest
+```
+
+Install Locust with MQTT support:
+
+```bash
+# Using pip
+pip install locust[mqtt]
+```
+
+## Usage
+
+```bash
+# Run simple example without web UI
+locust -f examples/mqtt/locustfile.py --headless
+# Run simple custom client example without web UI
+locust -f examples/mqtt/locustfile_custom_mqtt_client.py --headless
+```
+

examples/mqtt/locustfile.py

@@ -0,0 +1,34 @@
+from locust import task
+from locust.contrib.mqtt import MqttUser
+from locust.user.wait_time import between
+
+import time
+
+
+class MyUser(MqttUser):
+    host = "localhost"
+    port = 1883
+
+    # We could uncomment below to use the WebSockets transport
+    # transport = "websockets"
+
+    # ws_path = "/mqtt/custom/path"
+
+    # We'll probably want to throttle our publishing a bit: let's limit it to
+    # 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
+
+    # Sleep for a while to allow the client time to connect.
+    # This is probably not the most "correct" way to do this: a better method
+    # might be to add a gevent.event.Event to the MqttClient's on_connect
+    # callback and wait for that (with a timeout) here.
+    # However, this works well enough for the sake of an example.
+    def on_start(self):
+        time.sleep(5)
+
+    @task
+    def say_hello(self):
+        self.client.publish("hello/locust", b"hello world")

examples/mqtt/locustfile_custom_mqtt_client.py

@@ -0,0 +1,41 @@
+from locust import task
+from locust.contrib.mqtt import MqttClient, MqttUser
+from locust.user.wait_time import between
+
+import time
+
+
+# extend the MqttClient class with your own custom implementation
+class MyMqttClient(MqttClient):
+    # you can override the event name with your custom implementation
+    def _generate_event_name(self, event_type: str, qos: int, topic: str):
+        return f"mqtt:{event_type}:{qos}"
+
+
+class MyUser(MqttUser):
+    host = "localhost"
+    port = 1883
+
+    # We could uncomment below to use the WebSockets transport
+    # transport = "websockets"
+
+    # ws_path = "/mqtt/custom/path"
+
+    # We'll probably want to throttle our publishing a bit: let's limit it to
+    # 10-100 messages per second.
+    wait_time = between(0.01, 0.1)
+
+    # override the client_cls with your custom MqttClient implementation
+    client_cls = MyMqttClient
+
+    # Sleep for a while to allow the client time to connect.
+    # This is probably not the most "correct" way to do this: a better method
+    # might be to add a gevent.event.Event to the MqttClient's on_connect
+    # callback and wait for that (with a timeout) here.
+    # However, this works well enough for the sake of an example.
+    def on_start(self):
+        time.sleep(5)
+
+    @task
+    def say_hello(self):
+        self.client.publish("hello/locust", b"hello world locust custom client")

examples/mqtt/mosquitto_config/mosquitto.conf

@@ -0,0 +1,4 @@
+# mosquitto_config/mosquitto.conf
+
+listener 1883
+allow_anonymous true