tdEngine_mqtt_mock/mqtt_sender.py

import argparse
import json
import random
import time
from datetime import datetime
from urllib.parse import urlparse

import paho.mqtt.client as mqtt

from config.dependencies import build_sender_dependencies


DATA_KEYS = [
    "ts",
    "wellid",
    "stknum",
    "recid",
    "seqid",
    "actual_date",
    "actual_time",
    "actcod",
    "deptbitm",
    "deptbitv",
    "deptmeas",
    "deptvert",
    "blkpos",
    "ropa",
    "hkla",
    "hklx",
    "woba",
    "wobx",
    "torqa",
    "torqx",
    "rpma",
    "sppa",
    "chkp",
    "spm1",
    "spm2",
    "spm3",
    "tvolact",
    "tvolcact",
    "mfop",
    "mfoa",
    "mfia",
    "mdoa",
    "mdia",
    "mtoa",
    "mtia",
    "mcoa",
    "mcia",
    "stkc",
    "lagstks",
    "deptretm",
    "gasa",
    "space1",
    "space2",
    "space3",
    "space4",
    "space5",
]


def parse_broker(broker):
    if not broker:
        raise ValueError("broker is required")
    if "://" not in broker:
        broker = "tcp://" + broker
    parsed = urlparse(broker)
    host = parsed.hostname or "localhost"
    port = parsed.port or 1883
    scheme = (parsed.scheme or "tcp").lower()
    return scheme, host, port


def rand_int(a, b):
    return random.randint(a, b)


def build_random_payload(equipment_code):
    data = {key: 0 for key in DATA_KEYS}
    data["ts"] = int(time.time() * 1000)
    data["wellid"] = random.choice(["", f"WELL-{rand_int(1, 9999):04d}"])
    data["stknum"] = rand_int(0, 500)
    data["recid"] = rand_int(0, 100000)
    data["seqid"] = rand_int(0, 100000)
    data["actual_date"] = int(datetime.utcnow().strftime("%Y%m%d"))
    data["actual_time"] = int(datetime.utcnow().strftime("%H%M%S"))
    data["actcod"] = rand_int(0, 9)
    data["deptbitm"] = rand_int(0, 5000)
    data["deptbitv"] = rand_int(0, 5000)
    data["deptmeas"] = rand_int(0, 5000)
    data["deptvert"] = rand_int(0, 5000)
    data["blkpos"] = rand_int(0, 100)
    data["ropa"] = rand_int(0, 200)
    data["hkla"] = rand_int(0, 500)
    data["hklx"] = rand_int(0, 500)
    data["woba"] = rand_int(0, 200)
    data["wobx"] = rand_int(0, 200)
    data["torqa"] = rand_int(0, 200)
    data["torqx"] = rand_int(0, 200)
    data["rpma"] = rand_int(0, 300)
    data["sppa"] = rand_int(0, 5000)
    data["chkp"] = rand_int(0, 5000)
    data["spm1"] = rand_int(0, 200)
    data["spm2"] = rand_int(0, 200)
    data["spm3"] = rand_int(0, 200)
    data["tvolact"] = rand_int(0, 20000)
    data["tvolcact"] = rand_int(0, 20000)
    data["mfop"] = rand_int(0, 1000)
    data["mfoa"] = rand_int(0, 1000)
    data["mfia"] = rand_int(0, 1000)
    data["mdoa"] = rand_int(0, 1000)
    data["mdia"] = rand_int(0, 1000)
    data["mtoa"] = rand_int(0, 1000)
    data["mtia"] = rand_int(0, 1000)
    data["mcoa"] = rand_int(0, 1000)
    data["mcia"] = rand_int(0, 1000)
    data["stkc"] = rand_int(0, 200)
    data["lagstks"] = rand_int(0, 200)
    data["deptretm"] = rand_int(0, 5000)
    data["gasa"] = rand_int(0, 100)
    data["space1"] = rand_int(0, 10)
    data["space2"] = rand_int(0, 10)
    data["space3"] = rand_int(0, 10)
    data["space4"] = rand_int(0, 10)
    data["space5"] = rand_int(0, 10)
    return {
        "meta": {
            "equipment_code": equipment_code,
            "equipment_sn": equipment_code,
        },
        "data": data,
    }


def run_sender(args, deps):
    mqtt_config = deps.config.mqtt
    tms_config = deps.config.tms
    scheme, host, port = parse_broker(mqtt_config.broker)

    print("MQTT sender config:")
    print(f"  broker: {scheme}://{host}:{port}")
    print(f"  client-id: {mqtt_config.sender_client_id}")
    print(f"  pub-topic: {mqtt_config.pub_topic}")
    print(f"  interval: {args.interval}s")

    client = mqtt.Client(client_id=mqtt_config.sender_client_id, clean_session=True)
    if mqtt_config.username is not None:
        client.username_pw_set(mqtt_config.username, mqtt_config.password)

    if scheme in ("ssl", "tls", "mqtts"):
        client.tls_set()

    def on_disconnect(c, userdata, rc):
        print(f"Disconnected callback rc={rc}")

    client.on_disconnect = on_disconnect
    client.connect(host, port, keepalive=tms_config.keepalive)
    client.loop_start()

    try:
        if not mqtt_config.pub_topic:
            print("pub-topic is empty; nothing to publish")
            return
        seq = 0
        while True:
            seq += 1
            payload = build_random_payload(tms_config.device_code)
            client.publish(mqtt_config.pub_topic, json.dumps(payload, ensure_ascii=True))
            print(f"TX {mqtt_config.pub_topic}: {payload}")
            if args.count and seq >= args.count:
                break
            time.sleep(args.interval)
    except KeyboardInterrupt:
        pass
    finally:
        client.loop_stop()
        client.disconnect()
        print("Disconnected")


def main():
    ap = argparse.ArgumentParser(description="MQTT random data sender")
    ap.add_argument("--config", default="config.yaml", help="Path to config yaml")
    ap.add_argument("--interval", type=float, default=3.0, help="Publish interval (seconds)")
    ap.add_argument("--count", type=int, default=0, help="Publish count (0 = forever)")
    args = ap.parse_args()
    deps = build_sender_dependencies(args.config)
    run_sender(args, deps)


if __name__ == "__main__":
    main()