MineNasAI/tests/test_gateway.py

"""Gateway 模块测试"""

from __future__ import annotations

import pytest

from minenasai.gateway.protocol import (
    ChatMessage,
    MessageType,
    TaskComplexity,
    parse_message,
)
from minenasai.gateway.router import SmartRouter


class TestProtocol:
    """协议测试"""

    def test_chat_message(self):
        """测试聊天消息"""
        msg = ChatMessage(content="你好")

        assert msg.type == MessageType.CHAT
        assert msg.content == "你好"
        assert msg.id is not None

    def test_parse_message(self):
        """测试消息解析"""
        data = {
            "type": "chat",
            "content": "测试消息",
        }
        msg = parse_message(data)

        assert isinstance(msg, ChatMessage)
        assert msg.content == "测试消息"

    def test_parse_invalid_type(self):
        """测试无效消息类型"""
        data = {"type": "invalid"}

        with pytest.raises(ValueError, match="未知的消息类型"):
            parse_message(data)

    def test_parse_missing_type(self):
        """测试缺少类型"""
        data = {"content": "test"}

        with pytest.raises(ValueError, match="缺少 type"):
            parse_message(data)


class TestRouter:
    """智能路由测试"""

    def setup_method(self):
        """初始化路由器"""
        self.router = SmartRouter()

    def test_simple_question(self):
        """测试简单问题"""
        result = self.router.evaluate("今天天气怎么样？")

        assert result["complexity"] == TaskComplexity.SIMPLE
        assert result["suggested_handler"] == "quick_response"

    def test_complex_task(self):
        """测试复杂任务"""
        # 包含多个复杂关键词：重构、实现、优化
        result = self.router.evaluate(
            "请帮我重构这个项目的数据库模块，实现异步操作支持，"
            "优化连接池管理，同时要保持向后兼容。这是一个架构设计任务。"
        )

        # 复杂任务应该识别为 COMPLEX 或 MEDIUM
        assert result["complexity"] in [TaskComplexity.COMPLEX, TaskComplexity.MEDIUM]

    def test_medium_task(self):
        """测试中等任务"""
        result = self.router.evaluate("查看当前目录下的文件列表")

        assert result["complexity"] in [TaskComplexity.SIMPLE, TaskComplexity.MEDIUM]

    def test_command_override_simple(self):
        """测试命令覆盖 - 简单"""
        result = self.router.evaluate("/快速 帮我重构整个项目")

        assert result["complexity"] == TaskComplexity.SIMPLE
        assert result["confidence"] == 1.0
        assert result["content"] == "帮我重构整个项目"

    def test_command_override_complex(self):
        """测试命令覆盖 - 复杂"""
        result = self.router.evaluate("/深度 你好")

        assert result["complexity"] == TaskComplexity.COMPLEX
        assert result["confidence"] == 1.0

    def test_code_detection(self):
        """测试代码检测"""
        result = self.router.evaluate(
            "请帮我实现这个函数：\n```python\ndef hello():\n    pass\n```"
        )

        assert result["complexity"] == TaskComplexity.COMPLEX

    def test_multi_step_detection(self):
        """测试多步骤检测"""
        # 使用英文 step 模式和更多内容
        result = self.router.evaluate(
            "Step 1: 创建数据库表结构\n"
            "Step 2: 实现数据导入功能\n"
            "Step 3: 开发验证脚本\n"
            "Step 4: 部署到生产环境"
        )

        # 多步骤任务应该识别为复杂任务
        assert result["complexity"] in [TaskComplexity.COMPLEX, TaskComplexity.MEDIUM]


class TestRouterEdgeCases:
    """路由器边界情况测试"""

    def setup_method(self):
        self.router = SmartRouter()

    def test_empty_content(self):
        """测试空内容"""
        result = self.router.evaluate("")

        assert result["complexity"] == TaskComplexity.SIMPLE

    def test_very_long_content(self):
        """测试超长内容"""
        long_content = "请帮我分析 " + "这段代码 " * 200
        result = self.router.evaluate(long_content)

        assert result["complexity"] == TaskComplexity.COMPLEX

    def test_special_characters(self):
        """测试特殊字符"""
        result = self.router.evaluate("查看 /tmp/test.txt 文件内容")

        assert result["complexity"] in [TaskComplexity.SIMPLE, TaskComplexity.MEDIUM]


