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

from __future__ import annotations

from unittest.mock import AsyncMock, patch

import pytest

from yuxi.channels.adapters.yuanbao.adapter import YuanbaoAdapter
from yuxi.channels.adapters.yuanbao.format import format_outbound
from yuxi.channels.adapters.yuanbao.security import (
    check_dm_policy,
    check_group_policy,
    check_mention_required,
)
from yuxi.channels.adapters.yuanbao.session import (
    resolve_agent_route,
    resolve_chat_type_str,
    resolve_thread_key,
)
from yuxi.channels.adapters.yuanbao.streaming import send_blocks_stream
from yuxi.channels.adapters.yuanbao.token import YuanbaoTokenManager
from yuxi.channels.adapters.yuanbao.chunking import chunk_text
from yuxi.channels.models import (
    Attachment,
    ChannelIdentity,
    ChannelMessage,
    ChannelResponse,
    ChannelType,
    ChatType,
    MessageType,
    MentionsInfo,
)

YUANBAO = ChannelType.YUANBAO


def _make_identity(
    channel_chat_id: str = "user_openid_001",
    channel_user_id: str = "user_openid_001",
    channel_message_id: str = "msg_001",
) -> ChannelIdentity:
    return ChannelIdentity(
        channel_id="yuanbao",
        channel_type=YUANBAO,
        channel_user_id=channel_user_id,
        channel_chat_id=channel_chat_id,
        channel_message_id=channel_message_id,
    )


def _make_response(
    content: str = "Hello from Yuanbao",
    channel_chat_id: str = "user_openid_001",
    **kwargs,
) -> ChannelResponse:
    identity = _make_identity(channel_chat_id=channel_chat_id)
    return ChannelResponse(identity=identity, content=content, **kwargs)


# ==================== Token Manager Tests ====================


class TestYuanbaoTokenManager:
    def test_token_initial_state(self):
        tm = YuanbaoTokenManager(app_key="test_key", app_secret="test_secret", bot_app_id="bot_001")
        assert tm._access_token is None
        assert tm._is_expired() is True

    def test_api_base(self):
        tm = YuanbaoTokenManager(app_key="test_key", app_secret="test_secret", bot_app_id="bot_001")
        assert tm.api_base == "https://open-api.yuanbao.tencent.com"

    def test_bot_app_id(self):
        tm = YuanbaoTokenManager(app_key="test_key", app_secret="test_secret", bot_app_id="my_bot")
        assert tm.bot_app_id == "my_bot"

    def test_sign_request(self):
        tm = YuanbaoTokenManager(app_key="test_key", app_secret="test_secret", bot_app_id="bot_001")
        signature = tm._sign_request(1700000000)
        assert isinstance(signature, str)
        assert len(signature) == 64

    def test_is_expired_with_no_token(self):
        tm = YuanbaoTokenManager(app_key="test_key", app_secret="test_secret", bot_app_id="bot_001")
        assert tm._is_expired() is True

    @pytest.mark.asyncio
    async def test_get_token_refreshes(self):
        tm = YuanbaoTokenManager(app_key="test_key", app_secret="test_secret", bot_app_id="bot_001")
        tm._access_token = None
        tm._expires_at = None

        with patch.object(tm, "_refresh", new_callable=AsyncMock) as mock_refresh:
            mock_refresh.return_value = None
            await tm.get_token()
            mock_refresh.assert_called_once()

    @pytest.mark.asyncio
    async def test_refresh_token(self):
        tm = YuanbaoTokenManager(app_key="test_key", app_secret="test_secret", bot_app_id="bot_001")

        with patch.object(tm, "get_token", new_callable=AsyncMock) as mock_get:
            mock_get.return_value = "new_token"
            result = await tm.refresh_token()
            assert result == "new_token"


# ==================== Format Tests ====================


