qianbian/tests/test_mcp_tools.py

"""mcp_tools 测试：图源解码 + 工具契约（内存 MCP 会话端到端）。"""
import base64
import io
import json

import pytest
from bleak.exc import BleakError
from mcp.server.fastmcp import FastMCP
from mcp.shared.memory import create_connected_server_and_client_session
from PIL import Image

from ppclock.client import PPClient
from ppclock.mcp_tools import decode_image_source, register_tools
from ppclock.transports.base import TransportError
from ppclock.transports.local import TransportError as LocalTransportError


class FakeTransport:
    def __init__(self):
        self.address = "AA:BB:CC:DD:EE:FF"
        self.epd_writes = []
        self.rxtx_writes = []

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc):
        return False

    async def write_epd(self, data, response=True):
        self.epd_writes.append(data)

    async def write_rxtx(self, data, response=True):
        self.rxtx_writes.append(data)

    async def read_rxtx(self):
        return b"\x00"

    async def request_device_id(self, timeout=8.0):
        return "81233F3C267112"

    async def run_ota(self, image, on_progress=None):
        return {}

    async def delay(self, seconds):
        pass


class FakeManager:
    """模拟 DeviceManager：直接对 fake client 执行 op，或抛预置错误。"""

    def __init__(self, error=None, status_error=None):
        self.client = PPClient(FakeTransport())
        self.error = error
        self.status_error = status_error
        self.ops = 0

    async def run(self, op):
        self.ops += 1
        if self.error is not None:
            raise self.error
        return await op(self.client)

    async def status(self):
        if self.status_error is not None:
            raise self.status_error
        return {"connected": True, "mac": "AA:BB:CC:DD:EE:FF",
                "idle_seconds": 0.0, "connect_count": 1,
                "connect_timeout": 90.0, "idle_timeout": 300.0,
                "device_id": "81233F3C267112"}


def make_session_coro(manager):
    mcp = FastMCP("ppclock-test")
    register_tools(mcp, manager)
    return mcp


async def call(mcp, tool, args=None):
    async with create_connected_server_and_client_session(mcp._mcp_server) as s:
        result = await s.call_tool(tool, args or {})
        assert not result.isError, result.content
        return json.loads(result.content[0].text)


def png_b64() -> str:
    buf = io.BytesIO()
    Image.new("RGB", (400, 300), "white").save(buf, "PNG")
    return base64.b64encode(buf.getvalue()).decode()


class TestDecodeImageSource:
    def test_path(self, tmp_path):
        p = tmp_path / "a.png"
        Image.new("RGB", (10, 10)).save(p)
        assert decode_image_source(str(p)).size == (10, 10)

    def test_base64(self):
        assert decode_image_source(png_b64()).size == (400, 300)

    def test_data_uri(self):
        src = "data:image/png;base64," + png_b64()
        assert decode_image_source(src).size == (400, 300)

    def test_garbage(self):
        with pytest.raises(ValueError):
            decode_image_source("!!!not-base64!!!")

    def test_base64_not_image(self):
        src = base64.b64encode(b"not an image").decode()
        with pytest.raises(ValueError):
            decode_image_source(src)


