ForcePilot/backend/test/unit/channels/test_runtime.py

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.registry._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.registry._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.registry._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.registry._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.registry._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.registry._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}}