class TestFormatOutbound:
    def test_text_payload_direct(self):
        resp = _make_response(content="Hello World", channel_chat_id="user_openid_001")
        payload = format_outbound(resp)
        assert payload["content"] == "Hello World"
        assert payload["msg_type"] == "text"
        assert payload["open_id"] == "user_openid_001"

    def test_text_payload_group(self):
        resp = _make_response(content="Group message", channel_chat_id="group_openid_abc")
        resp.metadata = {"group_open_id": "group_openid_abc"}
        payload = format_outbound(resp)
        assert payload["content"] == "Group message"
        assert payload["group_open_id"] == "group_openid_abc"

    def test_text_payload_channel(self):
        resp = _make_response(content="Channel message", channel_chat_id="ch_001")
        resp.metadata = {"channel_id": "ch_001"}
        payload = format_outbound(resp)
        assert payload["content"] == "Channel message"
        assert payload["channel_id"] == "ch_001"

    def test_image_payload(self):
        resp = _make_response(content="Image caption", channel_chat_id="user_001")
        resp.message_type = MessageType.IMAGE
        resp.attachments = [Attachment(type="image", url="https://example.com/img.png")]
        payload = format_outbound(resp)
        assert payload["msg_type"] == "image"
        assert payload["media_url"] == "https://example.com/img.png"

    def test_file_payload(self):
        resp = _make_response(content="File description", channel_chat_id="user_001")
        resp.message_type = MessageType.FILE
        resp.attachments = [Attachment(type="file", url="https://example.com/doc.pdf", filename="doc.pdf")]
        payload = format_outbound(resp)
        assert payload["msg_type"] == "file"
        assert payload["media_url"] == "https://example.com/doc.pdf"
        assert payload["filename"] == "doc.pdf"

    def test_payload_with_reply(self):
        resp = _make_response(content="Reply", reply_to_message_id="msg_orig")
        payload = format_outbound(resp)
        assert payload["reply_to_msg_id"] == "msg_orig"

    def test_audio_as_file(self):
        resp = _make_response(content="Audio file", channel_chat_id="user_001")
        resp.message_type = MessageType.AUDIO
        payload = format_outbound(resp)
        assert payload["msg_type"] == "audio"

    def test_extra_metadata(self):
        resp = _make_response(content="Extra", channel_chat_id="user_001")
        resp.metadata = {"extra": {"key": "value"}}
        payload = format_outbound(resp)
        assert payload["extra"] == {"key": "value"}


# ==================== Normalize Inbound Tests ====================


