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_query_processor_integration(self): """Test that QueryProcessor is properly integrated.""" from mcp_kql_server.utils import QueryProcessor from mcp_kql_server.memory import get_memory_manager # Test that QueryProcessor can be instantiated memory_manager = get_memory_manager() processor = QueryProcessor(memory_manager) # Test basic functionality self.assertIsNotNone(processor) self.assertIsNotNone(processor.memory_manager) 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"]) @patch('mcp_kql_server.mcp_server.kusto_manager_global', {'authenticated': True}) def test_natural_language_query_processing(self): """Test natural language query processing logic.""" from mcp_kql_server.mcp_server import _generate_kql_from_natural_language # This would normally be an async function, but we're testing the logic # In a real test, you'd use asyncio.run() or similar self.assertTrue(callable(_generate_kql_from_natural_language)) 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_query_validation_logic(self): """Test query validation and preprocessing.""" from mcp_kql_server.utils import QueryProcessor from mcp_kql_server.memory import get_memory_manager processor = QueryProcessor(get_memory_manager()) # Test query cleaning dirty_query = " // Comment\n MyTable | take 10 " clean_query = processor.clean(dirty_query) self.assertIsInstance(clean_query, str) self.assertNotIn("//", clean_query) # Test query parsing test_query = "MyTable | take 10 | project Column1, Column2" parsed = processor.parse(test_query) self.assertIsInstance(parsed, dict) self.assertIn("tables", parsed) self.assertIn("operations", parsed) 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')) if __name__ == "__main__": unittest.main()