from __future__ import annotations import asyncio from unittest.mock import AsyncMock, MagicMock, patch import pytest from yuxi.channels.infra.broadcast import EventBroadcaster from yuxi.channels.services.context import ChatAbortEntry, ChatRunBuffer, GatewayRequestContext from yuxi.channels.manager import ChannelManager from yuxi.channels.registry import ChannelRegistry from yuxi.channels.router import MessageRouter class TestChatAbortEntryV2: def test_abort_sets_event(self): entry = ChatAbortEntry() assert not entry.is_aborted() entry.abort() assert entry.is_aborted() def test_abort_with_none_task_sets_event(self): entry = ChatAbortEntry(task=None) entry.abort() assert entry.is_aborted() def test_abort_cancels_task(self): async def dummy(): await asyncio.sleep(10) loop = asyncio.new_event_loop() task = loop.create_task(dummy()) entry = ChatAbortEntry(task=task) entry.abort() loop.run_until_complete(asyncio.sleep(0)) assert task.cancelled() assert entry.is_aborted() loop.close() def test_abort_already_done_task_does_not_error(self): async def noop(): pass loop = asyncio.new_event_loop() task = loop.create_task(noop()) loop.run_until_complete(task) entry = ChatAbortEntry(task=task) entry.abort() assert entry.is_aborted() loop.close() class TestChatRunBufferV2: def test_get_chunks(self): buffer = ChatRunBuffer(run_id="test-1") buffer.append_chunk("a") buffer.append_chunk("b") assert buffer.get_chunks() == ["a", "b"] def test_get_chunks_returns_copy(self): buffer = ChatRunBuffer(run_id="test-2") buffer.append_chunk("a") chunks = buffer.get_chunks() chunks.append("b") assert buffer.get_chunks() == ["a"] def test_chunk_count(self): buffer = ChatRunBuffer(run_id="test-3") assert buffer.chunk_count() == 0 buffer.append_chunk("a") assert buffer.chunk_count() == 1 buffer.append_chunk("b") assert buffer.chunk_count() == 2 def test_append_and_get_full_text(self): buffer = ChatRunBuffer(run_id="test-4") buffer.append_chunk("Hello ") buffer.append_chunk("World") assert buffer.get_full_text() == "Hello World" def test_mark_finished(self): buffer = ChatRunBuffer(run_id="test-5") assert not buffer.is_finished() buffer.mark_finished() assert buffer.is_finished() class TestGatewayRequestContextV2: def test_runtime_snapshot(self): ctx = GatewayRequestContext(runtime_snapshot=None) assert ctx.runtime_snapshot is None def test_mark_channel_logged_out_callable(self): called = [] async def mock_logout(channel_id: str, account_id: str): called.append((channel_id, account_id)) ctx = GatewayRequestContext(mark_channel_logged_out=mock_logout) assert ctx.mark_channel_logged_out is mock_logout def test_set_runtime_config(self): ctx = GatewayRequestContext() ctx.set_runtime_config("key1", "val1") ctx.set_runtime_config("key2", "val2") assert ctx.get_runtime_config("key1") == "val1" assert ctx.get_runtime_config("key2") == "val2" def test_add_chat_run_with_buffer(self): ctx = GatewayRequestContext() entry = ChatAbortEntry() buffer = ChatRunBuffer(run_id="r1") ctx.add_chat_run("r1", entry, buffer) assert ctx.chat_abort_controllers["r1"] is entry assert ctx.chat_run_buffers["r1"] is buffer def test_add_chat_run_without_buffer(self): ctx = GatewayRequestContext() entry = ChatAbortEntry() ctx.add_chat_run("r1", entry) assert ctx.chat_abort_controllers["r1"] is entry assert "r1" not in ctx.chat_run_buffers def test_remove_chat_run_clears_both(self): ctx = GatewayRequestContext() ctx.chat_abort_controllers["r1"] = ChatAbortEntry() ctx.chat_run_buffers["r1"] = ChatRunBuffer(run_id="r1") ctx.remove_chat_run("r1") assert "r1" not in ctx.chat_abort_controllers assert "r1" not in ctx.chat_run_buffers def test_remove_chat_run_not_exist_no_error(self): ctx = GatewayRequestContext() ctx.remove_chat_run("nonexistent") def test_abort_chat_run_success(self): ctx = GatewayRequestContext() entry = ChatAbortEntry() ctx.chat_abort_controllers["r1"] = entry assert ctx.abort_chat_run("r1") is True def test_abort_chat_run_not_found(self): ctx = GatewayRequestContext() assert ctx.abort_chat_run("nonexistent") is False def test_get_chat_buffer_found(self): ctx = GatewayRequestContext() buffer = ChatRunBuffer(run_id="r1") ctx.chat_run_buffers["r1"] = buffer assert ctx.get_chat_buffer("r1") is buffer def test_get_chat_buffer_not_found(self): ctx = GatewayRequestContext() assert ctx.get_chat_buffer("nonexistent") is None class TestChannelManagerStartup: def test_initial_state(self): manager = ChannelManager() assert manager.phase == "not_started" assert manager._initialized is False assert manager.context is None assert manager.broadcaster is None assert manager.watcher is None assert manager.doctor is None @pytest.mark.asyncio async def test_startup_full_flow(self): with ( patch("yuxi.channels.manager._load_builtin_adapters"), patch("yuxi.channels.manager._BUILTIN_ADAPTERS", {}), patch("yuxi.config") as mock_config, ): mock_config.channels = {} mock_pg = MagicMock() mock_db = AsyncMock() mock_pg.__aenter__ = AsyncMock(return_value=mock_db) mock_pg.__aexit__ = AsyncMock(return_value=None) with ( patch( "yuxi.channels.manager.pg_manager.get_async_session_context", return_value=mock_pg, ), patch.object(ChannelManager, "_ensure_virtual_department", AsyncMock()), patch.object(ChannelManager, "_ensure_default_agent_config", AsyncMock()), ): manager = ChannelManager() await manager.startup() assert manager.phase == "fully_running" assert manager._initialized is True assert manager.context is not None assert manager.broadcaster is not None assert manager.watcher is not None assert manager.doctor is not None assert manager.runtime_state.started_at > 0 @pytest.mark.asyncio async def test_startup_idempotent(self): manager = ChannelManager() manager._initialized = True await manager.startup() assert manager.phase == "not_started" @pytest.mark.asyncio async def test_shutdown_stops_watcher(self): with ( patch("yuxi.channels.manager._load_builtin_adapters"), patch("yuxi.channels.manager._BUILTIN_ADAPTERS", {}), patch("yuxi.config") as mock_config, ): mock_config.channels = {} mock_pg = MagicMock() mock_db = AsyncMock() mock_pg.__aenter__ = AsyncMock(return_value=mock_db) mock_pg.__aexit__ = AsyncMock(return_value=None) with ( patch( "yuxi.channels.manager.pg_manager.get_async_session_context", return_value=mock_pg, ), patch.object(ChannelManager, "_ensure_virtual_department", AsyncMock()), patch.object(ChannelManager, "_ensure_default_agent_config", AsyncMock()), ): manager = ChannelManager() await manager.startup() assert manager._watcher is not None assert manager._watcher._running is True await manager.shutdown() assert manager._initialized is False @pytest.mark.asyncio async def test_diagnose_channels_after_startup(self): with ( patch("yuxi.channels.manager._load_builtin_adapters"), patch("yuxi.channels.manager._BUILTIN_ADAPTERS", {}), patch("yuxi.config") as mock_config, ): mock_config.channels = {} mock_pg = MagicMock() mock_db = AsyncMock() mock_pg.__aenter__ = AsyncMock(return_value=mock_db) mock_pg.__aexit__ = AsyncMock(return_value=None) with ( patch( "yuxi.channels.manager.pg_manager.get_async_session_context", return_value=mock_pg, ), patch.object(ChannelManager, "_ensure_virtual_department", AsyncMock()), patch.object(ChannelManager, "_ensure_default_agent_config", AsyncMock()), ): manager = ChannelManager() await manager.startup() issues = await manager.diagnose_channels() assert isinstance(issues, list) @pytest.mark.asyncio async def test_startup_context_has_runtime_config(self): with ( patch("yuxi.channels.manager._load_builtin_adapters"), patch("yuxi.channels.manager._BUILTIN_ADAPTERS", {}), patch("yuxi.config") as mock_config, ): mock_config.channels = {"wx": {"enabled": True, "appid": "test"}} mock_pg = MagicMock() mock_db = AsyncMock() mock_pg.__aenter__ = AsyncMock(return_value=mock_db) mock_pg.__aexit__ = AsyncMock(return_value=None) with ( patch( "yuxi.channels.manager.pg_manager.get_async_session_context", return_value=mock_pg, ), patch.object(ChannelManager, "_ensure_virtual_department", AsyncMock()), patch.object(ChannelManager, "_ensure_default_agent_config", AsyncMock()), ): manager = ChannelManager() await manager.startup() ctx = manager.context assert ctx is not None assert "wx" in ctx.runtime_config class TestMessageRouterRuntime: def test_chat_abort_controllers_initialized(self): router = MessageRouter() assert router.chat_abort_controllers == {} assert router.chat_run_buffers == {} def test_abort_chat_returns_false_for_unknown_run(self): router = MessageRouter() assert router.abort_chat("nonexistent") is False @pytest.mark.asyncio async def test_abort_chat_cancels_running_task(self): router = MessageRouter() run_id = "test:msg-abort" async def slow_task(): await asyncio.sleep(10) loop = asyncio.get_running_loop() task = loop.create_task(slow_task()) router.chat_abort_controllers[run_id] = ChatAbortEntry(task=task) assert router.abort_chat(run_id) is True await asyncio.sleep(0) assert task.cancelled() def test_chat_run_buffer_registration(self): router = MessageRouter() buffer = ChatRunBuffer(run_id="test:buf-1") buffer.append_chunk("hello") router.chat_run_buffers["test:buf-1"] = buffer assert router.chat_run_buffers["test:buf-1"].get_full_text() == "hello" def test_abort_entry_has_event(self): router = MessageRouter() entry = ChatAbortEntry() router.chat_abort_controllers["run-1"] = entry router.abort_chat("run-1") assert entry.is_aborted() def test_multiple_abort_entries(self): router = MessageRouter() e1 = ChatAbortEntry() e2 = ChatAbortEntry() router.chat_abort_controllers["r1"] = e1 router.chat_abort_controllers["r2"] = e2 router.abort_chat("r1") assert e1.is_aborted() assert not e2.is_aborted() class TestEventBroadcasterRuntime: @pytest.mark.asyncio async def test_broadcast_to_multiple_queue_subscribers(self): broadcaster = EventBroadcaster() q1 = broadcaster.subscribe("ev1") q2 = broadcaster.subscribe("ev1") await broadcaster.broadcast("ev1", {"x": 1}) r1 = await asyncio.wait_for(q1.get(), timeout=1.0) r2 = await asyncio.wait_for(q2.get(), timeout=1.0) assert r1["payload"] == {"x": 1} assert r2["payload"] == {"x": 1} @pytest.mark.asyncio async def test_broadcast_to_callback_subscribers(self): broadcaster = EventBroadcaster() received = [] def handler(event, payload): received.append(payload) broadcaster.subscribe_callback("ev1", handler) await broadcaster.broadcast("ev1", {"data": "hello"}) assert received == [{"data": "hello"}] @pytest.mark.asyncio async def test_wildcard_queue_and_callback(self): broadcaster = EventBroadcaster() q = broadcaster.subscribe("*") cb_received = [] broadcaster.subscribe_callback("*", lambda e, p: cb_received.append(e)) await broadcaster.broadcast("any.event", {}) r = await asyncio.wait_for(q.get(), timeout=1.0) assert r["event"] == "any.event" assert cb_received == ["any.event"] @pytest.mark.asyncio async def test_node_send_to_session(self): broadcaster = EventBroadcaster() q = broadcaster.subscribe("session:s1") await broadcaster.node_send_to_session("s1", "ev", {"v": 1}) r = await asyncio.wait_for(q.get(), timeout=1.0) assert r == {"event": "ev", "payload": {"v": 1}} def test_subscriber_count_with_callbacks(self): broadcaster = EventBroadcaster() broadcaster.subscribe("e1") broadcaster.subscribe_callback("e1", lambda e, p: None) broadcaster.subscribe_callback("e2", lambda e, p: None) assert broadcaster.subscriber_count("e1") == 2 assert broadcaster.subscriber_count("e2") == 1 assert broadcaster.subscriber_count() == 3 class TestManagerContextIntegration: @pytest.mark.asyncio async def test_context_broadcast_delegates_to_broadcaster(self): with ( patch("yuxi.channels.manager._load_builtin_adapters"), patch("yuxi.channels.manager._BUILTIN_ADAPTERS", {}), patch("yuxi.config") as mock_config, ): mock_config.channels = {} mock_pg = MagicMock() mock_db = AsyncMock() mock_pg.__aenter__ = AsyncMock(return_value=mock_db) mock_pg.__aexit__ = AsyncMock(return_value=None) with ( patch( "yuxi.channels.manager.pg_manager.get_async_session_context", return_value=mock_pg, ), patch.object(ChannelManager, "_ensure_virtual_department", AsyncMock()), patch.object(ChannelManager, "_ensure_default_agent_config", AsyncMock()), ): manager = ChannelManager() await manager.startup() q = manager.broadcaster.subscribe("test.event") await manager.context.broadcast("test.event", {"key": "val"}) r = await asyncio.wait_for(q.get(), timeout=1.0) assert r["payload"] == {"key": "val"} @pytest.mark.asyncio async def test_context_node_send_to_session_delegates(self): with ( patch("yuxi.channels.manager._load_builtin_adapters"), patch("yuxi.channels.manager._BUILTIN_ADAPTERS", {}), patch("yuxi.config") as mock_config, ): mock_config.channels = {} mock_pg = MagicMock() mock_db = AsyncMock() mock_pg.__aenter__ = AsyncMock(return_value=mock_db) mock_pg.__aexit__ = AsyncMock(return_value=None) with ( patch( "yuxi.channels.manager.pg_manager.get_async_session_context", return_value=mock_pg, ), patch.object(ChannelManager, "_ensure_virtual_department", AsyncMock()), patch.object(ChannelManager, "_ensure_default_agent_config", AsyncMock()), ): manager = ChannelManager() await manager.startup() q = manager.broadcaster.subscribe("session:s1") await manager.context.node_send_to_session("s1", "ev", {"v": 1}) r = await asyncio.wait_for(q.get(), timeout=1.0) assert r == {"event": "ev", "payload": {"v": 1}}