class TestNormalizeInbound:
    @pytest.fixture
    def adapter(self):
        config = {"app_key": "test_key", "app_secret": "test_secret", "dm_policy": "open"}
        return YuanbaoAdapter(config=config)

    def test_direct_message(self, adapter):
        raw = {
            "type": "message",
            "open_id": "user_123",
            "msg_id": "msg_001",
            "msg_type": "text",
            "content": "你好",
            "chat_type": "direct",
            "timestamp": 1700000000,
        }
        msg = adapter.normalize_inbound(raw)
        assert msg.identity.channel_chat_id == "user_123"
        assert msg.chat_type == ChatType.DIRECT
        assert msg.content == "你好"
        assert msg.message_type == MessageType.TEXT

    def test_group_message(self, adapter):
        raw = {
            "type": "message",
            "open_id": "user_456",
            "group_open_id": "group_abc",
            "msg_id": "msg_002",
            "msg_type": "text",
            "content": "群聊消息",
            "chat_type": "group",
            "timestamp": 1700000000,
        }
        msg = adapter.normalize_inbound(raw)
        assert msg.identity.channel_chat_id == "group_abc"
        assert msg.chat_type == ChatType.GROUP
        assert msg.metadata.get("group_open_id") == "group_abc"

    def test_channel_message(self, adapter):
        raw = {
            "type": "message",
            "open_id": "user_789",
            "channel_id": "ch_guild_001",
            "msg_id": "msg_003",
            "msg_type": "text",
            "content": "频道消息",
            "chat_type": "channel",
            "timestamp": 1700000000,
        }
        msg = adapter.normalize_inbound(raw)
        assert msg.identity.channel_chat_id == "ch_guild_001"
        assert msg.chat_type == ChatType.GUILD_CHANNEL
        assert msg.metadata.get("channel_id") == "ch_guild_001"

    def test_image_message(self, adapter):
        raw = {
            "type": "message",
            "open_id": "user_123",
            "msg_id": "msg_004",
            "msg_type": "image",
            "content": "",
            "chat_type": "direct",
            "timestamp": 1700000000,
            "attachments": [
                {
                    "type": "image",
                    "url": "https://example.com/photo.jpg",
                    "filename": "photo.jpg",
                    "size": 102400,
                }
            ],
        }
        msg = adapter.normalize_inbound(raw)
        assert msg.message_type == MessageType.IMAGE
        assert len(msg.attachments) == 1
        assert msg.attachments[0].url == "https://example.com/photo.jpg"

    def test_file_message(self, adapter):
        raw = {
            "type": "message",
            "open_id": "user_123",
            "msg_id": "msg_005",
            "msg_type": "file",
            "content": "文档",
            "chat_type": "direct",
            "timestamp": 1700000000,
            "attachments": [
                {
                    "type": "file",
                    "url": "https://example.com/doc.pdf",
                    "filename": "doc.pdf",
                    "size": 204800,
                }
            ],
        }
        msg = adapter.normalize_inbound(raw)
        assert msg.message_type == MessageType.FILE
        assert len(msg.attachments) == 1

    def test_markdown_message(self, adapter):
        raw = {
            "type": "message",
            "open_id": "user_123",
            "msg_id": "msg_006",
            "msg_type": "markdown",
            "content": "**Bold** text",
            "chat_type": "direct",
            "timestamp": 1700000000,
        }
        msg = adapter.normalize_inbound(raw)
        assert msg.message_type == MessageType.TEXT
        assert msg.content == "**Bold** text"

    def test_command_detection(self, adapter):
        raw = {
            "type": "message",
            "open_id": "user_123",
            "msg_id": "msg_007",
            "msg_type": "text",
            "content": "/reset",
            "chat_type": "direct",
            "timestamp": 1700000000,
        }
        msg = adapter.normalize_inbound(raw)
        assert msg.message_type == MessageType.COMMAND

    def test_metadata_includes_yuanbao_chat_type(self, adapter):
        raw = {
            "type": "message",
            "open_id": "user_123",
            "msg_id": "msg_008",
            "msg_type": "text",
            "content": "Hello",
            "chat_type": "direct",
            "reply_to_msg_id": "msg_prev",
            "timestamp": 1700000000,
        }
        msg = adapter.normalize_inbound(raw)
        assert msg.metadata["yuanbao_chat_type"] == "direct"
        assert msg.metadata["reply_to_msg_id"] == "msg_prev"

    def test_group_message_mention(self, adapter):
        adapter._bot_info = {"username": "mybot"}
        raw = {
            "type": "message",
            "open_id": "user_456",
            "group_open_id": "group_abc",
            "msg_id": "msg_009",
            "msg_type": "text",
            "content": "@mybot 你好",
            "chat_type": "group",
            "timestamp": 1700000000,
        }
        msg = adapter.normalize_inbound(raw)
        assert msg.mentions is not None
        assert msg.mentions.is_bot_mentioned is True


# ==================== Security Tests ====================


