mcp/tests/test_public_gateway.py

import unittest

from public_gateway import PublicGatewayApp


class FakeSessionStore:
    def __init__(self):
        self.sessions = {}
        self.touched = []

    def get(self, gateway_session_id):
        return self.sessions.get(gateway_session_id)

    def touch_session(self, gateway_session_id):
        self.touched.append(gateway_session_id)


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

    def call_tool(self, token, tool_code, route_path, payload, request_id, client_ip=''):
        self.calls.append((token, tool_code, route_path, payload, request_id, client_ip))
        return {'code': 'MCP_0000', 'data': {'token_used': token}}

    def list_enabled_tools(self, token, request_id=''):
        return {
            'code': 'MCP_0000',
            'data': {
                'tool_codes': [
                    'query_order',
                    'query_track',
                    'query_order_exact',
                    'list_order_filter_options',
                ],
            },
        }


class PublicGatewayAppTest(unittest.TestCase):
    def test_two_employees_use_isolated_tokens(self):
        store = FakeSessionStore()
        store.sessions['GWS_A'] = {'mcp_token': 'MT_A'}
        store.sessions['GWS_B'] = {'mcp_token': 'MT_B'}
        api_client = FakeApiClient()
        app = PublicGatewayApp(session_store=store, api_client=api_client, auth_client=None)

        result_a = app.call_tool('GWS_A', 'query_order', {'keyword': 'A'}, request_id='rq_a')
        result_b = app.call_tool('GWS_B', 'query_order', {'keyword': 'B'}, request_id='rq_b')

        self.assertEqual('MT_A', result_a['data']['token_used'])
        self.assertEqual('MT_B', result_b['data']['token_used'])
        self.assertEqual('MT_A', api_client.calls[0][0])
        self.assertEqual('MT_B', api_client.calls[1][0])

    def test_public_tools_do_not_include_bind_auth_code(self):
        store = FakeSessionStore()
        store.sessions['GWS_A'] = {'mcp_token': 'MT_A'}
        app = PublicGatewayApp(session_store=store, api_client=FakeApiClient(), auth_client=None)

        tool_names = [tool['name'] for tool in app.list_tools('GWS_A')]

        self.assertNotIn('bind_auth_code', tool_names)
        self.assertIn('query_order', tool_names)
        self.assertIn('query_track', tool_names)
        self.assertIn('query_order_exact', tool_names)
        self.assertIn('list_order_filter_options', tool_names)

    def test_missing_session_returns_human_device_message(self):
        app = PublicGatewayApp(session_store=FakeSessionStore(), api_client=FakeApiClient(), auth_client=None)

        with self.assertRaises(RuntimeError) as error:
            app.call_tool('GWS_missing', 'query_order', {'keyword': 'A'}, request_id='rq_missing')

        self.assertIn('这台设备的 Workbuddy 配置已失效，请重新生成配置', str(error.exception))

    def test_tool_call_forwards_client_ip_to_api_client(self):
        store = FakeSessionStore()
        store.sessions['GWS_A'] = {'mcp_token': 'MT_A'}
        api_client = FakeApiClient()
        app = PublicGatewayApp(session_store=store, api_client=api_client, auth_client=None)

        app.call_tool('GWS_A', 'query_order', {'keyword': 'A'}, request_id='rq_a', client_ip='203.0.113.9')

        self.assertEqual('203.0.113.9', api_client.calls[0][5])

    def test_successful_tool_call_touches_gateway_session(self):
        store = FakeSessionStore()
        store.sessions['GWS_A'] = {'mcp_token': 'MT_A'}
        app = PublicGatewayApp(session_store=store, api_client=FakeApiClient(), auth_client=None)

        app.call_tool('GWS_A', 'query_order', {'keyword': 'A'}, request_id='rq_a')

        self.assertEqual(['GWS_A'], store.touched)


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