class TestTools:
    @pytest.mark.asyncio
    async def test_set_time_ok(self):
        r = await call(make_session_coro(FakeManager()), "set_time", {"tz": 8.0})
        assert r["ok"] is True and r["tool"] == "set_time"

    @pytest.mark.asyncio
    async def test_upload_image_base64(self):
        r = await call(make_session_coro(FakeManager()), "upload_image",
                       {"source": png_b64(), "slot": 0})
        assert r["ok"] is True
        assert r["data"]["bytes_bw"] == 15000
        assert r["data"]["slot"] == 0

    @pytest.mark.asyncio
    async def test_upload_image_not_image(self):
        src = base64.b64encode(b"not an image").decode()
        r = await call(make_session_coro(FakeManager()), "upload_image",
                       {"source": src, "slot": 0})
        assert r["ok"] is False
        assert r["error"]["code"] == "INVALID_PARAM"

    @pytest.mark.asyncio
    async def test_set_mode_invalid_param(self):
        r = await call(make_session_coro(FakeManager()), "set_mode",
                       {"mode": "nonsense"})
        assert r["ok"] is False
        assert r["error"]["code"] == "INVALID_PARAM"

    @pytest.mark.asyncio
    async def test_transport_error_mapped(self):
        m = FakeManager(error=TransportError("Connection timed out"))
        r = await call(make_session_coro(m), "set_time", {})
        assert r["ok"] is False
        assert r["error"]["code"] == "CONNECT_TIMEOUT"
        assert "hint" in r["error"]

    @pytest.mark.asyncio
    async def test_local_not_found_mapped(self):
        """真实 BLE 扫描落空：local.TransportError（中文）→ DEVICE_NOT_FOUND。"""
        m = FakeManager(error=LocalTransportError("未找到名称前缀 'NRF-' 的设备"))
        r = await call(make_session_coro(m), "set_time", {})
        assert r["ok"] is False
        assert r["error"]["code"] == "DEVICE_NOT_FOUND"
        assert "hint" in r["error"]

    @pytest.mark.asyncio
    async def test_local_connect_failure_mapped(self):
        """已知 MAC 守候未连上（bleak 报 not found 被包装成"连接失败"）→ CONNECT_TIMEOUT。"""
        m = FakeManager(error=LocalTransportError(
            "连接失败: Device with address AA:BB:CC:DD:EE:FF was not found."))
        r = await call(make_session_coro(m), "set_time", {})
        assert r["ok"] is False
        assert r["error"]["code"] == "CONNECT_TIMEOUT"
        assert "hint" in r["error"]

    @pytest.mark.asyncio
    async def test_raw_bleak_error_mapped(self):
        """写路径 bleak 裸抛 BleakError（不包装）→ BLE_ERROR。"""
        m = FakeManager(error=BleakError("Not connected"))
        r = await call(make_session_coro(m), "set_time", {})
        assert r["ok"] is False
        assert r["error"]["code"] == "BLE_ERROR"

    @pytest.mark.asyncio
    async def test_device_status_transport_error_mapped(self):
        """device_status 的传输层错误与其它工具同映射，不兜底 INTERNAL。"""
        m = FakeManager(status_error=LocalTransportError("等待设备 ID 应答超时"))
        r = await call(make_session_coro(m), "device_status", {})
        assert r["ok"] is False
        assert r["error"]["code"] == "CONNECT_TIMEOUT"
        assert "hint" in r["error"]

    @pytest.mark.asyncio
    async def test_device_status(self):
        r = await call(make_session_coro(FakeManager()), "device_status", {})
        assert r["ok"] is True
        assert r["data"]["device_id"] == "81233F3C267112"

    @pytest.mark.asyncio
    async def test_raw_send(self):
        r = await call(make_session_coro(FakeManager()), "raw_send",
                       {"data_hex": "e201", "channel": "alt"})
        assert r["ok"] is True

    @pytest.mark.asyncio
    async def test_raw_send_bad_hex(self):
        r = await call(make_session_coro(FakeManager()), "raw_send",
                       {"data_hex": "zz", "channel": "alt"})
        assert r["ok"] is False
        assert r["error"]["code"] == "INVALID_PARAM"

    @pytest.mark.asyncio
    async def test_render_template(self):
        r = await call(make_session_coro(FakeManager()), "render_template",
                       {"name": "custom", "payload": {"text": "你好"}, "slot": 0})
        assert r["ok"] is True
        assert r["data"]["template"] == "custom"

    @pytest.mark.asyncio
    async def test_countdown(self):
        r = await call(make_session_coro(FakeManager()), "countdown",
                       {"date": "2026-12-31", "mode": "clock", "prefix": "目标"})
        assert r["ok"] is True

    @pytest.mark.asyncio
    async def test_scan_devices(self, monkeypatch):
        async def fake_scan(uri="local", timeout=5.0):
            return [{"name": "NRF-5DBF28", "address": "18:BC:5A:5D:BF:28",
                     "rssi": -60}]
        monkeypatch.setattr(PPClient, "scan_devices", staticmethod(fake_scan))
        r = await call(make_session_coro(FakeManager()), "scan_devices", {})
        assert r["ok"] is True
        assert r["data"]["devices"][0]["address"] == "18:BC:5A:5D:BF:28"