class TestSecurityPolicy:
    @pytest.mark.asyncio
    async def test_dm_policy_open(self):
        config = {"dm_policy": "open"}
        result = await check_dm_policy("user_001", config)
        assert result is True

    @pytest.mark.asyncio
    async def test_dm_policy_disabled(self):
        config = {"dm_policy": "disabled"}
        result = await check_dm_policy("user_001", config)
        assert result is False

    @pytest.mark.asyncio
    async def test_dm_policy_allowlist_match(self):
        config = {"dm_policy": "allowlist", "allow_from": ["yb:user_001", "yb:user_002"]}
        result = await check_dm_policy("yb:user_001", config)
        assert result is True

    @pytest.mark.asyncio
    async def test_dm_policy_allowlist_no_match(self):
        config = {"dm_policy": "allowlist", "allow_from": ["yb:user_002"]}
        result = await check_dm_policy("user_001", config)
        assert result is False

    @pytest.mark.asyncio
    async def test_dm_policy_unknown(self):
        config = {"dm_policy": "unknown_policy"}
        result = await check_dm_policy("user_001", config)
        assert result is False

    @pytest.mark.asyncio
    async def test_group_policy_open(self):
        config = {"group_policy": "open"}
        result = await check_group_policy("group_abc", "user_001", config)
        assert result is True

    @pytest.mark.asyncio
    async def test_group_policy_disabled(self):
        config = {"group_policy": "disabled"}
        result = await check_group_policy("group_abc", "user_001", config)
        assert result is False

    @pytest.mark.asyncio
    async def test_group_policy_allowlist_global(self):
        config = {"group_policy": "allowlist", "group_allow_from": ["yb:user_001"]}
        result = await check_group_policy("group_abc", "yb:user_001", config)
        assert result is True

    @pytest.mark.asyncio
    async def test_group_policy_allowlist_per_group(self):
        config = {
            "group_policy": "allowlist",
            "groups": {"group_abc": {"allow_from": ["yb:user_001"]}},
        }
        result = await check_group_policy("group_abc", "yb:user_001", config)
        assert result is True

    @pytest.mark.asyncio
    async def test_mention_required_disabled(self):
        config = {"group_require_mention": False}
        msg = ChannelMessage(
            identity=_make_identity(),
            content="Hello",
            chat_type=ChatType.GROUP,
        )
        result = await check_mention_required("group_abc", msg, config)
        assert result is True

    @pytest.mark.asyncio
    async def test_mention_required_with_bot_mentioned(self):
        config = {"group_require_mention": True}
        msg = ChannelMessage(
            identity=_make_identity(),
            content="Hello",
            chat_type=ChatType.GROUP,
            mentions=MentionsInfo(is_bot_mentioned=True),
        )
        result = await check_mention_required("group_abc", msg, config)
        assert result is True

    @pytest.mark.asyncio
    async def test_mention_required_not_mentioned(self):
        config = {"group_require_mention": True}
        msg = ChannelMessage(
            identity=_make_identity(),
            content="Hello",
            chat_type=ChatType.GROUP,
        )
        result = await check_mention_required("group_abc", msg, config)
        assert result is False

    @pytest.mark.asyncio
    async def test_mention_required_by_name(self):
        config = {"group_require_mention": True}
        msg = ChannelMessage(
            identity=_make_identity(),
            content="@mybot help",
            chat_type=ChatType.GROUP,
        )
        result = await check_mention_required("group_abc", msg, config, bot_names=["mybot"])
        assert result is True


# ==================== Session Tests ====================


class TestSessionRouting:
    def test_resolve_thread_key_direct(self):
        identity = _make_identity(channel_chat_id="user_openid_123")
        thread_key = resolve_thread_key(identity, chat_type="direct")
        assert thread_key == "yuanbao:direct:user_openid_123"

    def test_resolve_thread_key_group(self):
        identity = _make_identity(channel_chat_id="group_openid_abc")
        thread_key = resolve_thread_key(identity, chat_type="group")
        assert thread_key == "yuanbao:group:group_openid_abc"

    def test_resolve_thread_key_channel(self):
        identity = _make_identity(channel_chat_id="ch_guild_001")
        thread_key = resolve_thread_key(identity, chat_type="channel")
        assert thread_key == "yuanbao:channel:ch_guild_001"

    def test_resolve_chat_type_str_direct(self):
        assert resolve_chat_type_str("direct") == "direct"

    def test_resolve_chat_type_str_group(self):
        assert resolve_chat_type_str("group") == "group"

    def test_resolve_chat_type_str_channel(self):
        assert resolve_chat_type_str("channel") == "guild_channel"

    def test_resolve_agent_route_direct(self):
        identity = _make_identity(channel_chat_id="user_123")
        route = resolve_agent_route(identity, chat_type="direct", default_agent_id="my_bot")
        assert route == "agent:my_bot:yuanbao:direct:user_123"

    def test_resolve_agent_route_group(self):
        identity = _make_identity(channel_chat_id="group_abc")
        route = resolve_agent_route(identity, chat_type="group", default_agent_id="default")
        assert route == "agent:default:yuanbao:group:group_abc"

    def test_resolve_agent_route_group_custom_agent(self):
        identity = _make_identity(channel_chat_id="group_abc")
        groups_config = {"group_abc": {"agent_id": "custom_agent"}}
        route = resolve_agent_route(
            identity, chat_type="group",
            default_agent_id="default", groups_config=groups_config,
        )
        assert route == "agent:custom_agent:yuanbao:group:group_abc"


