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