chongming-core/test/test_mqtt_custom_topics.py

"""
MQTT Bridge 自定义订阅主题单元测试

验证:
  1. _compile_topic 对 + / # 通配符的匹配与捕获
  2. 自定义主题路由: 主题匹配 → worker/subject/response_topic 解析 → worker_manager.call 调用
  3. 内置 chongming/... 主题兜底逻辑不变
"""
import asyncio
import orjson
import sys
from pathlib import Path
from unittest.mock import MagicMock

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "gateway" / "src"))

from chongming.app.core.mqtt_bridge import (
    MQTTBridge,
    MQTT_TOPIC_PREFIX,
    _compile_topic,
)


def make_msg(topic, payload):
    m = MagicMock()
    m.topic = topic
    m.data = orjson.dumps(payload)
    return m


class FakeWM:
    def __init__(self):
        self.calls = []
        self.result = {"result": "ok"}

    async def call(self, worker_type, subject, payload, timeout=30, cacheable=False,
                   instance_id=None):
        self.calls.append(
            {
                "worker_type": worker_type,
                "subject": subject,
                "payload": payload,
                "timeout": timeout,
                "cacheable": cacheable,
                "instance_id": instance_id,
            }
        )
        return self.result


def make_bridge(wm, custom_topics):
    """构造一个不真正连接 Broker 的 Bridge，并捕获发布消息"""
    bridge = MQTTBridge(worker_manager=wm, custom_topics=custom_topics)
    bridge._client = MagicMock()
    bridge._running = True
    published = []

    async def fake_publish(topic, data, qos=None):
        published.append((topic, data.decode("utf-8")))

    bridge._client.publish = fake_publish
    return bridge, published


def test_compile_topic():
    # 单级通配符
    p = _compile_topic("devices/+/telemetry")
    m = p.match("devices/dev-01/telemetry")
    assert m and m.groups() == ("dev-01",)
    assert not p.match("devices/a/b/telemetry")

    # 多级通配符
    p = _compile_topic("data/#")
    m = p.match("data/a/b/c")
    assert m and m.groups() == ("a/b/c",)

    # 无通配符
    p = _compile_topic("fixed/topic")
    assert p.match("fixed/topic")
    assert not p.match("fixed/topic/x")

    # 前缀必须完整匹配
    p = _compile_topic("devices/+/telemetry")
    assert not p.match("xdevices/a/telemetry")
    print("✅ _compile_topic 通配符编译/匹配通过")


async def test_custom_routing():
    wm = FakeWM()
    bridge, published = make_bridge(
        wm,
        [
            {"topic": "devices/+/telemetry", "worker": "device"},
            {
                "topic": "foo/#",
                "worker": "foo",
                "subject": "foo.handle.{1}",
                "response_topic": "out/{1}/resp",
            },
        ],
    )

    # 1. 默认 subject = <worker>.<最后一段>；响应 = <实际主题>/response
    await bridge._handle(make_msg("devices/dev-01/telemetry", {"meta": {}, "data": {"v": 1}}))
    assert wm.calls[-1]["worker_type"] == "device"
    assert wm.calls[-1]["subject"] == "device.telemetry"
    assert published[-1][0] == "devices/dev-01/telemetry/response"

    # 2. 带占位符 + 自定义响应主题
    await bridge._handle(make_msg("foo/a/b", {"meta": {}, "data": {"v": 2}}))
    assert wm.calls[-1]["worker_type"] == "foo"
    assert wm.calls[-1]["subject"] == "foo.handle.a/b"
    assert published[-1][0] == "out/a/b/resp"

    # 3. 内置 chongming 兜底
    await bridge._handle(
        make_msg(f"{MQTT_TOPIC_PREFIX}/example/add", {"meta": {}, "data": {"a": 1, "b": 2}})
    )
    assert wm.calls[-1]["worker_type"] == "example"
    assert wm.calls[-1]["subject"] == "example.add"
    assert published[-1][0] == f"{MQTT_TOPIC_PREFIX}/example/add/response"
    print("✅ 自定义主题路由 + 内置兜底通过")


async def test_envelope_meta_fields():
    wm = FakeWM()
    bridge, published = make_bridge(
        wm, [{"topic": "devices/+/telemetry", "worker": "device"}]
    )
    await bridge._handle(
        make_msg(
            "devices/d1/telemetry",
            {"meta": {"timeout": 5, "cacheable": True, "method": "POST"}, "data": {}},
        )
    )
    assert wm.calls[-1]["timeout"] == 5.0
    assert wm.calls[-1]["cacheable"] is True
    payload = wm.calls[-1]["payload"]
    assert payload["meta"]["method"] == "POST"
    assert payload["data"] == {}
    print("✅ 信封 meta 字段 (method/timeout/cacheable) 透传通过")


async def test_flat_payload_routed():
    """非信封（扁平）载荷兼容：整包视为 data、meta 置空照常路由（与 worker 运行时容错一致）"""
    wm = FakeWM()
    bridge, published = make_bridge(
        wm, [{"topic": "devices/+/telemetry", "worker": "device"}]
    )
    await bridge._handle(make_msg("devices/d1/telemetry", {"_method": "POST", "v": 1}))
    assert wm.calls, "扁平载荷不应被丢弃，应照常路由到 worker"
    assert wm.calls[-1]["worker_type"] == "device"
    assert wm.calls[-1]["payload"] == {"meta": {}, "data": {"_method": "POST", "v": 1}}
    assert published[-1][0] == "devices/d1/telemetry/response"
    print("✅ 非信封（扁平）载荷按原始报文路由通过")


