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}} 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): app = PublicGatewayApp(session_store=FakeSessionStore(), api_client=FakeApiClient(), auth_client=None) tool_names = [tool['name'] for tool in app.list_tools()] self.assertNotIn('bind_auth_code', tool_names) self.assertIn('query_order', tool_names) self.assertIn('query_track', 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()