"""
Worker 框架流式回复单测（不依赖真实 zenoh 会话）

直接测 runtime.stream_replies_to_query（fake query 仅需 key_expr 与 reply）:
  1. 异步生成器：多块顺序 reply，payload 为 orjson JSON
  2. 同步生成器：同样逐块 reply
  3. 生成器中途异常：{"error": ...} 收尾块后停流
  4. reply 失败（查询已终止）：立即停流且不抛异常

运行: uv run python test_stream_worker.py
"""
import asyncio
import sys
from pathlib import Path

import orjson

REPO_ROOT = Path(__file__).resolve().parent
# 优先使用本地框架 checkout（即使 chongming_worker 已通过 uv 安装也以本地为准）
sys.path.insert(0, str(REPO_ROOT / "utils" / "worker" / "src"))

from chongming_worker.runtime import stream_replies_to_query  # noqa: E402

failures = []


def check(name, cond, detail=""):
    status = "PASS" if cond else "FAIL"
    print(f"[{status}] {name}" + (f" — {detail}" if detail else ""))
    if not cond:
        failures.append(name)


class FakeQuery:
    """仅实现 stream_replies_to_query 用到的两个接口（key_expr / reply）"""

    def __init__(self, key_expr="worker/example/worker-0/example.stream_chat"):
        self.key_expr = key_expr
        self.replies = []
        self.fail_reply = False

    def reply(self, key_expr, payload):
        if self.fail_reply:
            raise RuntimeError("query terminated")
        self.replies.append((key_expr, payload))


def test_async_generator():
    q = FakeQuery()

    async def gen():
        yield {"token": "a"}
        yield {"token": "b"}

    asyncio.run(stream_replies_to_query(q, gen()))
    check("异步生成器逐块回复", len(q.replies) == 2)
    check("回复 key_expr 正确", all(k == q.key_expr for k, _ in q.replies))
    check(
        "payload 为 JSON dict",
        [orjson.loads(p) for _, p in q.replies] == [{"token": "a"}, {"token": "b"}],
    )


def test_sync_generator():
    q = FakeQuery()

    def gen():
        yield {"n": 1}
        yield {"n": 2}

    asyncio.run(stream_replies_to_query(q, gen()))
    check(
        "同步生成器逐块回复",
        [orjson.loads(p) for _, p in q.replies] == [{"n": 1}, {"n": 2}],
    )


def test_mid_stream_exception():
    q = FakeQuery()

    async def gen():
        yield {"token": "a"}
        raise ValueError("boom")

    asyncio.run(stream_replies_to_query(q, gen()))
    check("中途异常收尾 error 块", len(q.replies) == 2)
    check("error 块内容", orjson.loads(q.replies[1][1]) == {"error": "boom"})


def test_reply_failure_stops():
    q = FakeQuery()
    q.fail_reply = True
    yielded = []

    async def gen():
        yielded.append("yield-1")
        yield {"token": "a"}
        yielded.append("yield-2")  # 首个块 reply 失败后不应被执行到
        yield {"token": "b"}

    # 不应抛异常，且首个块 reply 失败后立即停流
    asyncio.run(stream_replies_to_query(q, gen()))
    check("reply 失败不抛异常", True)
    check("reply 失败后立即停流", yielded == ["yield-1"])


def run():
    test_async_generator()
    test_sync_generator()
    test_mid_stream_exception()
    test_reply_failure_stops()

    print(f"\n结果: {len(failures)} 个失败")
    return 1 if failures else 0


if __name__ == "__main__":
    sys.exit(run())