tdEngine_mqtt_mock/app/mqtt_mock.py

import argparse
import json
import logging
import os
import time
from datetime import datetime
from urllib.parse import urlparse

import paho.mqtt.client as mqtt

from config import build_mock_dependencies
from model import DrillingRealtimeData


logger = logging.getLogger(__name__)


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 ensure_dir(path):
    if not path:
        return
    dir_path = os.path.dirname(path)
    if dir_path:
        os.makedirs(dir_path, exist_ok=True)


def build_server_payload(equipment_code):
    return DrillingRealtimeData.empty(wellid=equipment_code).to_payload(equipment_code)


def parse_payload(raw_payload):
    try:
        return json.loads(raw_payload)
    except Exception:
        return raw_payload


def build_ack_payload(topic, raw_payload):
    payload = {
        "ack": True,
        "topic": topic,
        "ts": datetime.utcnow().isoformat(timespec="seconds") + "Z",
    }
    try:
        obj = json.loads(raw_payload)
        if isinstance(obj, dict):
            meta = obj.get("meta") if isinstance(obj.get("meta"), dict) else {}
            if "equipment_code" in meta:
                payload["equipment_code"] = meta["equipment_code"]
            if "equipment_sn" in meta:
                payload["equipment_sn"] = meta["equipment_sn"]
            if "id" in obj:
                payload["id"] = obj["id"]
            payload["src"] = obj
    except Exception:
        payload["raw"] = raw_payload
    return payload


def write_message(path, topic, raw_payload):
    ensure_dir(path)
    record = {
        "ts": datetime.utcnow().isoformat(timespec="seconds") + "Z",
        "topic": topic,
        "payload": parse_payload(raw_payload),
    }
    payload_obj = record["payload"]
    if isinstance(payload_obj, dict):
        data = payload_obj.get("data")
        if isinstance(data, dict) and isinstance(data.get("ts"), (int, float)):
            record["ts_iso"] = datetime.utcfromtimestamp(data["ts"] / 1000.0).isoformat(timespec="seconds") + "Z"
    with open(path, "a", encoding="utf-8") as f:
        f.write(json.dumps(record, ensure_ascii=True))
        f.write("\n")


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

    logger.info("MQTT mock config broker=%s://%s:%s client_id=%s mode=%s", scheme, host, port, mqtt_config.mock_client_id, args.mode)
    logger.info("Topics ingest=%s forward=%s ack=%s", mqtt_config.pub_topic, mqtt_config.sub_topic, mqtt_config.ack_topic)
    logger.info("Data file=%s device_code=%s tdengine_enabled=%s", deps.data_file, tms_config.device_code, tdengine_writer.enabled)
    if tdengine_writer.enabled:
        logger.info("TDengine host=%s database=%s stable=%s pool_size=%s", tdengine_config.base_url, tdengine_config.database, tdengine_config.stable, tdengine_config.pool_size)

    client = mqtt.Client(client_id=mqtt_config.mock_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_connect(c, userdata, flags, rc):
        if rc == 0:
            logger.info("Connected")
        else:
            logger.error("Connect failed rc=%s", rc)
            return
        if args.mode in ("listen", "both") and mqtt_config.pub_topic:
            c.subscribe(mqtt_config.pub_topic)
            logger.info("Subscribed %s", mqtt_config.pub_topic)

    def on_disconnect(c, userdata, rc):
        logger.info("Disconnected callback rc=%s", rc)

    def on_message(c, userdata, msg):
        payload = msg.payload.decode("utf-8", errors="replace")
        logger.info("RX topic=%s bytes=%s", msg.topic, len(msg.payload))
        if msg.topic != mqtt_config.pub_topic:
            return
        if deps.data_file:
            write_message(deps.data_file, msg.topic, payload)
            logger.info("Wrote file %s", deps.data_file)
        if tdengine_writer.enabled:
            try:
                tdengine_writer.write_payload(parse_payload(payload))
                logger.info("Wrote TDengine")
            except Exception:
                logger.exception("Write TDengine failed")
        if mqtt_config.sub_topic:
            c.publish(mqtt_config.sub_topic, payload)
            logger.info("Forwarded %s", mqtt_config.sub_topic)
        if mqtt_config.ack_topic:
            ack_payload = build_ack_payload(msg.topic, payload)
            c.publish(mqtt_config.ack_topic, json.dumps(ack_payload, ensure_ascii=True))
            logger.info("TX ack %s", mqtt_config.ack_topic)

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

    try:
        if args.mode in ("publish", "both"):
            if not mqtt_config.pub_topic:
                logger.warning("pub-topic is empty; nothing to publish")
            else:
                seq = 0
                while True:
                    seq += 1
                    payload = build_server_payload(tms_config.device_code)
                    client.publish(mqtt_config.pub_topic, json.dumps(payload, ensure_ascii=True))
                    logger.info("TX %s #%s", mqtt_config.pub_topic, seq)
                    if args.count and seq >= args.count:
                        break
                    time.sleep(args.interval)
        else:
            while True:
                time.sleep(1)
    except KeyboardInterrupt:
        logger.info("Mock interrupted")
    finally:
        client.loop_stop()
        client.disconnect()
        logger.info("Mock stopped")


def add_arguments(parser):
    parser.add_argument("--config", default="config.yaml", help="Path to config yaml")
    parser.add_argument("--mode", choices=["publish", "listen", "both"], default="listen")
    parser.add_argument("--interval", type=float, default=2.0, help="Publish interval (seconds)")
    parser.add_argument("--count", type=int, default=0, help="Publish count (0 = forever)")
    parser.add_argument("--data-file", default="", help="Override data-file in config")


def main(argv=None):
    parser = argparse.ArgumentParser(description="MQTT mock service")
    add_arguments(parser)
    args = parser.parse_args(argv)
    deps = build_mock_dependencies(args.config, data_file_override=args.data_file)
    run_mock_service(args, deps)


if __name__ == "__main__":
    main()