mcp/tests/test_mcp_protocol.py

import io
import json
import unittest

from app import GatewayApp
from services.token_store import InMemoryTokenStore
from mcp_protocol import McpProtocolHandler


class DummyApiClient:
    def __init__(self):
        self.calls = []

    def call_tool(self, tool_code, route_path, payload, request_id):
        self.calls.append(
            {
                'tool_code': tool_code,
                'route_path': route_path,
                'payload': payload,
                'request_id': request_id,
            }
        )
        return {
            'code': 'MCP_0000',
            'msg': 'success',
            'data': {
                'summary': 'matched 1 order',
                'records': [
                    {
                        'order_no': 'SO20260706001',
                    }
                ],
                'tips': ['scoped by employee permissions'],
            },
            'meta': {
                'request_id': request_id,
            },
        }

class FullColumnsApiClient(DummyApiClient):
    def call_tool(self, tool_code, route_path, payload, request_id):
        response = super().call_tool(tool_code, route_path, payload, request_id)
        columns = [
            ('order_number', '订单号'),
            ('reference_number', '客户参考号'),
            ('status_txt_name', '状态'),
            ('check_status_txt_name', '是否已查验'),
            ('customer_name', '客户名称'),
            ('customer_account_type_name', '客户属性'),
            ('inbound_date', '入库时间'),
            ('wo_num', '未完成工单'),
            ('product_name', '物流产品'),
            ('inbound_pieces', '件数'),
            ('inbound_volume', '体积(CBM)'),
            ('inbound_weight', '重量(KG)'),
            ('pro_cn_name', '品名'),
            ('export_declaration_type', '报关方式'),
            ('merge_declare_number', '合并报关单号'),
            ('delivery_address', '派送地址'),
            ('container_code', '柜号'),
            ('out_status_txt', '排舱单状态'),
            ('hinge_of_destination', '目的港'),
            ('etd', 'ETD'),
            ('atd', 'ATD'),
            ('eta', 'ETA'),
            ('ata', 'ATA'),
            ('release_time', '清关放行时间'),
            ('oversea_inbound_date', '海外入库时间'),
            ('appt_time', 'APPT时间'),
            ('est_loading_time', '预计装柜时间'),
            ('pickup_time', '海外提柜时间'),
            ('delivery_way_title', '派送方式'),
            ('tracking_number', '快递单号'),
            ('shipment_id', 'SHIPMENT ID'),
            ('goods_attribute', '商品属性'),
            ('sku', 'SKU'),
            ('sales_user', '商务经理'),
            ('service_user', '客户经理'),
            ('department_name', '事业部'),
            ('remark', '订单备注'),
            ('importer_name', '进口商'),
            ('warehouse_name', '交货仓库'),
            ('paid_status_name', '付款状态'),
        ]
        response['data']['columns'] = [
            {'key': key, 'name': name, 'check': True} for key, name in columns
        ]
        response['data']['records'] = [
            {key: '{0}-value'.format(key) for key, _name in columns}
        ]
        return response


