"""
Tests for db_client module.

These tests assume a StarRocks cluster is running on localhost with default configurations:
- Host: localhost
- Port: 9030 (MySQL protocol)
- User: root
- Password: (empty)
- No default database set

Run tests with: pytest tests/test_db_client.py -v
"""

import os
import pytest
import pandas as pd
from unittest.mock import patch, MagicMock

# Set up test environment variables
os.environ.pop('STARROCKS_FE_ARROW_FLIGHT_SQL_PORT', None)  # Force MySQL mode for tests
os.environ.pop('STARROCKS_DB', None)  # No default database

from src.mcp_server_starrocks.db_client import (
    DBClient, 
    ResultSet, 
    get_db_client, 
    reset_db_connections,
    parse_connection_url
)


class TestDBClient:
    """Test cases for DBClient class."""
    
    @pytest.fixture
    def db_client(self):
        """Create a fresh DBClient instance for each test."""
        # Reset global state
        reset_db_connections()
        return DBClient()
    
    def test_client_initialization(self, db_client):
        """Test DBClient initialization with default settings."""
        assert db_client.enable_arrow_flight_sql is False
        assert db_client.default_database is None
        assert db_client._connection_pool is None
        assert db_client._adbc_connection is None
    
    def test_singleton_pattern(self):
        """Test that get_db_client returns the same instance."""
        client1 = get_db_client()
        client2 = get_db_client()
        assert client1 is client2
    
    def test_execute_show_databases(self, db_client):
        """Test executing SHOW DATABASES query."""
        result = db_client.execute("SHOW DATABASES")
        
        assert isinstance(result, ResultSet)
        assert result.success is True
        assert result.column_names is not None
        assert len(result.column_names) == 1
        assert result.rows is not None
        assert len(result.rows) > 0
        assert result.execution_time is not None
        assert result.execution_time > 0
        
        # Check that information_schema is present (standard in StarRocks)
        database_names = [row[0] for row in result.rows]
        assert 'information_schema' in database_names
    
    def test_execute_show_databases_pandas(self, db_client):
        """Test executing SHOW DATABASES with pandas return format."""
        result = db_client.execute("SHOW DATABASES", return_format="pandas")
        
        assert isinstance(result, ResultSet)
        assert result.success is True
        assert result.pandas is not None
        assert isinstance(result.pandas, pd.DataFrame)
        assert len(result.pandas.columns) == 1
        assert len(result.pandas) > 0
        
        # Test that to_pandas() returns the same DataFrame
        df = result.to_pandas()
        assert df is result.pandas
    
    def test_execute_invalid_query(self, db_client):
        """Test executing an invalid SQL query."""
        result = db_client.execute("SELECT * FROM nonexistent_table_12345")
        
        assert isinstance(result, ResultSet)
        assert result.success is False
        assert result.error_message is not None
        assert "nonexistent_table_12345" in result.error_message or "doesn't exist" in result.error_message.lower()
        assert result.execution_time is not None
    
    def test_execute_create_and_drop_database(self, db_client):
        """Test creating and dropping a test database."""
        test_db_name = "test_mcp_db_client"
        
        # Clean up first (in case previous test failed)
        db_client.execute(f"DROP DATABASE IF EXISTS {test_db_name}")
        
        # Create database
        create_result = db_client.execute(f"CREATE DATABASE {test_db_name}")
        assert create_result.success is True
        assert create_result.rows_affected is not None  # DDL returns row count (usually 0)
        
        # Verify database exists
        show_result = db_client.execute("SHOW DATABASES")
        database_names = [row[0] for row in show_result.rows]
        assert test_db_name in database_names
        
        # Drop database
        drop_result = db_client.execute(f"DROP DATABASE {test_db_name}")
        assert drop_result.success is True
        
        # Verify database is gone
        show_result = db_client.execute("SHOW DATABASES")
        database_names = [row[0] for row in show_result.rows]
        assert test_db_name not in database_names
    
    def test_execute_with_specific_database(self, db_client):
        """Test executing query with specific database context."""
        # Use information_schema which should always be available
        result = db_client.execute("SHOW TABLES", db="information_schema")
        
        assert result.success is True
        assert result.column_names is not None
        assert result.rows is not None
        assert len(result.rows) > 0  # information_schema should have tables
        
        # Check for expected information_schema tables
        table_names = [row[0] for row in result.rows]
        expected_tables = ['tables', 'columns', 'schemata']
        found_expected = any(table in table_names for table in expected_tables)
        assert found_expected, f"Expected at least one of {expected_tables} in {table_names}"
    
    def test_execute_with_invalid_database(self, db_client):
        """Test executing query with non-existent database."""
        result = db_client.execute("SHOW TABLES", db="nonexistent_db_12345")
        
        assert result.success is False
        assert result.error_message is not None
        assert "nonexistent_db_12345" in result.error_message
    
    def test_execute_table_operations(self, db_client):
        """Test creating, inserting, querying, and dropping a table."""
        test_db = "test_mcp_table_ops"
        test_table = "test_table"
        
        try:
            # Create database
            create_db_result = db_client.execute(f"CREATE DATABASE IF NOT EXISTS {test_db}")
            assert create_db_result.success is True
            
            # Create table (with replication_num=1 for single-node setup)
            create_table_sql = f"""
            CREATE TABLE {test_db}.{test_table} (
                id INT,
                name STRING,
                value DOUBLE
            )
            PROPERTIES ("replication_num" = "1")
            """
            create_result = db_client.execute(create_table_sql)
            assert create_result.success is True
            
            # Insert data
            insert_sql = f"""
            INSERT INTO {test_db}.{test_table} VALUES 
            (1, 'test1', 1.5),
            (2, 'test2', 2.5),
            (3, 'test3', 3.5)
            """
            insert_result = db_client.execute(insert_sql)
            assert insert_result.success is True
            assert insert_result.rows_affected == 3
            
            # Query data
            select_result = db_client.execute(f"SELECT * FROM {test_db}.{test_table} ORDER BY id")
            assert select_result.success is True
            assert len(select_result.column_names) == 3
            assert select_result.column_names == ['id', 'name', 'value']
            assert len(select_result.rows) == 3
            # MySQL connector returns tuples, convert to lists for comparison
            assert list(select_result.rows[0]) == [1, 'test1', 1.5]
            assert list(select_result.rows[1]) == [2, 'test2', 2.5]
            assert list(select_result.rows[2]) == [3, 'test3', 3.5]
            
            # Test COUNT query
            count_result = db_client.execute(f"SELECT COUNT(*) as cnt FROM {test_db}.{test_table}")
            assert count_result.success is True
            assert count_result.rows[0][0] == 3
            
            # Test with specific database context
            ctx_result = db_client.execute(f"SELECT * FROM {test_table}", db=test_db)
            assert ctx_result.success is True
            assert len(ctx_result.rows) == 3
            
        finally:
            # Clean up
            db_client.execute(f"DROP DATABASE IF EXISTS {test_db}")
    
    def test_execute_pandas_format_with_data(self, db_client):
        """Test pandas format with actual data."""
        test_db = "test_mcp_pandas"
        
        try:
            # Setup test data
            db_client.execute(f"CREATE DATABASE IF NOT EXISTS {test_db}")
            db_client.execute(f"""
                CREATE TABLE {test_db}.pandas_test (
                    id INT,
                    category STRING,
                    amount DECIMAL(10,2)
                )
                PROPERTIES ("replication_num" = "1")
            """)
            db_client.execute(f"""
                INSERT INTO {test_db}.pandas_test VALUES 
                (1, 'A', 100.50),
                (2, 'B', 200.75),
                (3, 'A', 150.25)
            """)
            
            # Test executing query with pandas format
            result = db_client.execute(f"SELECT * FROM {test_db}.pandas_test ORDER BY id", return_format="pandas")
            
            assert isinstance(result, ResultSet)
            assert result.success is True
            assert result.pandas is not None
            assert isinstance(result.pandas, pd.DataFrame)
            assert len(result.pandas) == 3
            assert list(result.pandas.columns) == ['id', 'category', 'amount']
            assert result.pandas.iloc[0]['id'] == 1
            assert result.pandas.iloc[0]['category'] == 'A'
            assert float(result.pandas.iloc[0]['amount']) == 100.50
            
            # Test that to_pandas() returns the same DataFrame
            df = result.to_pandas()
            assert df is result.pandas
        
        finally:
            db_client.execute(f"DROP DATABASE IF EXISTS {test_db}")
    
    def test_connection_error_handling(self, db_client):
        """Test error handling when connection fails."""
        # Mock a connection failure
        with patch.object(db_client, '_get_connection', side_effect=Exception("Connection failed")):
            result = db_client.execute("SHOW DATABASES")
            
            assert result.success is False
            assert "Connection failed" in result.error_message
            assert result.execution_time is not None
    
    def test_reset_connections(self, db_client):
        """Test connection reset functionality."""
        # First execute a query to establish connection
        result1 = db_client.execute("SHOW DATABASES")
        assert result1.success is True
        
        # Reset connections
        db_client.reset_connections()
        
        # Should still work after reset
        result2 = db_client.execute("SHOW DATABASES")
        assert result2.success is True
    
    def test_describe_table(self, db_client):
        """Test DESCRIBE table functionality."""
        test_db = "test_mcp_describe"
        test_table = "describe_test"
        
        try:
            # Create test table
            db_result = db_client.execute(f"CREATE DATABASE IF NOT EXISTS {test_db}")
            assert db_result.success, f"Failed to create database: {db_result.error_message}"
            
            table_result = db_client.execute(f"""
                CREATE TABLE {test_db}.{test_table} (
                    id BIGINT NOT NULL COMMENT 'Primary key',
                    name VARCHAR(100) COMMENT 'Name field',
                    created_at DATETIME,
                    is_active BOOLEAN
                )
                PROPERTIES ("replication_num" = "1")
            """)
            assert table_result.success, f"Failed to create table: {table_result.error_message}"
            
            # Verify table exists first
            show_result = db_client.execute(f"SHOW TABLES", db=test_db)
            assert show_result.success, f"Failed to show tables: {show_result.error_message}"
            table_names = [row[0] for row in show_result.rows]
            assert test_table in table_names, f"Table {test_table} not found in {table_names}"
            
            # Describe table (use full table name for clarity)
            result = db_client.execute(f"DESCRIBE {test_db}.{test_table}")
            
            assert result.success is True
            assert result.column_names is not None
            assert len(result.rows) == 4  # 4 columns
            
            # Check column names in result (should include Field, Type, etc.)
            expected_columns = ['Field', 'Type', 'Null', 'Key', 'Default', 'Extra']
            for expected_col in expected_columns[:len(result.column_names)]:
                assert expected_col in result.column_names
            
            # Check that our table columns are present
            field_names = [row[0] for row in result.rows]
            assert 'id' in field_names
            assert 'name' in field_names
            assert 'created_at' in field_names
            assert 'is_active' in field_names
        
        finally:
            db_client.execute(f"DROP DATABASE IF EXISTS {test_db}")


