mcp/tests/test_public_server_stream.py

import json
import unittest
from http.client import HTTPConnection
from threading import Thread
from time import sleep
from unittest.mock import MagicMock

from constants import DEVICE_INVALID_MESSAGE
from public_gateway import PublicGatewayApp
from public_server import serve_public
from utils.security import generate_gateway_session_id


class TestPublicServerStream(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.mock_session_store = MagicMock()
        cls.mock_api_client = MagicMock()
        cls.mock_auth_client = MagicMock()
        cls.gateway_app = PublicGatewayApp(
            session_store=cls.mock_session_store,
            api_client=cls.mock_api_client,
            auth_client=cls.mock_auth_client,
        )
        cls.server_thread = Thread(
            target=serve_public,
            args=(cls.gateway_app,),
            kwargs={'host': '127.0.0.1', 'port': 18766, 'enable_rate_limit': False},
            daemon=True,
        )
        cls.server_thread.start()
        sleep(0.5)

    def setUp(self):
        self.mock_session_store.reset_mock()
        self.mock_api_client.reset_mock()
        self.mock_auth_client.reset_mock()
        self.mock_session_store.get.return_value = None
        self.mock_api_client.list_enabled_tools.return_value = {
            'code': 'MCP_0000',
            'data': {'tool_codes': ['query_order']},
        }
        self.mock_api_client.call_tool.return_value = {
            'code': '0',
            'msg': 'success',
            'data': {'order_no': 'ABC123'},
        }

    def _exchange(self, method, path, body=None, headers=None):
        conn = HTTPConnection('127.0.0.1', 18766, timeout=5)
        try:
            request_headers = {}
            if headers:
                request_headers.update(headers)
            conn.request(method, path, body, request_headers)
            response = conn.getresponse()
            return response.status, response.getheader('Allow'), response.read()
        finally:
            conn.close()

    def _post_json(self, path, payload, headers=None):
        raw = json.dumps(payload).encode('utf-8')
        status, allow, body = self._exchange('POST', path, raw, headers)
        parsed = json.loads(body.decode('utf-8')) if body else None
        return status, allow, parsed

    def test_existing_mcp_initialize_keeps_protocol_version(self):
        status, _allow, body = self._post_json('/mcp', {
            'jsonrpc': '2.0',
            'id': 1,
            'method': 'initialize',
            'params': {'protocolVersion': '2025-03-26'},
        })
        self.assertEqual(200, status)
        self.assertEqual('2025-06-18', body['result']['protocolVersion'])

    def test_existing_mcp_notification_stays_method_not_found(self):
        status, _allow, body = self._post_json('/mcp', {
            'jsonrpc': '2.0',
            'method': 'notifications/initialized',
        })
        self.assertEqual(200, status)
        self.assertEqual(-32601, body['error']['code'])
        self.assertIn('Method not found', body['error']['message'])

    def test_existing_get_mcp_stays_404(self):
        status, _allow, body = self._exchange('GET', '/mcp')
        self.assertEqual(404, status)
        self.assertEqual(b'', body)

    def test_health_stays_200(self):
        status, _allow, body = self._exchange('GET', '/health')
        self.assertEqual(200, status)
        self.assertEqual({'ok': True}, json.loads(body.decode('utf-8')))

    def test_delete_mcp_stays_unimplemented(self):
        status, _allow, _body = self._exchange('DELETE', '/mcp')
        self.assertEqual(501, status)

    def test_stream_lookalike_paths_stay_404(self):
        status, _allow, _body = self._post_json('/mcp/stream-extra', {
            'jsonrpc': '2.0',
            'id': 1,
            'method': 'ping',
        })
        self.assertEqual(404, status)
        status, _allow, _body = self._post_json('/mcp/stream?x=1', {
            'jsonrpc': '2.0',
            'id': 1,
            'method': 'ping',
        })
        self.assertEqual(404, status)

    def test_get_stream_returns_405(self):
        status, allow, body = self._exchange('GET', '/mcp/stream')
        self.assertEqual(405, status)
        self.assertEqual('POST', allow)
        self.assertEqual(b'', body)

    def test_delete_stream_returns_405(self):
        status, allow, body = self._exchange('DELETE', '/mcp/stream')
        self.assertEqual(405, status)
        self.assertEqual('POST', allow)
        self.assertEqual(b'', body)

    def test_stream_notification_returns_202(self):
        status, allow, body = self._post_json('/mcp/stream', {
            'jsonrpc': '2.0',
            'method': 'notifications/initialized',
        })
        self.assertEqual(202, status)
        self.assertIsNone(allow)
        self.assertIsNone(body)

    def test_stream_invalid_json_returns_400(self):
        status, _allow, body = self._exchange(
            'POST',
            '/mcp/stream',
            b'this is not json',
        )
        self.assertEqual(400, status)
        parsed = json.loads(body.decode('utf-8'))
        self.assertEqual(-32700, parsed['error']['code'])
        self.assertEqual('Parse error', parsed['error']['message'])

    def test_stream_empty_body_returns_parse_error(self):
        status, _allow, body = self._exchange('POST', '/mcp/stream', b'')
        self.assertEqual(400, status)
        parsed = json.loads(body.decode('utf-8'))
        self.assertEqual(-32700, parsed['error']['code'])

    def test_stream_non_object_returns_invalid_request(self):
        status, _allow, body = self._exchange('POST', '/mcp/stream', b'[1]')
        self.assertEqual(400, status)
        parsed = json.loads(body.decode('utf-8'))
        self.assertEqual(-32600, parsed['error']['code'])
        self.assertEqual('Invalid Request', parsed['error']['message'])

    def test_stream_ping(self):
        status, _allow, body = self._post_json('/mcp/stream', {
            'jsonrpc': '2.0',
            'id': 7,
            'method': 'ping',
        })
        self.assertEqual(200, status)
        self.assertEqual({'jsonrpc': '2.0', 'id': 7, 'result': {}}, body)

    def test_stream_initialize_echoes_known_versions(self):
        for version in ('2025-03-26', '2025-06-18'):
            status, _allow, body = self._post_json('/mcp/stream', {
                'jsonrpc': '2.0',
                'id': 1,
                'method': 'initialize',
                'params': {'protocolVersion': version},
            })
            self.assertEqual(200, status)
            self.assertEqual(version, body['result']['protocolVersion'])
            self.assertEqual('fms-mcp-gateway', body['result']['serverInfo']['name'])
            self.assertEqual(False, body['result']['capabilities']['tools']['listChanged'])

    def test_stream_initialize_unknown_version_falls_back(self):
        status, _allow, body = self._post_json('/mcp/stream', {
            'jsonrpc': '2.0',
            'id': 1,
            'method': 'initialize',
            'params': {'protocolVersion': '2099-01-01'},
        })
        self.assertEqual(200, status)
        self.assertEqual('2025-03-26', body['result']['protocolVersion'])

    def test_stream_initialize_without_params_falls_back(self):
        status, _allow, body = self._post_json('/mcp/stream', {
            'jsonrpc': '2.0',
            'id': 1,
            'method': 'initialize',
        })
        self.assertEqual(200, status)
        self.assertEqual('2025-03-26', body['result']['protocolVersion'])

    def test_stream_initialize_with_empty_params_falls_back(self):
        status, _allow, body = self._post_json('/mcp/stream', {
            'jsonrpc': '2.0',
            'id': 1,
            'method': 'initialize',
            'params': {},
        })
        self.assertEqual(200, status)
        self.assertEqual('2025-03-26', body['result']['protocolVersion'])

    def test_stream_initialize_with_non_object_params_falls_back(self):
        status, _allow, body = self._post_json('/mcp/stream', {
            'jsonrpc': '2.0',
            'id': 1,
            'method': 'initialize',
            'params': ['2025-06-18'],
        })
        self.assertEqual(200, status)
        self.assertEqual('2025-03-26', body['result']['protocolVersion'])

    def test_stream_request_without_method_uses_existing_error(self):
        status, _allow, body = self._post_json('/mcp/stream', {
            'jsonrpc': '2.0',
            'id': 3,
        })
        self.assertEqual(200, status)
        self.assertEqual(-32601, body['error']['code'])

    def test_stream_tools_list_without_session_matches_existing_error(self):
        status, _allow, body = self._post_json('/mcp/stream', {
            'jsonrpc': '2.0',
            'id': 1,
            'method': 'tools/list',
        })
        self.assertEqual(200, status)
        self.assertEqual(-32001, body['error']['code'])
        self.assertIn(DEVICE_INVALID_MESSAGE, body['error']['message'])

    def test_stream_tools_call_uses_existing_session(self):
        session_id = generate_gateway_session_id()
        self.mock_session_store.get.return_value = {
            'mcp_token': 'MT_valid_token',
            'admin_id': 456,
            'company_id': 200,
        }
        status, _allow, body = self._post_json(
            '/mcp/stream',
            {
                'jsonrpc': '2.0',
                'id': 1,
                'method': 'tools/call',
                'params': {
                    'name': 'query_order',
                    'arguments': {'order_no': 'ABC123'},
                },
            },
            {'X-Gateway-Session': session_id},
        )
        self.assertEqual(200, status)
        self.assertIn('result', body)
        self.assertFalse(body['result']['isError'])
        self.mock_api_client.call_tool.assert_called_once()


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