class McpProtocolTest(unittest.TestCase):
    def build_handler(self):
        token_store = InMemoryTokenStore(refresh_skew_seconds=60)
        token_store.save('MT_demo', '2099-01-01T00:00:00')
        app = GatewayApp(
            auth_client=None,
            api_client=DummyApiClient(),
            token_store=token_store,
        )
        return McpProtocolHandler(app)

    def test_initialize_returns_server_capabilities(self):
        handler = self.build_handler()

        response = handler.handle_request(
            {
                'jsonrpc': '2.0',
                'id': 1,
                'method': 'initialize',
                'params': {
                    'protocolVersion': '2025-06-18',
                    'capabilities': {},
                    'clientInfo': {'name': 'workbuddy', 'version': '1.0.0'},
                },
            }
        )

        self.assertEqual('2.0', response['jsonrpc'])
        self.assertEqual(1, response['id'])
        self.assertEqual('2025-06-18', response['result']['protocolVersion'])
        self.assertIn('tools', response['result']['capabilities'])
        self.assertEqual('fms-mcp-gateway', response['result']['serverInfo']['name'])

    def test_initialized_notification_does_not_emit_response(self):
        handler = self.build_handler()

        response = handler.handle_message(
            {
                'jsonrpc': '2.0',
                'method': 'notifications/initialized',
            }
        )

        self.assertIsNone(response)

    def test_tools_list_returns_registered_tools(self):
        handler = self.build_handler()

        response = handler.handle_request(
            {
                'jsonrpc': '2.0',
                'id': 2,
                'method': 'tools/list',
                'params': {},
            }
        )

        self.assertEqual('2.0', response['jsonrpc'])
        self.assertEqual(2, response['id'])
        by_name = {tool['name']: tool for tool in response['result']['tools']}
        self.assertIn('query_order', by_name)
        self.assertIn('bind_auth_code', by_name)
        self.assertIn('inputSchema', by_name['query_order'])

    def test_tools_call_wraps_gateway_result_as_structured_content(self):
        handler = self.build_handler()

        response = handler.handle_request(
            {
                'jsonrpc': '2.0',
                'id': 3,
                'method': 'tools/call',
                'params': {
                    'name': 'query_order',
                    'arguments': {
                        'keyword': 'SO20260706001',
                        'page': 1,
                        'limit': 20,
                    },
                },
            }
        )

        self.assertEqual('2.0', response['jsonrpc'])
        self.assertEqual(3, response['id'])
        self.assertFalse(response['result']['isError'])
        self.assertEqual('matched 1 order', response['result']['structuredContent']['summary'])
        self.assertEqual('SO20260706001', response['result']['structuredContent']['records'][0]['order_no'])
        self.assertEqual('text', response['result']['content'][0]['type'])
        self.assertIn('matched 1 order', response['result']['content'][0]['text'])

    def test_tools_call_renders_all_query_order_columns_in_text_content(self):
        token_store = InMemoryTokenStore(refresh_skew_seconds=60)
        token_store.save('MT_demo', '2099-01-01T00:00:00')
        app = GatewayApp(
            auth_client=None,
            api_client=FullColumnsApiClient(),
            token_store=token_store,
        )
        handler = McpProtocolHandler(app)

        response = handler.handle_request(
            {
                'jsonrpc': '2.0',
                'id': 4,
                'method': 'tools/call',
                'params': {
                    'name': 'query_order',
                    'arguments': {
                        'keyword': 'SO20260706001',
                    },
                },
            }
        )

        text = response['result']['content'][0]['text']
        self.assertIn('表头共 40 列', text)
        self.assertIn('1. 订单号 (order_number)', text)
        self.assertIn('40. 付款状态 (paid_status_name)', text)
        self.assertIn('- 付款状态: paid_status_name-value', text)
        self.assertEqual(40, len(response['result']['structuredContent']['columns']))
    def test_run_stdio_writes_only_request_responses(self):
        handler = self.build_handler()
        stdin = io.StringIO(
            json.dumps(
                {
                    'jsonrpc': '2.0',
                    'id': 1,
                    'method': 'initialize',
                    'params': {
                        'protocolVersion': '2025-06-18',
                        'capabilities': {},
                        'clientInfo': {'name': 'workbuddy', 'version': '1.0.0'},
                    },
                }
            )
            + '\n'
            + json.dumps(
                {
                    'jsonrpc': '2.0',
                    'method': 'notifications/initialized',
                }
            )
            + '\n'
        )
        stdout = io.StringIO()

        handler.run_stdio(stdin=stdin, stdout=stdout)

        lines = [line for line in stdout.getvalue().splitlines() if line.strip()]
        self.assertEqual(1, len(lines))
        response = json.loads(lines[0])
        self.assertEqual(1, response['id'])
        self.assertEqual('2025-06-18', response['result']['protocolVersion'])


if __name__ == '__main__':
    unittest.main()