MCP KQL Server

""" Unit tests for the MCP server module. Author: Arjun Trivedi Email: arjuntrivedi42@yahoo.com """ import unittest import json from unittest.mock import patch from mcp_kql_server.constants import TEST_CONFIG class TestMCPServerFunctions(unittest.TestCase): """Test cases for MCP server functions and operations.""" def setUp(self): """Set up test fixtures.""" self.valid_query = f"cluster('{TEST_CONFIG['mock_cluster_uri']}').database('{TEST_CONFIG['mock_database']}').{TEST_CONFIG['mock_table']} | take 10" self.test_cluster_uri = TEST_CONFIG["mock_cluster_uri"] self.test_database = TEST_CONFIG["mock_database"] self.test_table = TEST_CONFIG["mock_table"] def test_error_handler_enhancement(self): """Test enhanced ErrorHandler functionality.""" from mcp_kql_server.utils import ErrorHandler # Test basic error handling test_error = Exception("Test error") result = ErrorHandler.handle_kusto_error(test_error) self.assertIsInstance(result, dict) self.assertIn("success", result) self.assertIn("error", result) self.assertIn("suggestions", result) self.assertIn("recovery_actions", result) self.assertFalse(result["success"]) def test_schema_manager_integration(self): """Test SchemaManager integration.""" from mcp_kql_server.utils import SchemaManager from mcp_kql_server.memory import get_memory_manager memory_manager = get_memory_manager() schema_manager = SchemaManager(memory_manager) self.assertIsNotNone(schema_manager) self.assertIsNotNone(schema_manager.memory_manager) def test_bracket_if_needed_function(self): """Test the bracket_if_needed utility function.""" from mcp_kql_server.utils import bracket_if_needed # Test normal column name (should not be bracketed) normal_column = "RegularColumn" self.assertEqual(bracket_if_needed(normal_column), normal_column) # Test column with spaces (should be bracketed) spaced_column = "Column With Spaces" result = bracket_if_needed(spaced_column) self.assertTrue(result.startswith("['") and result.endswith("']")) # Test reserved word reserved_word = "table" # This is a KQL reserved word result = bracket_if_needed(reserved_word) self.assertTrue(result.startswith("['") and result.endswith("']")) def test_configuration_constants(self): """Test that required configuration constants are properly defined.""" from mcp_kql_server.constants import ( SERVER_NAME, SERVER_DESCRIPTION, CONNECTION_CONFIG, ERROR_HANDLING_CONFIG, KQL_RESERVED_WORDS ) self.assertIsInstance(SERVER_NAME, str) self.assertIsInstance(SERVER_DESCRIPTION, str) self.assertIsInstance(CONNECTION_CONFIG, dict) self.assertIsInstance(ERROR_HANDLING_CONFIG, dict) self.assertIsInstance(KQL_RESERVED_WORDS, (list, frozenset)) self.assertGreater(len(KQL_RESERVED_WORDS), 0) def test_json_serialization_helper(self): """Test the JSON serialization helper from ErrorHandler.""" from mcp_kql_server.utils import ErrorHandler # Test with simple data simple_data = {"key": "value", "number": 42} result = ErrorHandler.safe_json_dumps(simple_data) self.assertIsInstance(result, str) parsed = json.loads(result) self.assertEqual(parsed["key"], "value") self.assertEqual(parsed["number"], 42) # Test with complex data that might cause issues complex_data = {"datetime": "2023-01-01T00:00:00", "none_value": None} result = ErrorHandler.safe_json_dumps(complex_data) self.assertIsInstance(result, str) def test_memory_manager_availability(self): """Test that memory manager is properly available.""" from mcp_kql_server.memory import get_memory_manager manager = get_memory_manager() self.assertIsNotNone(manager) # Test basic memory manager functionality self.assertTrue(hasattr(manager, 'get_memory_stats')) # Note: clear_schema_cache is available via SchemaManager, not directly on memory manager self.assertTrue(hasattr(manager, 'corpus') or hasattr(manager, 'memory_path')) @patch("mcp_kql_server.mcp_server.kql_execute_tool") @patch("mcp_kql_server.mcp_server.schema_manager") @patch("mcp_kql_server.mcp_server.kusto_manager_global", {"authenticated": True}) def test_execute_kql_query_schema_context_on_error(self, mock_schema_manager, mock_execute_tool): """Test that schema context is added to suggestions on SEM0100 error.""" from mcp_kql_server.mcp_server import _execute_kql_query_logic import asyncio # Mock execution to raise SEM0100 error mock_execute_tool.side_effect = Exception("Request is invalid and cannot be processed: Semantic error: SEM0100: 'project' operator: Failed to resolve scalar expression named 'InvalidCol'") # Mock schema manager to return schema async def mock_get_schema(*args, **kwargs): return { "columns": {"ValidCol1": "string", "ValidCol2": "int"}, "last_updated": "2023-01-01" } mock_schema_manager.get_table_schema.side_effect = mock_get_schema # Execute query result_json = asyncio.run(_execute_kql_query_logic( query="MyTable | project InvalidCol", cluster_url="https://test.kusto.windows.net", database="TestDB" )) result = json.loads(result_json) # Verify results self.assertFalse(result["success"]) self.assertIn("suggestions", result) # Check if schema context is in suggestions suggestions_str = str(result["suggestions"]) self.assertIn("Schema Context for referenced tables", suggestions_str) self.assertIn("ValidCol1", suggestions_str) self.assertIn("ValidCol2", suggestions_str) if __name__ == "__main__": unittest.main()