class TestDBClientWithArrowFlight:
    """Test cases for DBClient with Arrow Flight SQL (if configured)."""
    
    @pytest.fixture
    def arrow_client(self):
        """Create DBClient with Arrow Flight SQL if available."""
        # Check if Arrow Flight SQL port is configured (either from env or default test port)
        arrow_port = os.getenv('STARROCKS_FE_ARROW_FLIGHT_SQL_PORT', '9408')
        
        # Test if Arrow Flight SQL is actually available by trying to connect
        try:
            with patch.dict(os.environ, {'STARROCKS_FE_ARROW_FLIGHT_SQL_PORT': arrow_port}):
                reset_db_connections()
                client = DBClient()
                assert client.enable_arrow_flight_sql is True
                
                # Test basic connectivity
                result = client.execute("SHOW DATABASES")
                if not result.success:
                    pytest.skip(f"Arrow Flight SQL not available on port {arrow_port}: {result.error_message}")
                
                return client
        except Exception as e:
            pytest.skip(f"Arrow Flight SQL not available: {e}")
    
    def test_arrow_flight_basic_query(self, arrow_client):
        """Test basic query with Arrow Flight SQL."""
        result = arrow_client.execute("SHOW DATABASES")
        
        assert isinstance(result, ResultSet)
        assert result.success is True
        assert result.column_names is not None
        assert result.rows is not None
        assert len(result.rows) > 0
        
        # Verify we're actually using Arrow Flight SQL
        assert arrow_client.enable_arrow_flight_sql is True
    
    def test_arrow_flight_pandas_format(self, arrow_client):
        """Test pandas format with Arrow Flight SQL."""
        result = arrow_client.execute("SHOW DATABASES", return_format="pandas")
        
        assert isinstance(result, ResultSet)
        assert result.success is True
        assert result.pandas is not None
        assert isinstance(result.pandas, pd.DataFrame)
        assert len(result.pandas) > 0
        assert len(result.pandas.columns) == 1
        
        # Test that to_pandas() returns the same DataFrame
        df = result.to_pandas()
        assert df is result.pandas
        
        # Verify we're actually using Arrow Flight SQL
        assert arrow_client.enable_arrow_flight_sql is True
    
    def test_arrow_flight_table_operations(self, arrow_client):
        """Test table operations with Arrow Flight SQL."""
        test_db = "test_arrow_flight"
        test_table = "arrow_test"
        
        try:
            # Create database
            create_db_result = arrow_client.execute(f"CREATE DATABASE IF NOT EXISTS {test_db}")
            assert create_db_result.success is True
            
            # Create table
            create_table_sql = f"""
            CREATE TABLE {test_db}.{test_table} (
                id INT,
                name STRING,
                value DOUBLE
            )
            PROPERTIES ("replication_num" = "1")
            """
            create_result = arrow_client.execute(create_table_sql)
            assert create_result.success is True
            
            # Insert data
            insert_sql = f"""
            INSERT INTO {test_db}.{test_table} VALUES 
            (1, 'arrow1', 1.1),
            (2, 'arrow2', 2.2)
            """
            insert_result = arrow_client.execute(insert_sql)
            assert insert_result.success is True
            # Note: StarRocks Arrow Flight SQL always returns 0 for rows_affected due to implementation limitations
            assert insert_result.rows_affected == 0
            
            # Query data with pandas format
            select_result = arrow_client.execute(f"SELECT * FROM {test_db}.{test_table} ORDER BY id", return_format="pandas")
            assert isinstance(select_result, ResultSet)
            assert select_result.success is True
            assert select_result.pandas is not None
            assert isinstance(select_result.pandas, pd.DataFrame)
            assert len(select_result.pandas) == 2
            # Note: StarRocks Arrow Flight SQL loses column names in SELECT results (known limitation)
            # The columns come back as empty strings, but the data is correct
            assert len(select_result.pandas.columns) == 3
            # Since column names are empty, access by position instead
            assert select_result.pandas.iloc[0, 0] == 1    # id column
            assert select_result.pandas.iloc[0, 1] == 'arrow1'  # name column  
            assert select_result.pandas.iloc[0, 2] == 1.1  # value column
            
            # Test that to_pandas() returns the same DataFrame
            df = select_result.to_pandas()
            assert df is select_result.pandas
            
            # Query data with raw format
            raw_result = arrow_client.execute(f"SELECT * FROM {test_db}.{test_table} ORDER BY id")
            assert raw_result.success is True
            assert len(raw_result.rows) == 2
            # Note: Column names are empty due to StarRocks Arrow Flight SQL limitation
            assert raw_result.column_names == ['', '', '']
            # But the data is correct
            assert raw_result.rows[0] == [1, 'arrow1', 1.1]
            assert raw_result.rows[1] == [2, 'arrow2', 2.2]
            
        finally:
            # Clean up
            arrow_client.execute(f"DROP DATABASE IF EXISTS {test_db}")
    
    def test_arrow_flight_error_handling(self, arrow_client):
        """Test error handling with Arrow Flight SQL."""
        # Test invalid query
        result = arrow_client.execute("SELECT * FROM nonexistent_arrow_table")
        assert result.success is False
        assert result.error_message is not None
        
        # Test invalid database - Note: Arrow Flight SQL may fail with connection errors
        # before database validation, so we just check that it fails
        result = arrow_client.execute("SHOW TABLES", db="nonexistent_arrow_db")
        assert result.success is False
        assert result.error_message is not None


