chongming-core/test/test_mqtt_e2e.py

"""
MQTT 端到端测试脚本

测试流程:
  1. 启动内嵌 amqtt Broker
  2. 启动 subscriber（监听响应）
  3. 启动 publisher（发送 MQTT 请求 → chongming/example/add）
  4. 验证 subscriber 收到正确的响应

注意: 此脚本独立运行，不依赖 chongming 服务端。
      使用时需确保 chongming 服务已通过 `uv run serve` 启动。
"""

import asyncio
import orjson
import sys
import logging
from amqtt.client import MQTTClient
from amqtt.mqtt.constants import QOS_1

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(name)s] %(message)s",
)

PREFIX = "chongming"
WORKER = "example"
OPERATION = "add"
REQUEST_TOPIC = f"{PREFIX}/{WORKER}/{OPERATION}"
RESPONSE_TOPIC = f"{PREFIX}/{WORKER}/{OPERATION}/response"
BROKER_URL = "mqtt://localhost:1883/"

# 测试用例（信封协议: meta.method 控制分派，data 携带业务参数）
TEST_CASES = [
    {"a": 3, "b": 5, "expected": 8},
    {"a": -1, "b": 10, "expected": 9},
    {"a": 3.14, "b": 2.86, "expected": 6.0},
]

PASS = 0
FAIL = 0


async def run_test():
    global PASS, FAIL
    client = MQTTClient()
    await client.connect(BROKER_URL)
    await client.subscribe([(RESPONSE_TOPIC, QOS_1)])
    print(f"📡 已订阅: {RESPONSE_TOPIC}\n")

    for i, case in enumerate(TEST_CASES):
        a, b, expected = case["a"], case["b"], case["expected"]
        payload = {"meta": {"method": "MQTT"}, "data": {"a": a, "b": b}}
        payload_bytes = orjson.dumps(payload)

        await client.publish(REQUEST_TOPIC, payload_bytes, QOS_1)
        print(f"📤 [{i+1}/{len(TEST_CASES)}] 发送: {REQUEST_TOPIC} → {payload}")

        # 等待响应（超时 5 秒）
        try:
            message = await asyncio.wait_for(client.deliver_message(), timeout=5.0)
            data = orjson.loads(message.data)  # type: ignore
            actual = data.get("result")
            status = "✅ PASS" if abs(actual - expected) < 0.01 else "❌ FAIL"
            if status == "✅ PASS":
                PASS += 1
            else:
                FAIL += 1
            print(f"📥 [{i+1}/{len(TEST_CASES)}] 响应: {orjson.dumps(data).decode()} → {status} (expect={expected}, got={actual})")
        except asyncio.TimeoutError:
            FAIL += 1
            print(f"📥 [{i+1}/{len(TEST_CASES)}] ❌ FAIL: 超时无响应")

    await client.unsubscribe([RESPONSE_TOPIC])
    await client.disconnect()

    print(f"\n{'='*50}")
    print(f"测试结果: ✅ {PASS} 通过  ❌ {FAIL} 失败")
    print(f"{'='*50}")
    return FAIL == 0


async def main():
    # 先检查 Broker 是否可用
    test_client = MQTTClient()
    try:
        await asyncio.wait_for(test_client.connect(BROKER_URL), timeout=3.0)
        await test_client.disconnect()
    except Exception:
        print("❌ 无法连接到 MQTT Broker (localhost:1883)")
        print("   请先启动 chongming 服务: uv run serve")
        sys.exit(1)

    success = await run_test()
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n⏹️  测试被中断")