0.3.0 Release Version

test/utils/test_query_executor.py

@@ -0,0 +1,186 @@
+#!/usr/bin/env python3
+"""
+Query executor tests
+"""
+
+import pytest
+from unittest.mock import Mock, AsyncMock, patch
+
+from doris_mcp_server.utils.query_executor import DorisQueryExecutor
+from doris_mcp_server.utils.config import DorisConfig
+
+
+class TestDorisQueryExecutor:
+    """Doris query executor tests"""
+
+    @pytest.fixture
+    def mock_config(self):
+        """Create mock configuration"""
+        from doris_mcp_server.utils.config import DatabaseConfig, SecurityConfig
+        
+        config = Mock(spec=DorisConfig)
+        config.doris_host = "localhost"
+        config.doris_port = 9030
+        config.doris_user = "test_user"
+        config.doris_password = "test_password"
+        config.doris_database = "test_db"
+        
+        # Add database config
+        config.database = Mock(spec=DatabaseConfig)
+        config.database.host = "localhost"
+        config.database.port = 9030
+        config.database.user = "test_user"
+        config.database.password = "test_password"
+        config.database.database = "test_db"
+        config.database.health_check_interval = 60
+        config.database.min_connections = 5
+        config.database.max_connections = 20
+        config.database.connection_timeout = 30
+        config.database.max_connection_age = 3600
+        
+        return config
+
+    @pytest.fixture
+    def query_executor(self, mock_config):
+        """Create query executor instance"""
+        # Create a mock connection manager
+        mock_connection_manager = Mock()
+        return DorisQueryExecutor(mock_connection_manager, mock_config)
+
+    @pytest.mark.asyncio
+    async def test_execute_query_success(self, query_executor):
+        """Test successful query execution using MCP interface"""
+        with patch.object(query_executor, 'execute_sql_for_mcp') as mock_execute:
+            mock_execute.return_value = {
+                "success": True,
+                "data": [
+                    {"id": 1, "name": "张三", "email": "zhangsan@example.com"},
+                    {"id": 2, "name": "李四", "email": "lisi@example.com"}
+                ],
+                "row_count": 2,
+                "execution_time": 0.15,
+                "columns": ["id", "name", "email"]
+            }
+            
+            sql = "SELECT id, name, email FROM users LIMIT 2"
+            result = await query_executor.execute_sql_for_mcp(sql)
+            
+            # Verify results
+            assert result["success"] is True
+            assert result["row_count"] == 2
+            assert len(result["data"]) == 2
+            assert result["data"][0]["id"] == 1
+            assert result["data"][0]["name"] == "张三"
+            assert result["data"][1]["email"] == "lisi@example.com"
+
+    @pytest.mark.asyncio
+    async def test_execute_query_with_parameters(self, query_executor):
+        """Test query execution with parameters"""
+        with patch.object(query_executor, 'execute_sql_for_mcp') as mock_execute:
+            mock_execute.return_value = {
+                "success": True,
+                "data": [{"id": 1, "name": "张三"}],
+                "row_count": 1,
+                "execution_time": 0.1
+            }
+            
+            sql = "SELECT id, name FROM users WHERE department = 'sales'"
+            result = await query_executor.execute_sql_for_mcp(sql)
+            
+            # Verify results
+            assert result["success"] is True
+            assert result["row_count"] == 1
+            assert len(result["data"]) == 1
+
+    @pytest.mark.asyncio
+    async def test_execute_query_connection_error(self, query_executor):
+        """Test query execution with connection error"""
+        with patch.object(query_executor, 'execute_sql_for_mcp') as mock_execute:
+            mock_execute.return_value = {
+                "success": False,
+                "error": "Connection failed",
+                "data": None
+            }
+            
+            sql = "SELECT * FROM users"
+            result = await query_executor.execute_sql_for_mcp(sql)
+            
+            assert result["success"] is False
+            assert "Connection failed" in result["error"]
+
+    @pytest.mark.asyncio
+    async def test_execute_query_sql_error(self, query_executor):
+        """Test query execution with SQL error"""
+        with patch.object(query_executor, 'execute_sql_for_mcp') as mock_execute:
+            mock_execute.return_value = {
+                "success": False,
+                "error": "SQL syntax error",
+                "data": None
+            }
+            
+            sql = "SELECT * FROM non_existent_table"
+            result = await query_executor.execute_sql_for_mcp(sql)
+            
+            assert result["success"] is False
+            assert "SQL syntax error" in result["error"]
+
+    @pytest.mark.asyncio
+    async def test_execute_query_empty_result(self, query_executor):
+        """Test query execution with empty result"""
+        with patch.object(query_executor, 'execute_sql_for_mcp') as mock_execute:
+            mock_execute.return_value = {
+                "success": True,
+                "data": [],
+                "row_count": 0,
+                "execution_time": 0.05
+            }
+            
+            sql = "SELECT * FROM users WHERE id = 999"
+            result = await query_executor.execute_sql_for_mcp(sql)
+            
+            assert result["success"] is True
+            assert result["data"] == []
+            assert result["row_count"] == 0
+
+    @pytest.mark.asyncio
+    async def test_execute_query_max_rows_limit(self, query_executor):
+        """Test query execution with max rows limit"""
+        with patch.object(query_executor, 'execute_sql_for_mcp') as mock_execute:
+            # Mock large result set limited to 100 rows
+            limited_result = [{"id": i, "name": f"user_{i}"} for i in range(100)]
+            mock_execute.return_value = {
+                "success": True,
+                "data": limited_result,
+                "row_count": 100,
+                "execution_time": 0.2
+            }
+            
+            sql = "SELECT id, name FROM users"
+            result = await query_executor.execute_sql_for_mcp(sql, limit=100)
+            
+            # Should be limited to max_rows
+            assert result["success"] is True
+            assert len(result["data"]) == 100
+
+    @pytest.mark.asyncio
+    async def test_execute_sql_for_mcp_interface(self, query_executor):
+        """Test the MCP interface method directly"""
+        with patch.object(query_executor.connection_manager, 'get_connection') as mock_get_conn:
+            # Mock connection and result
+            mock_connection = AsyncMock()
+            mock_connection.execute.return_value = Mock(
+                data=[{"id": 1, "name": "张三"}],
+                row_count=1,
+                execution_time=0.1,
+                metadata={}
+            )
+            mock_get_conn.return_value = mock_connection
+            
+            sql = "SELECT id, name FROM users LIMIT 1"
+            result = await query_executor.execute_sql_for_mcp(sql)
+            
+            # Should return success format
+            assert "success" in result
+            if result["success"]:
+                assert "data" in result
+                assert "row_count" in result