# ==================== Chunking Tests ====================


class TestChunking:
    def test_short_text(self):
        result = chunk_text("Short text", limit=20000)
        assert result == ["Short text"]

    def test_exact_limit(self):
        text = "A" * 5000
        result = chunk_text(text, limit=6000)
        assert len(result) == 1

    def test_multiple_chunks(self):
        text = "\n\n".join([f"Paragraph {i}" for i in range(200)])
        result = chunk_text(text, limit=500)
        assert len(result) > 1

    def test_single_very_long_paragraph(self):
        text = "A" * 10000
        result = chunk_text(text, limit=1000)
        assert len(result) > 1
        for chunk in result:
            assert len(chunk) <= 1000


# ==================== Adapter Tests ====================


class TestYuanbaoAdapter:
    @pytest.fixture
    def adapter(self):
        config = {"app_key": "test_key", "app_secret": "test_secret", "dm_policy": "open"}
        return YuanbaoAdapter(config=config)

    def test_channel_id(self, adapter):
        assert adapter.channel_id == "yuanbao"

    def test_channel_type(self, adapter):
        assert adapter.channel_type == YUANBAO

    def test_text_chunk_limit(self, adapter):
        assert adapter.text_chunk_limit == 3000

    def test_supports_markdown(self, adapter):
        assert adapter.supports_markdown is True

    def test_supports_streaming(self, adapter):
        assert adapter.supports_streaming is True

    def test_streaming_modes(self, adapter):
        assert "off" in adapter.streaming_modes
        assert "block" in adapter.streaming_modes

    def test_max_media_size_mb(self, adapter):
        assert adapter.max_media_size_mb == 20

    def test_initial_status(self, adapter):
        assert adapter.status == "disconnected"

    def test_webhook_path(self, adapter):
        assert adapter.webhook_path is None

    def test_format_outbound_text(self, adapter):
        resp = _make_response(content="Hello")
        payload = adapter.format_outbound(resp)
        assert payload["content"] == "Hello"
        assert payload["msg_type"] == "text"

    def test_format_outbound_group(self, adapter):
        resp = _make_response(content="Group msg", channel_chat_id="group_abc")
        resp.metadata = {"group_open_id": "group_abc"}
        payload = adapter.format_outbound(resp)
        assert payload["group_open_id"] == "group_abc"

    @pytest.mark.asyncio
    async def test_health_check_without_token(self, adapter):
        status = await adapter.health_check()
        assert status.status == "unhealthy"

    @pytest.mark.asyncio
    async def test_refresh_token_no_manager(self, adapter):
        result = await adapter._refresh_token_if_needed()
        assert result is False

    @pytest.mark.asyncio
    async def test_send_without_client(self, adapter):
        resp = _make_response(content="Test")
        result = await adapter.send(resp)
        assert result.success is False

    def test_receive_yields_nothing(self, adapter):
        async def _collect():
            collected = []
            async for msg in adapter.receive():
                collected.append(msg)
            return collected

        import asyncio
        result = asyncio.run(_collect())
        assert result == []

    @pytest.mark.asyncio
    async def test_check_security_dm_open(self, adapter):
        msg = ChannelMessage(
            identity=_make_identity(channel_chat_id="user_123"),
            content="Hello",
            chat_type=ChatType.DIRECT,
        )
        result = await adapter._check_security(msg)
        assert result is True

    @pytest.mark.asyncio
    async def test_check_security_group_open(self, adapter):
        adapter.config["group_policy"] = "open"
        adapter.config["group_require_mention"] = False
        msg = ChannelMessage(
            identity=_make_identity(channel_chat_id="group_abc"),
            content="Hello",
            chat_type=ChatType.GROUP,
            metadata={"group_open_id": "group_abc"},
        )
        result = await adapter._check_security(msg)
        assert result is True

    @pytest.mark.asyncio
    async def test_check_security_channel(self, adapter):
        msg = ChannelMessage(
            identity=_make_identity(channel_chat_id="ch_001"),
            content="Hello",
            chat_type=ChatType.GUILD_CHANNEL,
        )
        result = await adapter._check_security(msg)
        assert result is True

    def test_resolve_chat_context_direct(self, adapter):
        raw = {"open_id": "user_001", "chat_type": "direct"}
        chat_type, chat_id = adapter._resolve_chat_context(raw)
        assert chat_type == ChatType.DIRECT
        assert chat_id == "user_001"

    def test_resolve_chat_context_group(self, adapter):
        raw = {"open_id": "user_002", "group_open_id": "group_abc", "chat_type": "group"}
        chat_type, chat_id = adapter._resolve_chat_context(raw)
        assert chat_type == ChatType.GROUP
        assert chat_id == "group_abc"

    def test_resolve_chat_context_channel(self, adapter):
        raw = {"open_id": "user_003", "channel_id": "ch_001", "chat_type": "channel"}
        chat_type, chat_id = adapter._resolve_chat_context(raw)
        assert chat_type == ChatType.GUILD_CHANNEL
        assert chat_id == "ch_001"


