k3s/rabbitmq-chart/test_send_mqtt_lb.py

#!/usr/bin/env python3
"""Send a single MQTT message to the Traefik LB IP to verify external MQTT connectivity.

Defaults to host 153.92.153.32, port 1883, topic mqtt/test, message "test message", user: user, pass: changeme.
Exits 0 on success, non-zero on failure.
"""
import sys
import time
import argparse

import paho.mqtt.client as mqtt


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--host', default='153.92.153.32')
    p.add_argument('--port', type=int, default=1883)
    p.add_argument('--topic', default='mqtt/test')
    p.add_argument('--message', default='hello from test_send_mqtt_lb')
    p.add_argument('--username', default='user')
    p.add_argument('--password', default='changeme')
    p.add_argument('--timeout', type=int, default=10)
    args = p.parse_args()

    connected = False
    published = False

    def on_connect(client, userdata, flags, rc):
        nonlocal connected
        connected = (rc == 0)
        print(f'on_connect rc={rc}, success={connected}')

    def on_publish(client, userdata, mid):
        nonlocal published
        published = True
        print(f'on_publish mid={mid}')

    client = mqtt.Client()
    client.username_pw_set(args.username, args.password)
    client.on_connect = on_connect
    client.on_publish = on_publish

    try:
        client.connect(args.host, args.port, keepalive=60)
    except Exception as e:
        print('connect failed:', e)
        return 2

    client.loop_start()

    # wait for connection
    waited = 0
    while not connected and waited < args.timeout:
        time.sleep(0.2)
        waited += 0.2

    if not connected:
        print('failed to connect within timeout')
        client.loop_stop()
        return 3

    try:
        result, mid = client.publish(args.topic, args.message, qos=1)
        print('publish result:', result, 'mid:', mid)
    except Exception as e:
        print('publish failed:', e)
        client.loop_stop()
        return 4

    # wait for on_publish
    waited = 0
    while not published and waited < args.timeout:
        time.sleep(0.2)
        waited += 0.2

    client.loop_stop()
    client.disconnect()

    if published:
        print('message sent')
        return 0
    else:
        print('message not confirmed published')
        return 5


if __name__ == '__main__':
    sys.exit(main())