"""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()