zcbot/tests/test_broker_redis.py

"""RedisRunBroker 单测(fakeredis,无需真 redis 服务)。

覆盖跨进程语义的关键面:
- publish → 订阅 queue 收到(经 pubsub reader 路由)
- close 的 done key + publish 双通道:晚到订阅者(错过 publish)靠 key 立即收 done
- start 清 done key(新 run 的订阅者不被上一轮 done 打断)
- cancel key 的 request/is/clear 三态(跨"实例"可见:第二个 broker 实例读得到)
- emit 在 redis 挂掉时降级不抛(丢帧,run 不受影响)

用两个 broker 实例共享同一 FakeServer 模拟蓝绿两进程。
"""
import asyncio
import unittest
from uuid import uuid4

import fakeredis
import fakeredis.aioredis

from web.broker import LocalRunBroker, RedisRunBroker, _make_broker


def _make_pair(server):
    """一个共享 FakeServer 的 RedisRunBroker(sync + async 两个客户端注入)。"""
    sync_c = fakeredis.FakeRedis(server=server, decode_responses=True)
    async_c = fakeredis.aioredis.FakeRedis(server=server, decode_responses=True)
    return RedisRunBroker("redis://fake:6379/0", sync_client=sync_c, async_client=async_c)


async def _drain_until(q: asyncio.Queue, timeout: float = 2.0):
    return await asyncio.wait_for(q.get(), timeout=timeout)


class TestRedisBroker(unittest.TestCase):
    def _run(self, coro):
        return asyncio.run(coro)

    def test_emit_reaches_subscriber(self):
        async def main():
            server = fakeredis.FakeServer()
            b = _make_pair(server)
            b.bind_loop(asyncio.get_running_loop())
            tid = uuid4()
            q = b.subscribe(tid)
            await asyncio.sleep(0.1)  # 等 channel 挂载
            b.emit(tid, {"type": "text", "delta": "你好"})
            ev = await _drain_until(q)
            self.assertEqual(ev, {"type": "text", "delta": "你好"})
            b.unsubscribe(tid, q)
            await b.shutdown()
        self._run(main())

    def test_cross_instance_emit(self):
        """模拟蓝绿:A 实例上的 run emit,B 实例上的 SSE 订阅者收到。"""
        async def main():
            server = fakeredis.FakeServer()
            b_a = _make_pair(server)   # 跑 run 的实例
            b_b = _make_pair(server)   # 挂 SSE 的实例
            loop = asyncio.get_running_loop()
            b_a.bind_loop(loop)
            b_b.bind_loop(loop)
            tid = uuid4()
            q = b_b.subscribe(tid)
            await asyncio.sleep(0.1)
            b_a.emit(tid, {"type": "llm_start"})
            ev = await _drain_until(q)
            self.assertEqual(ev["type"], "llm_start")
            await b_a.shutdown()
            await b_b.shutdown()
        self._run(main())

    def test_late_subscriber_gets_done_via_key(self):
        """close 之后才来的订阅者:publish 早已错过,靠 done key 立即断流。"""
        async def main():
            server = fakeredis.FakeServer()
            b = _make_pair(server)
            b.bind_loop(asyncio.get_running_loop())
            tid = uuid4()
            b.close(tid)                     # SETEX done + publish(无人订阅,publish 丢)
            q = b.subscribe(tid)             # 晚到
            ev = await _drain_until(q)
            self.assertEqual(ev, {"type": "done"})
            self.assertTrue(b.is_done(tid))
            await b.shutdown()
        self._run(main())

    def test_start_clears_done(self):
        async def main():
            server = fakeredis.FakeServer()
            b = _make_pair(server)
            b.bind_loop(asyncio.get_running_loop())
            tid = uuid4()
            b.close(tid)
            self.assertTrue(b.is_done(tid))
            b.start(tid)                     # 新 run
            self.assertFalse(b.is_done(tid))
            q = b.subscribe(tid)
            await asyncio.sleep(0.1)
            self.assertTrue(q.empty())       # 不该被上一轮 done 打断
            b.emit(tid, {"type": "run_start"})
            ev = await _drain_until(q)
            self.assertEqual(ev["type"], "run_start")
            await b.shutdown()
        self._run(main())

    def test_cancel_cross_instance(self):
        """A 实例(SSE/cancel 端点)request_cancel,B 实例(跑 run 的)poll 到。"""
        async def main():
            server = fakeredis.FakeServer()
            b_a = _make_pair(server)
            b_b = _make_pair(server)
            b_a.bind_loop(asyncio.get_running_loop())
            b_b.bind_loop(asyncio.get_running_loop())
            tid = uuid4()
            self.assertFalse(b_b.is_cancelled(tid))
            b_a.request_cancel(tid)
            self.assertTrue(b_b.is_cancelled(tid))    # 跨实例可见
            b_b.clear_cancel(tid)
            self.assertFalse(b_a.is_cancelled(tid))
            await b_a.shutdown()
            await b_b.shutdown()
        self._run(main())

    def test_emit_degrades_when_redis_down(self):
        """redis 操作抛异常 → emit/is_cancelled 降级不抛(log 节流另测无必要)。"""
        class Boom:
            def __getattr__(self, name):
                def _raise(*a, **k):
                    raise ConnectionError("redis down")
                return _raise

        b = RedisRunBroker("redis://fake:6379/0", sync_client=Boom(),
                           async_client=fakeredis.aioredis.FakeRedis(decode_responses=True))
        tid = uuid4()
        b.emit(tid, {"type": "text", "delta": "x"})       # 不抛
        self.assertFalse(b.is_cancelled(tid))              # 降级 False
        b.request_cancel(tid)                              # 不抛
        b.close(tid)                                       # 不抛

    def test_bind_loop_fail_fast_when_unreachable(self):
        class Boom:
            def ping(self):
                raise ConnectionError("refused")

        b = RedisRunBroker("redis://fake:6379/0", sync_client=Boom(),
                           async_client=fakeredis.aioredis.FakeRedis(decode_responses=True))

        async def main():
            with self.assertRaises(RuntimeError):
                b.bind_loop(asyncio.get_running_loop())
        self._run(main())


class TestFactory(unittest.TestCase):
    def test_default_is_local(self):
        import os
        old = os.environ.pop("ZCBOT_REDIS_URL", None)
        try:
            self.assertIsInstance(_make_broker(), LocalRunBroker)
        finally:
            if old is not None:
                os.environ["ZCBOT_REDIS_URL"] = old

    def test_env_selects_redis(self):
        import os
        old = os.environ.get("ZCBOT_REDIS_URL")
        os.environ["ZCBOT_REDIS_URL"] = "redis://127.0.0.1:6379/0"
        try:
            self.assertIsInstance(_make_broker(), RedisRunBroker)
        finally:
            if old is None:
                del os.environ["ZCBOT_REDIS_URL"]
            else:
                os.environ["ZCBOT_REDIS_URL"] = old


if __name__ == "__main__":
    unittest.main()