class TestResultSet:
    """Test cases for ResultSet dataclass."""
    
    def test_result_set_creation(self):
        """Test ResultSet creation with various parameters."""
        # Success case
        result = ResultSet(
            success=True,
            column_names=['id', 'name'],
            rows=[[1, 'test'], [2, 'test2']],
            execution_time=0.5
        )
        
        assert result.success is True
        assert result.column_names == ['id', 'name']
        assert result.rows == [[1, 'test'], [2, 'test2']]
        assert result.execution_time == 0.5
        assert result.rows_affected is None
        assert result.error_message is None
    
    def test_result_set_to_pandas_from_rows(self):
        """Test ResultSet to_pandas conversion from rows."""
        result = ResultSet(
            success=True,
            column_names=['id', 'name', 'value'],
            rows=[[1, 'test1', 10.5], [2, 'test2', 20.5]],
            execution_time=0.1
        )
        
        df = result.to_pandas()
        assert isinstance(df, pd.DataFrame)
        assert len(df) == 2
        assert list(df.columns) == ['id', 'name', 'value']
        assert df.iloc[0]['id'] == 1
        assert df.iloc[0]['name'] == 'test1'
        assert df.iloc[0]['value'] == 10.5
        assert df.iloc[1]['id'] == 2
        assert df.iloc[1]['name'] == 'test2'
        assert df.iloc[1]['value'] == 20.5
    
    def test_result_set_to_pandas_from_pandas_field(self):
        """Test ResultSet to_pandas returns existing pandas field if available."""
        original_df = pd.DataFrame({
            'id': [1, 2],
            'name': ['test1', 'test2'],
            'value': [10.5, 20.5]
        })
        
        result = ResultSet(
            success=True,
            column_names=['id', 'name', 'value'],
            rows=[[1, 'test1', 10.5], [2, 'test2', 20.5]],
            pandas=original_df,
            execution_time=0.1
        )
        
        df = result.to_pandas()
        assert df is original_df  # Should return the same object
    
    def test_result_set_to_string(self):
        """Test ResultSet to_string conversion."""
        result = ResultSet(
            success=True,
            column_names=['id', 'name', 'value'],
            rows=[[1, 'test1', 10.5], [2, 'test2', 20.5]],
            execution_time=0.1
        )
        
        string_output = result.to_string()
        expected_lines = [
            'id,name,value',
            '1,test1,10.5',
            '2,test2,20.5',
            ''
        ]
        assert string_output == '\n'.join(expected_lines)
    
    def test_result_set_to_string_with_limit(self):
        """Test ResultSet to_string with limit."""
        result = ResultSet(
            success=True,
            column_names=['id', 'name'],
            rows=[[1, 'very_long_test_string'], [2, 'another_long_string']],
            execution_time=0.1
        )
        
        # Test with very small limit
        string_output = result.to_string(limit=20)
        lines = string_output.split('\n')
        assert lines[0] == 'id,name'  # Header should always be included
        # Should stop before all rows due to limit
        assert len(lines) < 4  # Should be less than header + 2 rows + empty line
    
    def test_result_set_to_string_error_cases(self):
        """Test ResultSet to_string error handling."""
        # Test with failed result
        failed_result = ResultSet(
            success=False,
            error_message="Test error"
        )
        
        string_output = failed_result.to_string()
        assert string_output == "Error: Test error"
        
        # Test with no data
        no_data_result = ResultSet(
            success=True,
            column_names=None,
            rows=None
        )
        
        string_output = no_data_result.to_string()
        assert string_output == "No data"
    
    def test_result_set_to_pandas_error_cases(self):
        """Test ResultSet to_pandas error handling."""
        # Test with failed result
        failed_result = ResultSet(
            success=False,
            error_message="Test error"
        )
        
        with pytest.raises(ValueError, match="Cannot convert failed result to DataFrame"):
            failed_result.to_pandas()
        
        # Test with no data
        no_data_result = ResultSet(
            success=True,
            column_names=None,
            rows=None
        )
        
        with pytest.raises(ValueError, match="No data available to convert to DataFrame"):
            no_data_result.to_pandas()
    
    def test_result_set_error_case(self):
        """Test ResultSet for error cases."""
        result = ResultSet(
            success=False,
            error_message="Test error",
            execution_time=0.1
        )
        
        assert result.success is False
        assert result.error_message == "Test error"
        assert result.execution_time == 0.1
        assert result.column_names is None
        assert result.rows is None
        assert result.rows_affected is None
    
    def test_result_set_write_operation(self):
        """Test ResultSet for write operations."""
        result = ResultSet(
            success=True,
            rows_affected=5,
            execution_time=0.2
        )
        
        assert result.success is True
        assert result.rows_affected == 5
        assert result.execution_time == 0.2
        assert result.column_names is None
        assert result.rows is None
        assert result.error_message is None