class TestConnectionManager:
    """WebSocket 连接管理器测试"""

    def test_import_manager(self):
        """测试导入连接管理器"""
        from minenasai.gateway.server import ConnectionManager

        manager = ConnectionManager()
        assert manager.active_connections == {}

    @pytest.mark.asyncio
    async def test_connect_and_disconnect(self):
        """测试连接和断开"""
        from unittest.mock import AsyncMock, MagicMock

        from minenasai.gateway.server import ConnectionManager

        manager = ConnectionManager()

        # Mock WebSocket
        mock_ws = AsyncMock()
        mock_ws.accept = AsyncMock()

        # 连接
        await manager.connect(mock_ws, "client-1")
        assert "client-1" in manager.active_connections
        mock_ws.accept.assert_called_once()

        # 断开
        manager.disconnect("client-1")
        assert "client-1" not in manager.active_connections

    @pytest.mark.asyncio
    async def test_disconnect_nonexistent(self):
        """测试断开不存在的连接"""
        from minenasai.gateway.server import ConnectionManager

        manager = ConnectionManager()
        # 不应该抛出异常
        manager.disconnect("nonexistent")

    @pytest.mark.asyncio
    async def test_send_message(self):
        """测试发送消息"""
        from unittest.mock import AsyncMock

        from minenasai.gateway.server import ConnectionManager

        manager = ConnectionManager()

        # Mock WebSocket
        mock_ws = AsyncMock()
        mock_ws.accept = AsyncMock()
        mock_ws.send_json = AsyncMock()

        # 连接
        await manager.connect(mock_ws, "client-1")

        # 发送消息
        await manager.send_message("client-1", {"type": "test"})
        mock_ws.send_json.assert_called_once_with({"type": "test"})

    @pytest.mark.asyncio
    async def test_send_message_to_nonexistent(self):
        """测试发送消息给不存在的客户端"""
        from minenasai.gateway.server import ConnectionManager

        manager = ConnectionManager()
        # 不应该抛出异常
        await manager.send_message("nonexistent", {"type": "test"})

    @pytest.mark.asyncio
    async def test_broadcast(self):
        """测试广播消息"""
        from unittest.mock import AsyncMock

        from minenasai.gateway.server import ConnectionManager

        manager = ConnectionManager()

        # Mock 多个 WebSocket
        mock_ws1 = AsyncMock()
        mock_ws1.accept = AsyncMock()
        mock_ws1.send_json = AsyncMock()

        mock_ws2 = AsyncMock()
        mock_ws2.accept = AsyncMock()
        mock_ws2.send_json = AsyncMock()

        # 连接
        await manager.connect(mock_ws1, "client-1")
        await manager.connect(mock_ws2, "client-2")

        # 广播
        await manager.broadcast({"type": "broadcast"})
        mock_ws1.send_json.assert_called_once_with({"type": "broadcast"})
        mock_ws2.send_json.assert_called_once_with({"type": "broadcast"})


class TestGatewayServer:
    """Gateway 服务器测试"""

    def test_import_app(self):
        """测试导入应用"""
        from minenasai.gateway.server import app

        assert app is not None
        assert app.title == "MineNASAI Gateway"

    def test_import_endpoints(self):
        """测试导入端点函数"""
        from minenasai.gateway.server import list_agents, list_sessions, root

        assert callable(root)
        assert callable(list_agents)
        assert callable(list_sessions)


class TestMessageTypes:
    """消息类型测试"""

    def test_status_message(self):
        """测试状态消息"""
        from minenasai.gateway.protocol import StatusMessage

        msg = StatusMessage(status="thinking", message="处理中...")
        assert msg.type == MessageType.STATUS
        assert msg.status == "thinking"
        assert msg.message == "处理中..."

    def test_response_message(self):
        """测试响应消息"""
        from minenasai.gateway.protocol import ResponseMessage

        msg = ResponseMessage(content="Hello!", in_reply_to="msg-123")
        assert msg.type == MessageType.RESPONSE
        assert msg.content == "Hello!"
        assert msg.in_reply_to == "msg-123"

    def test_error_message(self):
        """测试错误消息"""
        from minenasai.gateway.protocol import ErrorMessage

        msg = ErrorMessage(message="Something went wrong", code="ERR_001")
        assert msg.type == MessageType.ERROR
        assert msg.message == "Something went wrong"
        assert msg.code == "ERR_001"

    def test_pong_message(self):
        """测试心跳响应消息"""
        from minenasai.gateway.protocol import PongMessage

        msg = PongMessage()
        assert msg.type == MessageType.PONG