# ==================== Streaming Tests ====================


class TestStreaming:
    @pytest.mark.asyncio
    async def test_stream_blocks(self):
        text = "Para 1\n\nPara 2\n\nPara 3"
        sent_messages = []

        async def mock_send(response):
            sent_messages.append(response.content)

        await send_blocks_stream(
            "chat_001", text, mock_send,
            chunk_size=1,
        )
        assert len(sent_messages) > 0
        full_text = " ".join(sent_messages)
        assert "Para 1" in full_text
        assert "Para 2" in full_text
        assert "Para 3" in full_text

    @pytest.mark.asyncio
    async def test_stream_with_metadata(self):
        text = "Single block"
        sent_messages = []

        async def mock_send(response):
            sent_messages.append(response)

        await send_blocks_stream(
            "chat_001", text, mock_send,
            metadata={"group_open_id": "group_abc"},
            chunk_size=1,
        )
        assert len(sent_messages) == 1
        assert sent_messages[0].metadata["group_open_id"] == "group_abc"


# ==================== Probe Tests ====================


class TestProbe:
    @pytest.mark.asyncio
    async def test_health_check_healthy(self):
        from yuxi.channels.adapters.yuanbao.probe import health_check_yuanbao

        mock_resp = AsyncMock()
        mock_resp.status = 200
        mock_resp.__aenter__ = AsyncMock(return_value=mock_resp)
        mock_resp.__aexit__ = AsyncMock(return_value=None)
        mock_resp.json = AsyncMock(return_value={"bot_app_id": "bot_001"})

        with patch("aiohttp.ClientSession.get", return_value=mock_resp):
            result = await health_check_yuanbao(
                "https://open-api.yuanbao.tencent.com", "test_token", ws_connected=True,
            )
            assert result.status == "healthy"
            assert result.metadata["ws_connected"] is True

    @pytest.mark.asyncio
    async def test_health_check_unhealthy_auth(self):
        from yuxi.channels.adapters.yuanbao.probe import health_check_yuanbao

        mock_resp = AsyncMock()
        mock_resp.status = 401
        mock_resp.__aenter__ = AsyncMock(return_value=mock_resp)
        mock_resp.__aexit__ = AsyncMock(return_value=None)

        with patch("aiohttp.ClientSession.get", return_value=mock_resp):
            result = await health_check_yuanbao(
                "https://open-api.yuanbao.tencent.com", "bad_token",
            )
            assert result.status == "unhealthy"
            assert "Token" in result.last_error

    @pytest.mark.asyncio
    async def test_health_check_exception(self):
        from yuxi.channels.adapters.yuanbao.probe import health_check_yuanbao

        with patch("aiohttp.ClientSession.get", side_effect=Exception("Network error")):
            result = await health_check_yuanbao(
                "https://open-api.yuanbao.tencent.com", "test_token",
            )
            assert result.status == "unhealthy"
            assert "Network error" in result.last_error