class TestParseConnectionUrl:
    """Test cases for parse_connection_url function."""
    
    def test_parse_basic_url(self):
        """Test parsing basic connection URL without schema."""
        url = "root:password123@localhost:9030/test_db"
        result = parse_connection_url(url)

        expected = {
            'user': 'root',
            'password': 'password123',
            'host': 'localhost',
            'port': '9030',
            'database': 'test_db'
        }
        assert result == expected
    
    def test_parse_url_with_schema(self):
        """Test parsing connection URL with schema."""
        url = "mysql://admin:secret@db.example.com:3306/production"
        result = parse_connection_url(url)

        expected = {
            'user': 'admin',
            'password': 'secret',
            'host': 'db.example.com',
            'port': '3306',
            'database': 'production'
        }
        assert result == expected
    
    def test_parse_url_with_different_schemas(self):
        """Test parsing URLs with various schema types."""
        test_cases = [
            ("starrocks://user:pass@host:9030/db", "starrocks"),
            ("jdbc+mysql://user:pass@host:3306/db", "jdbc+mysql"),
            ("postgresql://user:pass@host:5432/db", "postgresql"),
        ]

        for url, expected_schema in test_cases:
            result = parse_connection_url(url)
            # Schema is no longer returned in the result
            assert result['user'] == 'user'
            assert result['password'] == 'pass'
            assert result['host'] == 'host'
            assert result['database'] == 'db'
    
    def test_parse_url_empty_password_succeeds(self):
        """Test that URL with empty password now works."""
        url = "root:@localhost:9030/test_db"
        result = parse_connection_url(url)

        expected = {
            'user': 'root',
            'password': '',  # Empty password
            'host': 'localhost',
            'port': '9030',
            'database': 'test_db'
        }
        assert result == expected
    
    def test_parse_url_no_password_colon(self):
        """Test URL without password colon (e.g., root@localhost:9030)."""
        url = "root@localhost:9030"
        result = parse_connection_url(url)
        
        expected = {
            'user': 'root',
            'password': '',  # Default empty password
            'host': 'localhost',
            'port': '9030',
            'database': None
        }
        assert result == expected
    
    def test_parse_url_missing_port_uses_default(self):
        """Test URL without port uses default 9030."""
        url = "root:password@localhost/mydb"
        result = parse_connection_url(url)
        
        expected = {
            'user': 'root',
            'password': 'password',
            'host': 'localhost',
            'port': '9030',  # Default port
            'database': 'mydb'
        }
        assert result == expected
    
    def test_parse_url_minimal_format(self):
        """Test minimal URL format (just user@host)."""
        url = "user@host"
        result = parse_connection_url(url)
        
        expected = {
            'user': 'user',
            'password': '',  # Default empty password
            'host': 'host',
            'port': '9030',  # Default port
            'database': None
        }
        assert result == expected
    
    def test_parse_url_empty_string_password(self):
        """Test URL with explicit empty password using double colon."""
        url = "user::@host:9030/db"
        result = parse_connection_url(url)
        
        expected = {
            'user': 'user',
            'password': ':',  # Literal colon as password
            'host': 'host',
            'port': '9030',
            'database': 'db'
        }
        assert result == expected
    
    def test_parse_url_complex_password_limitation(self):
        """Test that password with @ symbol has regex limitation (parses incorrectly)."""
        url = "user:p@ssw0rd!@server:9030/mydb"
        result = parse_connection_url(url)
        
        # Due to regex limitation, @ in password causes incorrect parsing
        assert result['user'] == 'user'
        assert result['password'] == 'p'  # Only gets characters before first @
        assert result['host'] == 'ssw0rd!@server'  # Rest becomes host
        assert result['port'] == '9030'
        assert result['database'] == 'mydb'
    
    def test_parse_url_password_without_at_symbol(self):
        """Test parsing URL with complex password without @ symbol."""
        url = "user:p#ssw0rd!$%^&*()@server:9030/mydb"
        result = parse_connection_url(url)
        
        assert result['user'] == 'user'
        assert result['password'] == 'p#ssw0rd!$%^&*()'
        assert result['host'] == 'server'
        assert result['port'] == '9030'
        assert result['database'] == 'mydb'
    
    def test_parse_url_complex_username_with_at_symbol_limitation(self):
        """Test that username with @ symbol fails (regex limitation)."""
        url = "user.name+tag@domain:password123@host:9030/db"
        # This should fail because our regex cannot distinguish between 
        # the @ in username vs the @ separator for host
        with pytest.raises(ValueError, match="Invalid connection URL"):
            parse_connection_url(url)
    
    def test_parse_url_complex_username_without_at(self):
        """Test parsing URL with complex username without @ symbol."""
        url = "user.name+tag_domain:password123@host:9030/db"
        result = parse_connection_url(url)
        
        assert result['user'] == 'user.name+tag_domain'
        assert result['password'] == 'password123'
        assert result['host'] == 'host'
        assert result['port'] == '9030'
        assert result['database'] == 'db'
    
    def test_parse_url_numeric_database(self):
        """Test parsing URL with numeric database name."""
        url = "root:pass@localhost:9030/db123"
        result = parse_connection_url(url)
        
        assert result['database'] == 'db123'
    
    def test_parse_url_database_with_hyphens(self):
        """Test parsing URL with database name containing hyphens."""
        url = "root:pass@localhost:9030/test-db-name"
        result = parse_connection_url(url)
        
        assert result['database'] == 'test-db-name'
    
    def test_parse_url_ip_address_host(self):
        """Test parsing URL with IP address as host."""
        url = "root:pass@192.168.1.100:9030/testdb"
        result = parse_connection_url(url)
        
        assert result['host'] == '192.168.1.100'
        assert result['port'] == '9030'
        assert result['database'] == 'testdb'
    
    def test_parse_url_different_ports(self):
        """Test parsing URLs with different port numbers."""
        test_cases = [
            ("user:pass@host:3306/db", "3306"),
            ("user:pass@host:5432/db", "5432"),
            ("user:pass@host:27017/db", "27017"),
            ("user:pass@host:1/db", "1"),
            ("user:pass@host:65535/db", "65535"),
        ]
        
        for url, expected_port in test_cases:
            result = parse_connection_url(url)
            assert result['port'] == expected_port
    
    def test_parse_invalid_urls(self):
        """Test that invalid URLs raise ValueError."""
        invalid_urls = [
            # Missing required parts
            "@host:9030/db",  # Missing user
            "user:pass@:9030/db",  # Missing host
            
            # Malformed URLs
            "user:pass@host:port/db",  # Non-numeric port
            "user:pass@host:9030/",  # Empty database
            "user:pass@host:9030/db/extra",  # Extra path component
            "",  # Empty string
            "random-string-not-url",  # Not a URL format
            
            # Special cases
            "://user:pass@host:9030/db",  # Empty schema
            "user:pass@host:-1/db",  # Negative port
        ]
        
        for invalid_url in invalid_urls:
            with pytest.raises(ValueError, match="Invalid connection URL"):
                parse_connection_url(invalid_url)
    
    def test_parse_url_colon_in_password_works(self):
        """Test that colon in password actually works (unlike @ symbol)."""
        url = "user:pass:extra@host:9030/db"
        result = parse_connection_url(url)
        
        assert result['user'] == 'user'
        assert result['password'] == 'pass:extra'  # Colons in password are fine
        assert result['host'] == 'host'
        assert result['port'] == '9030'
        assert result['database'] == 'db'
    
    def test_parse_url_without_database(self):
        """Test parsing URL without database (database is optional)."""
        url = "user:password@host:9030"
        result = parse_connection_url(url)
        
        assert result['user'] == 'user'
        assert result['password'] == 'password'
        assert result['host'] == 'host'
        assert result['port'] == '9030'
        assert result['database'] == None  # Database should be None when omitted
    
    def test_parse_url_with_schema_without_database(self):
        """Test parsing URL with schema but without database."""
        url = "mysql://admin:secret@db.example.com:3306"
        result = parse_connection_url(url)
        
        assert result['user'] == 'admin'
        assert result['password'] == 'secret'
        assert result['host'] == 'db.example.com'
        assert result['port'] == '3306'
        assert result['database'] == None
    
    def test_parse_url_various_schemas_without_database(self):
        """Test parsing URLs with various schemas but no database."""
        test_cases = [
            ("starrocks://user:pass@host:9030", "starrocks"),
            ("jdbc+mysql://user:pass@host:3306", "jdbc+mysql"),
            ("postgresql://user:pass@host:5432", "postgresql"),
        ]
        
        for url, expected_schema in test_cases:
            result = parse_connection_url(url)
            # Schema is no longer returned in the result
            assert result['user'] == 'user'
            assert result['password'] == 'pass'
            assert result['host'] == 'host'
            assert result['database'] == None
    
    def test_parse_url_edge_cases(self):
        """Test edge cases that should work."""
        # Single character components
        url = "a:b@c:1/d"
        result = parse_connection_url(url)
        assert result['user'] == 'a'
        assert result['password'] == 'b'
        assert result['host'] == 'c'
        assert result['port'] == '1'
        assert result['database'] == 'd'
        
        # Long components
        long_user = "a" * 100
        long_pass = "b" * 100
        long_host = "c" * 50
        long_db = "d" * 50
        url = f"{long_user}:{long_pass}@{long_host}:9030/{long_db}"
        result = parse_connection_url(url)
        assert result['user'] == long_user
        assert result['password'] == long_pass
        assert result['host'] == long_host
        assert result['database'] == long_db
    
    def test_parse_url_returns_dict_with_all_keys(self):
        """Test that parse_connection_url always returns dict with all expected keys."""
        test_cases = [
            "root:pass@localhost:9030/db",
            "mysql://root:pass@localhost:3306/db",
        ]
        
        expected_keys = {'user', 'password', 'host', 'port', 'database'}
        
        for url in test_cases:
            result = parse_connection_url(url)
            assert isinstance(result, dict)
            assert set(result.keys()) == expected_keys
    
    def test_parse_url_regex_pattern_comprehensive(self):
        """Test comprehensive regex pattern matching."""
        # Test that the regex correctly captures each group
        url = "custom+schema://test_user:complex!pass@sub.domain.com:12345/my_db-name"
        result = parse_connection_url(url)
        
        # Schema is no longer returned in the result
        assert result['user'] == 'test_user'
        assert result['password'] == 'complex!pass'
        assert result['host'] == 'sub.domain.com'
        assert result['port'] == '12345'
        assert result['database'] == 'my_db-name'


