ForcePilot/backend/test/unit/channel/interfaces/websocket/test_manager.py

from __future__ import annotations

import asyncio
import pytest
from uuid import uuid4

from yuxi.channel.interfaces.websocket.manager import WsConnectionManager
from yuxi.channel.domain.port.ws_connection_port import WsConnectionPort
from yuxi.channel.domain.port.channel_adapter_port import ChannelAdapterPort
from yuxi.channel.domain.model.message.unified_message import UnifiedMessage
from yuxi.channel.domain.model.message.peer import Peer
from yuxi.channel.domain.model.shared.channel_type import ChannelType


class _FakeWsConnection(WsConnectionPort):
    def __init__(self, channel_type: str, connected: bool = True):
        self._channel_type = channel_type
        self._connected = connected
        self.started = False
        self.stopped = False

    @property
    def channel_type(self) -> str:
        return self._channel_type

    @property
    def is_connected(self) -> bool:
        return self._connected

    async def start(self, loop, on_message) -> None:
        self.started = True

    async def stop(self) -> None:
        self.stopped = True


class _FakeAdapter(ChannelAdapterPort):
    def __init__(self, channel_type: str):
        self._channel_type = channel_type

    @property
    def channel_type(self) -> str:
        return self._channel_type

    @property
    def ws_connection(self):
        return None

    async def open(self) -> None:
        pass

    async def close(self) -> None:
        pass

    async def receive_message(self, raw: dict) -> UnifiedMessage:
        return UnifiedMessage(
            message_id="test",
            channel_type=ChannelType.WEB,
            sender=Peer(id="user", name="User"),
            content="hello",
        )

    async def send_message(self, session_id: str, content: str, *, channel_type: str, metadata: dict):
        from yuxi.channel.domain.model.message.dispatch_result import SendResult
        return SendResult(success=True)

    async def send_typing(self, session_id: str) -> None:
        pass

    async def send_media(self, session_id: str, *, url: str, media_type: str, metadata: dict) -> bool:
        return True

    async def is_healthy(self) -> bool:
        return True

    @property
    def capabilities(self):
        from yuxi.channel.domain.model.shared.channel_capabilities import ChannelCapabilities
        return ChannelCapabilities()

    @classmethod
    def get_default_config(cls) -> dict:
        return {}


class _FakeInboundService:
    def __init__(self):
        self.submitted = []

    async def submit(self, message, *, channel_type: str, trace_id: str = ""):
        self.submitted.append({"message": message, "channel_type": channel_type, "trace_id": trace_id})


class _FakeMetrics:
    pass


@pytest.fixture
def manager():
    return WsConnectionManager(metrics=_FakeMetrics())


@pytest.mark.asyncio
async def test_ws_manager_register(manager):
    conn = _FakeWsConnection("feishu")
    manager.register(conn)
    assert "feishu" in manager._connections


@pytest.mark.asyncio
async def test_ws_manager_set_adapters(manager):
    adapters = {"feishu": _FakeAdapter("feishu")}
    manager.set_adapters(adapters)
    assert manager._adapters == adapters


@pytest.mark.asyncio
async def test_ws_manager_start_all(manager):
    conn = _FakeWsConnection("feishu")
    manager.register(conn)

    inbound = _FakeInboundService()
    loop = asyncio.get_event_loop()

    await manager.start_all(loop, inbound)
    assert conn.started is True

    await manager.stop_all()
    assert conn.stopped is True


@pytest.mark.asyncio
async def test_ws_manager_connections_status(manager):
    conn1 = _FakeWsConnection("feishu", connected=True)
    conn2 = _FakeWsConnection("dingtalk", connected=False)
    manager.register(conn1)
    manager.register(conn2)

    status = manager.connections_status
    assert status["feishu"] is True
    assert status["dingtalk"] is False


@pytest.mark.asyncio
async def test_ws_manager_on_message(manager):
    adapter = _FakeAdapter("feishu")
    manager.set_adapters({"feishu": adapter})

    inbound = _FakeInboundService()
    loop = asyncio.get_event_loop()
    await manager.start_all(loop, inbound)

    raw = {
        "channel_type": "feishu",
        "header": {"event_id": "test-event-id"},
        "content": "hello",
    }
    await manager._on_message(raw)

    assert len(inbound.submitted) == 1
    assert inbound.submitted[0]["channel_type"] == "feishu"
    assert inbound.submitted[0]["trace_id"] == "test-event-id"

    await manager.stop_all()


@pytest.mark.asyncio
async def test_ws_manager_on_message_no_adapter(manager):
    manager.set_adapters({})

    inbound = _FakeInboundService()
    loop = asyncio.get_event_loop()
    await manager.start_all(loop, inbound)

    raw = {
        "channel_type": "unknown",
        "header": {"event_id": "test-event-id"},
    }
    await manager._on_message(raw)

    assert len(inbound.submitted) == 0

    await manager.stop_all()


@pytest.mark.asyncio
async def test_ws_manager_on_message_receive_error(manager):
    class _BadAdapter(_FakeAdapter):
        async def receive_message(self, raw: dict) -> UnifiedMessage:
            raise ValueError("parse error")

    manager.set_adapters({"feishu": _BadAdapter("feishu")})

    inbound = _FakeInboundService()
    loop = asyncio.get_event_loop()
    await manager.start_all(loop, inbound)

    raw = {
        "channel_type": "feishu",
        "header": {"event_id": "test-event-id"},
    }
    await manager._on_message(raw)

    assert len(inbound.submitted) == 0

    await manager.stop_all()


@pytest.mark.asyncio
async def test_ws_manager_on_message_submit_error(manager):
    class _BadInboundService:
        async def submit(self, message, *, channel_type: str, trace_id: str = ""):
            raise ValueError("submit error")

    adapter = _FakeAdapter("feishu")
    manager.set_adapters({"feishu": adapter})

    inbound = _BadInboundService()
    loop = asyncio.get_event_loop()
    await manager.start_all(loop, inbound)

    raw = {
        "channel_type": "feishu",
        "header": {"event_id": "test-event-id"},
    }
    # Should not raise
    await manager._on_message(raw)

    await manager.stop_all()