unit_test_handlers.py•6.5 kB
import json
import unittest
import mcp.types as types
from unittest.mock import AsyncMock, patch
from codelogic_mcp_server.handlers import handle_call_tool, extract_relationships
class TestHandleCallTool(unittest.TestCase):
@patch('codelogic_mcp_server.handlers.server.request_context')
@patch('codelogic_mcp_server.handlers.get_mv_id')
@patch('codelogic_mcp_server.handlers.get_method_entity')
@patch('codelogic_mcp_server.handlers.get_impact')
@patch('codelogic_mcp_server.handlers.find_node_by_id')
async def test_handle_call_tool_method(self, mock_find_node_by_id, mock_get_impact, mock_get_method_entity, mock_get_mv_id, mock_request_context):
# Setup mocks
mock_request_context.session.send_log_message = AsyncMock()
mock_get_mv_id.return_value = 'mv_id'
mock_get_method_entity.return_value = [{'properties': {'id': 'node_id'}, 'name': 'node_name'}]
mock_get_impact.return_value = json.dumps({
'data': {
'relationships': [{'startId': 'start_id', 'endId': 'end_id', 'type': 'type'}],
'nodes': [{'id': 'start_id', 'name': 'start_name', 'primaryLabel': 'label'}, {'id': 'end_id', 'name': 'end_name', 'primaryLabel': 'label'}]
}
})
mock_find_node_by_id.side_effect = lambda nodes, id: next(node for node in nodes if node['id'] == id)
# Call the function
result = await handle_call_tool('codelogic-method-impact', {'method': 'method_name'})
# Assertions
mock_request_context.session.send_log_message.assert_any_call(level="info", data="Materialized view ID: mv_id")
mock_request_context.session.send_log_message.assert_any_call(level="info", data="Node ID: node_id, Node Name: node_name")
mock_request_context.session.send_log_message.assert_any_call(level="info", data="Impact analysis completed for method_name")
self.assertEqual(result, [types.TextContent(type="text", text="Impact analysis for method: method_name\n- start_name (type) -> end_name (label)")])
@patch('codelogic_mcp_server.handlers.server.request_context')
@patch('codelogic_mcp_server.handlers.get_mv_id')
@patch('codelogic_mcp_server.handlers.get_method_entity')
@patch('codelogic_mcp_server.handlers.get_impact')
@patch('codelogic_mcp_server.handlers.find_node_by_id')
async def test_handle_call_tool_function(self, mock_find_node_by_id, mock_get_impact, mock_get_method_entity, mock_get_mv_id, mock_request_context):
# Setup mocks
mock_request_context.session.send_log_message = AsyncMock()
mock_get_mv_id.return_value = 'mv_id'
mock_get_method_entity.return_value = [{'properties': {'id': 'node_id'}, 'name': 'node_name'}]
mock_get_impact.return_value = json.dumps({
'data': {
'relationships': [{'startId': 'start_id', 'endId': 'end_id', 'type': 'type'}],
'nodes': [{'id': 'start_id', 'name': 'start_name', 'primaryLabel': 'label'}, {'id': 'end_id', 'name': 'end_name', 'primaryLabel': 'label'}]
}
})
mock_find_node_by_id.side_effect = lambda nodes, id: next(node for node in nodes if node['id'] == id)
# Call the function
result = await handle_call_tool('codelogic-method-impact', {'function': 'function_name'})
# Assertions
mock_request_context.session.send_log_message.assert_any_call(level="info", data="Materialized view ID: mv_id")
mock_request_context.session.send_log_message.assert_any_call(level="info", data="Node ID: node_id, Node Name: node_name")
mock_request_context.session.send_log_message.assert_any_call(level="info", data="Impact analysis completed for function_name")
self.assertEqual(result, [types.TextContent(type="text", text="Impact analysis for function: function_name\n- start_name (type) -> end_name (label)")])
async def test_handle_call_tool_unknown_tool(self):
with self.assertRaises(ValueError) as context:
await handle_call_tool('unknown-tool', {'method': 'method_name'})
self.assertEqual(str(context.exception), "Unknown tool: unknown-tool")
async def test_handle_call_tool_missing_arguments(self):
with self.assertRaises(ValueError) as context:
await handle_call_tool('codelogic-method-impact', None)
self.assertEqual(str(context.exception), "Missing arguments")
async def test_handle_call_tool_missing_method_function(self):
with self.assertRaises(ValueError) as context:
await handle_call_tool('codelogic-method-impact', {})
self.assertEqual(str(context.exception), "At least one of method or function must be provided")
@patch('codelogic_mcp_server.handlers.server.request_context')
async def test_handle_call_tool_missing_request_context(self, mock_request_context):
mock_request_context.session = None
with self.assertRaises(LookupError) as context:
await handle_call_tool('codelogic-method-impact', {'method': 'method_name'})
self.assertEqual(str(context.exception), "Request context is not set")
class TestExtractRelationships(unittest.TestCase):
def setUp(self):
self.impact_data = {
'data': {
'nodes': [
{'id': '1', 'identity': 'identity1', 'name': 'Node1', 'primaryLabel': 'Class'},
{'id': '2', 'identity': 'identity2', 'name': 'Node2', 'primaryLabel': 'Method'},
],
'relationships': [
{'startId': '1', 'endId': '2', 'type': 'CALLS'},
]
}
}
def test_extract_relationships(self):
expected_output = ["- identity1 (CALLS) -> identity2"]
result = extract_relationships(self.impact_data)
self.assertEqual(result, expected_output)
def test_extract_relationships_no_relationships(self):
self.impact_data['data']['relationships'] = []
expected_output = []
result = extract_relationships(self.impact_data)
self.assertEqual(result, expected_output)
def test_extract_relationships_missing_node(self):
self.impact_data['data']['relationships'] = [
{'startId': '1', 'endId': '3', 'type': 'CALLS'},
]
expected_output = []
result = extract_relationships(self.impact_data)
self.assertEqual(result, expected_output)
if __name__ == '__main__':
unittest.main()