mcp/tests/test_public_gateway.py

import unittest

from public_gateway import PublicGatewayApp
from services.diagnostic_event import RequestDiagnosticEmitter


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_missing_redis_session_emits_gateway_session_failure(self):
        reporter = RecordingReporter()
        emitter = RequestDiagnosticEmitter(
            reporter,
            'rq_missing',
            defer_until_identity=True,
        )
        app = PublicGatewayApp(FakeSessionStore(), FakeApiClient())

        with self.assertRaises(RuntimeError):
            app.call_tool(
                'GWS_missing',
                'query_order',
                request_id='rq_missing',
                diagnostic_emitter=emitter,
            )
        emitter.flush()

        event = reporter.events[-1]
        self.assertEqual('gateway_session', event['stage'])
        self.assertEqual('failed', event['status'])
        self.assertEqual('GATEWAY_SESSION_NOT_FOUND', event['event_code'])
        self.assertIn('session_hash', event)

    def test_enabled_tool_lookup_failure_emits_backend_failure(self):
        class FailingListClient(FakeApiClient):
            def list_enabled_tools(self, token, request_id=''):
                raise OSError('registry unavailable secret')

        store = FakeSessionStore()
        store.sessions['GWS_A'] = {
            'mcp_token': 'MT_A',
            'admin_id': 88,
            'company_id': 1002,
        }
        reporter = RecordingReporter()
        emitter = RequestDiagnosticEmitter(
            reporter,
            'rq_lookup',
            defer_until_identity=True,
        )
        emitter.emit(
            stage='request_ingress',
            status='started',
            event_code='REQUEST_RECEIVED',
            context={'transport': 'http'},
        )
        app = PublicGatewayApp(store, FailingListClient())

        with self.assertRaises(OSError):
            app.call_tool(
                'GWS_A',
                'query_order',
                request_id='rq_lookup',
                diagnostic_emitter=emitter,
            )

        self.assertTrue(all(
            event['company_id'] == 1002 for event in reporter.events
        ))
        failed = reporter.events[-1]
        self.assertEqual('backend_call', failed['stage'])
        self.assertEqual('failed', failed['status'])
        self.assertEqual('ENABLED_TOOL_LOOKUP_FAILED', failed['event_code'])

    def test_unknown_and_disabled_tools_emit_backend_failures(self):
        store = FakeSessionStore()
        store.sessions['GWS_A'] = {
            'mcp_token': 'MT_A',
            'admin_id': 88,
            'company_id': 1002,
        }
        app = PublicGatewayApp(store, FakeApiClient())

        cases = (
            ('not_registered', 'TOOL_NOT_REGISTERED', KeyError),
            ('query_outbound_detail', 'TOOL_DISABLED', RuntimeError),
        )
        for tool_name, event_code, exception_class in cases:
            with self.subTest(tool_name=tool_name):
                reporter = RecordingReporter()
                emitter = RequestDiagnosticEmitter(
                    reporter,
                    'rq_tool_check',
                    defer_until_identity=True,
                )
                with self.assertRaises(exception_class):
                    app.call_tool(
                        'GWS_A',
                        tool_name,
                        request_id='rq_tool_check',
                        diagnostic_emitter=emitter,
                    )
                self.assertEqual(event_code, reporter.events[-1]['event_code'])
                self.assertEqual('failed', reporter.events[-1]['status'])

    def test_enabled_tool_lookup_failure_without_emitter_keeps_exception(self):
        class FailingListClient(FakeApiClient):
            def list_enabled_tools(self, token, request_id=''):
                raise OSError('registry unavailable')

        store = FakeSessionStore()
        store.sessions['GWS_A'] = {'mcp_token': 'MT_A'}
        app = PublicGatewayApp(store, FailingListClient())

        with self.assertRaises(OSError):
            app.call_tool('GWS_A', 'query_order', request_id='rq_lookup')
    def test_diagnostic_events_include_session_identity_and_backend_result(self):
        store = FakeSessionStore()
        store.sessions['GWS_A'] = {
            'mcp_token': 'MT_A',
            'admin_id': 88,
            'company_id': 1002,
        }
        reporter = RecordingReporter()
        emitter = RequestDiagnosticEmitter(reporter, 'rq_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',
            diagnostic_emitter=emitter,
        )

        self.assertEqual(
            ['gateway_session', 'backend_call', 'backend_call'],
            [event['stage'] for event in reporter.events],
        )
        self.assertEqual(
            ['succeeded', 'started', 'succeeded'],
            [event['status'] for event in reporter.events],
        )
        self.assertTrue(all(event['company_id'] == 1002 for event in reporter.events))
        self.assertNotIn('GWS_A', str(reporter.events))
        self.assertNotIn('MT_A', str(reporter.events))

    def test_backend_failure_emits_failed_event_without_changing_exception(self):
        class FailingApiClient(FakeApiClient):
            def call_tool(self, *args, **kwargs):
                raise OSError('backend secret unavailable')

        store = FakeSessionStore()
        store.sessions['GWS_A'] = {
            'mcp_token': 'MT_A',
            'admin_id': 88,
            'company_id': 1002,
        }
        reporter = RecordingReporter()
        emitter = RequestDiagnosticEmitter(reporter, 'rq_a')
        app = PublicGatewayApp(store, FailingApiClient(), auth_client=None)

        with self.assertRaisesRegex(OSError, 'backend secret unavailable'):
            app.call_tool(
                'GWS_A',
                'query_order',
                request_id='rq_a',
                diagnostic_emitter=emitter,
            )

        failed = reporter.events[-1]
        self.assertEqual('backend_call', failed['stage'])
        self.assertEqual('failed', failed['status'])
        self.assertEqual('UNEXPECTED_EXCEPTION', failed['event_code'])
        self.assertNotIn('backend secret', str(failed))

    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)


class RecordingReporter:
    def __init__(self):
        self.events = []

    def report(self, event):
        self.events.append(event)
        return True


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