async def test_target_instance():
    """实例定向：载荷顶层 target_instance 作为控制字段摘除并透传给 worker_manager.call"""
    wm = FakeWM()
    bridge, published = make_bridge(
        wm, [{"topic": "devices/+/telemetry", "worker": "device"}]
    )
    # 1. 带 target_instance：透传实例名，且控制字段不进业务 payload
    await bridge._handle(make_msg(
        "devices/d1/telemetry",
        {"meta": {}, "data": {"v": 1}, "target_instance": "edge-01"},
    ))
    assert wm.calls[-1]["instance_id"] == "edge-01", wm.calls[-1]
    assert "target_instance" not in wm.calls[-1]["payload"], wm.calls[-1]["payload"]
    assert wm.calls[-1]["payload"]["data"] == {"v": 1}

    # 2. 不带该字段：instance_id 为 None（跨实例随机均衡，与改动前一致）
    await bridge._handle(make_msg("devices/d1/telemetry", {"meta": {}, "data": {"v": 1}}))
    assert wm.calls[-1]["instance_id"] is None, wm.calls[-1]

    # 3. 空白值视为未指定
    await bridge._handle(make_msg(
        "devices/d1/telemetry",
        {"meta": {}, "data": {"v": 1}, "target_instance": "   "},
    ))
    assert wm.calls[-1]["instance_id"] is None, wm.calls[-1]

    # 4. 内置 chongming/... 主题同样支持
    await bridge._handle(make_msg(
        f"{MQTT_TOPIC_PREFIX}/example/add",
        {"meta": {}, "data": {"a": 1}, "target_instance": "edge-02"},
    ))
    assert wm.calls[-1]["instance_id"] == "edge-02", wm.calls[-1]
    print("✅ 实例定向（target_instance 控制字段）透传通过")


async def test_non_json_dropped():
    wm = FakeWM()
    bridge, published = make_bridge(
        wm, [{"topic": "devices/+/telemetry", "worker": "device"}]
    )
    bad = MagicMock()
    bad.topic = "devices/d1/telemetry"
    bad.data = b"not-json"
    await bridge._handle(bad)
    assert wm.calls == []
    assert published == []
    print("✅ 非 JSON 负载被丢弃通过")


async def test_board_ping_topic():
    """aibox_sophon_adapter 插件：设备发布到 /board_ping → 路由到 aibox_sophon_adapter.board_ping"""
    wm = FakeWM()
    bridge, published = make_bridge(
        wm,
        [
            {
                "topic": "/board_ping",
                "worker": "aibox_sophon_adapter",
                "subject": "aibox_sophon_adapter.board_ping",
            }
        ],
    )
    await bridge._handle(make_msg("/board_ping", {"meta": {}, "data": {"BoardId": "B001", "Status": 1}}))
    assert wm.calls[-1]["worker_type"] == "aibox_sophon_adapter"
    assert wm.calls[-1]["subject"] == "aibox_sophon_adapter.board_ping"
    assert published[-1][0] == "/board_ping/response"
    # 回归：盒子心跳为原始报文（无信封），不得被丢弃，应整包作为 data 路由
    wm.calls.clear(); published.clear()
    await bridge._handle(make_msg("/board_ping", {"BoardId": "B001", "Status": 1}))
    assert wm.calls and wm.calls[-1]["payload"] == {
        "meta": {}, "data": {"BoardId": "B001", "Status": 1}
    }, wm.calls
    assert published[-1][0] == "/board_ping/response"
    print("✅ /board_ping 自定义主题路由通过（含原始报文兼容）")


def test_plugin_config_driven_collection():
    """插件在其 config/workers/<type>/config.toml 中声明 [[workers.mqtt_topics]]，
    网关 _discover_plugins 自动收集，并默认 worker 为该插件类型。"""
    from chongming.app.core.lifespan_setup import _discover_plugins

    plugins = _discover_plugins()
    by_type = {p["type"]: p for p in plugins}
    assert "aibox_sophon_adapter" in by_type, "未发现 aibox_sophon_adapter 插件"
    topics = by_type["aibox_sophon_adapter"]["mqtt_topics"]
    assert any(
        t["topic"] == "box/aibox_sophon/board_ping"
        and t["worker"] == "aibox_sophon_adapter"
        and t["subject"] == "aibox_sophon_adapter.board_ping"
        for t in topics
    ), topics
    # 不应泄漏到插件运行时配置
    assert "mqtt_topics" not in by_type["aibox_sophon_adapter"]["config"]
    print("✅ 插件级 [[workers.mqtt_topics]] 收集通过")


async def main():
    test_compile_topic()
    await test_custom_routing()
    await test_envelope_meta_fields()
    await test_flat_payload_routed()
    await test_target_instance()
    await test_non_json_dropped()
    await test_board_ping_topic()
    test_plugin_config_driven_collection()
    print("\n🎉 全部自定义订阅主题测试通过")


if __name__ == "__main__":
    asyncio.run(main())