class TestDummyMode:
    """Test cases for STARROCKS_DUMMY_TEST environment variable."""
    
    def test_dummy_mode_enabled(self):
        """Test that dummy mode returns expected dummy data."""
        # Set dummy test environment variable
        with patch.dict(os.environ, {'STARROCKS_DUMMY_TEST': '1'}):
            client = DBClient()
            assert client.enable_dummy_test is True
            
            # Test basic query
            result = client.execute("SELECT * FROM any_table")
            
            assert result.success is True
            assert result.column_names == ['name']
            assert result.rows == [['aaa'], ['bbb'], ['ccc']]
            assert result.execution_time is not None
            assert result.execution_time > 0
            assert result.pandas is None  # pandas should be None for raw format
    
    def test_dummy_mode_with_pandas_format(self):
        """Test dummy mode with pandas return format."""
        with patch.dict(os.environ, {'STARROCKS_DUMMY_TEST': '1'}):
            client = DBClient()
            
            result = client.execute("SELECT * FROM any_table", return_format="pandas")
            
            assert result.success is True
            assert result.column_names == ['name']
            assert result.rows == [['aaa'], ['bbb'], ['ccc']]
            assert result.pandas is not None
            assert isinstance(result.pandas, pd.DataFrame)
            assert len(result.pandas) == 3
            assert list(result.pandas.columns) == ['name']
            assert result.pandas.iloc[0]['name'] == 'aaa'
            assert result.pandas.iloc[1]['name'] == 'bbb'
            assert result.pandas.iloc[2]['name'] == 'ccc'
    
    def test_dummy_mode_ignores_statement_and_db(self):
        """Test that dummy mode returns same data regardless of SQL statement or database."""
        with patch.dict(os.environ, {'STARROCKS_DUMMY_TEST': '1'}):
            client = DBClient()
            
            # Test different statements
            result1 = client.execute("SHOW DATABASES")
            result2 = client.execute("CREATE TABLE test (id INT)")
            result3 = client.execute("SELECT COUNT(*) FROM users", db="production")
            
            # All should return the same dummy data
            for result in [result1, result2, result3]:
                assert result.success is True
                assert result.column_names == ['name']
                assert result.rows == [['aaa'], ['bbb'], ['ccc']]
    
    def test_dummy_mode_disabled_by_default(self):
        """Test that dummy mode is disabled when environment variable is not set."""
        # Ensure STARROCKS_DUMMY_TEST is not set
        with patch.dict(os.environ, {}, clear=True):
            os.environ.pop('STARROCKS_DUMMY_TEST', None)  # Remove if exists
            client = DBClient()
            assert client.enable_dummy_test is False
    
    def test_dummy_mode_with_empty_string(self):
        """Test that empty string for STARROCKS_DUMMY_TEST disables dummy mode."""
        with patch.dict(os.environ, {'STARROCKS_DUMMY_TEST': ''}):
            client = DBClient()
            assert client.enable_dummy_test is False
    
    def test_dummy_mode_with_various_truthy_values(self):
        """Test that various truthy values enable dummy mode."""
        truthy_values = ['1', 'true', 'True', 'yes', 'on', 'any_non_empty_string']
        
        for value in truthy_values:
            with patch.dict(os.environ, {'STARROCKS_DUMMY_TEST': value}):
                client = DBClient()
                assert client.enable_dummy_test is True, f"Failed for value: {value}"
    
    def test_dummy_mode_to_pandas_conversion(self):
        """Test to_pandas() method works with dummy data."""
        with patch.dict(os.environ, {'STARROCKS_DUMMY_TEST': '1'}):
            client = DBClient()
            
            # Test raw format conversion
            result = client.execute("SELECT * FROM test")
            df = result.to_pandas()
            assert isinstance(df, pd.DataFrame)
            assert len(df) == 3
            assert list(df.columns) == ['name']
            assert df.iloc[0]['name'] == 'aaa'
            
            # Test pandas format (should return same DataFrame)
            result_pandas = client.execute("SELECT * FROM test", return_format="pandas")
            df_pandas = result_pandas.to_pandas()
            assert df_pandas is result_pandas.pandas
    
    def test_dummy_mode_to_string_conversion(self):
        """Test to_string() method works with dummy data."""
        with patch.dict(os.environ, {'STARROCKS_DUMMY_TEST': '1'}):
            client = DBClient()
            
            result = client.execute("SELECT * FROM test")
            string_output = result.to_string()
            
            expected_lines = [
                'name',
                'aaa',
                'bbb', 
                'ccc',
                ''
            ]
            assert string_output == '\n'.join(expected_lines)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])