test_error_recovery.py

#!/usr/bin/env python3
"""
错误自愈测试

测试系统的错误自愈能力，包括：
- SQL 语法错误的自动修正
- 表名/字段名错误的自动修正
- 重试机制（最多 3 次）
- 错误处理和日志记录

需求: 2.2, 2.3, 2.4, 2.5
"""

import sys
import logging
from unittest.mock import Mock, patch
from nl2sql_agent import NL2SQLAgent, DatabaseConfig, setup_logging
from nl2sql_agent.llm_client import MockLLMClient


# 配置日志
setup_logging(level="INFO")
logger = logging.getLogger(__name__)


class ErrorRecoveryMockLLM(MockLLMClient):
    """
    用于错误自愈测试的 Mock LLM
    
    可以模拟生成错误的 SQL，然后在重试时生成正确的 SQL
    """
    
    def __init__(self):
        super().__init__()
        self.call_count = {}  # 记录每个查询的调用次数
    
    def generate_sql_with_error(self, query: str, error_type: str = "syntax"):
        """
        生成包含错误的 SQL
        
        参数:
            query: 用户查询
            error_type: 错误类型（syntax, table_name, column_name）
        """
        # 记录调用次数
        if query not in self.call_count:
            self.call_count[query] = 0
        self.call_count[query] += 1
        
        # 第一次调用返回错误的 SQL
        if self.call_count[query] == 1:
            if error_type == "syntax":
                # 语法错误：FORM 而不是 FROM
                return "SELECT COUNT(*) as count FORM customers"
            elif error_type == "table_name":
                # 表名错误：customer 而不是 customers
                return "SELECT COUNT(*) as count FROM customer"
            elif error_type == "column_name":
                # 字段名错误：customer_name 而不是 name
                return "SELECT customer_name FROM customers LIMIT 10"
        
        # 第二次及以后调用返回正确的 SQL
        return "SELECT COUNT(*) as count FROM customers"


def test_syntax_error_recovery():
    """
    测试 SQL 语法错误的自动修正
    
    验证：系统能够识别并修正 SQL 语法错误
    """
    logger.info("=" * 70)
    logger.info("测试 1: SQL 语法错误自动修正")
    logger.info("=" * 70)
    
    try:
        # 加载数据库配置
        db_config = DatabaseConfig.from_file("rds_demo/db_config.txt")
        
        # 创建带错误恢复的 Mock LLM
        mock_llm = ErrorRecoveryMockLLM()
        
        # 初始化 Agent
        agent = NL2SQLAgent(
            db_config=db_config,
            use_mock_llm=True
        )
        
        # 替换 LLM 客户端
        agent.llm_client = mock_llm
        
        # 模拟第一次生成错误的 SQL
        user_query = "查询客户总数"
        logger.info(f"用户查询: {user_query}")
        
        # 第一次调用：生成错误的 SQL
        error_sql = mock_llm.generate_sql_with_error(user_query, "syntax")
        logger.info(f"第 1 次生成的 SQL（错误）: {error_sql}")
        
        # 验证错误 SQL 会被拒绝
        success, result = agent.executor.execute_query(error_sql)
        assert not success, "错误的 SQL 应该执行失败"
        logger.info(f"✓ 错误 SQL 被正确识别: {result}")
        
        # 第二次调用：生成正确的 SQL
        correct_sql = mock_llm.generate_sql_with_error(user_query, "syntax")
        logger.info(f"第 2 次生成的 SQL（修正后）: {correct_sql}")
        
        # 验证正确的 SQL 可以执行
        success, result = agent.executor.execute_query(correct_sql)
        assert success, "修正后的 SQL 应该执行成功"
        logger.info(f"✓ 修正后的 SQL 执行成功: {result}")
        
        # 清理
        agent.close()
        
        logger.info("✓ 测试 1 通过\n")
        return True
        
    except Exception as e:
        logger.error(f"✗ 测试 1 失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


def test_table_name_error_recovery():
    """
    测试表名错误的自动修正
    
    验证：系统能够识别并修正表名错误
    """
    logger.info("=" * 70)
    logger.info("测试 2: 表名错误自动修正")
    logger.info("=" * 70)
    
    try:
        # 加载数据库配置
        db_config = DatabaseConfig.from_file("rds_demo/db_config.txt")
        
        # 创建带错误恢复的 Mock LLM
        mock_llm = ErrorRecoveryMockLLM()
        
        # 初始化 Agent
        agent = NL2SQLAgent(
            db_config=db_config,
            use_mock_llm=True
        )
        
        # 替换 LLM 客户端
        agent.llm_client = mock_llm
        
        user_query = "查询客户总数"
        logger.info(f"用户查询: {user_query}")
        
        # 第一次：错误的表名
        error_sql = mock_llm.generate_sql_with_error(user_query, "table_name")
        logger.info(f"第 1 次生成的 SQL（错误表名）: {error_sql}")
        
        success, result = agent.executor.execute_query(error_sql)
        assert not success, "错误的表名应该导致执行失败"
        logger.info(f"✓ 表名错误被正确识别: {result}")
        
        # 第二次：正确的表名
        correct_sql = mock_llm.generate_sql_with_error(user_query, "table_name")
        logger.info(f"第 2 次生成的 SQL（修正后）: {correct_sql}")
        
        success, result = agent.executor.execute_query(correct_sql)
        assert success, "修正后的 SQL 应该执行成功"
        logger.info(f"✓ 修正后的 SQL 执行成功")
        
        # 清理
        agent.close()
        
        logger.info("✓ 测试 2 通过\n")
        return True
        
    except Exception as e:
        logger.error(f"✗ 测试 2 失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


def test_column_name_error_recovery():
    """
    测试字段名错误的自动修正
    
    验证：系统能够识别并修正字段名错误
    """
    logger.info("=" * 70)
    logger.info("测试 3: 字段名错误自动修正")
    logger.info("=" * 70)
    
    try:
        # 加载数据库配置
        db_config = DatabaseConfig.from_file("rds_demo/db_config.txt")
        
        # 创建带错误恢复的 Mock LLM
        mock_llm = ErrorRecoveryMockLLM()
        
        # 初始化 Agent
        agent = NL2SQLAgent(
            db_config=db_config,
            use_mock_llm=True
        )
        
        # 替换 LLM 客户端
        agent.llm_client = mock_llm
        
        user_query = "查询客户姓名"
        logger.info(f"用户查询: {user_query}")
        
        # 第一次：错误的字段名
        error_sql = mock_llm.generate_sql_with_error(user_query, "column_name")
        logger.info(f"第 1 次生成的 SQL（错误字段名）: {error_sql}")
        
        success, result = agent.executor.execute_query(error_sql)
        assert not success, "错误的字段名应该导致执行失败"
        logger.info(f"✓ 字段名错误被正确识别: {result}")
        
        # 第二次：正确的字段名
        correct_sql = "SELECT name FROM customers LIMIT 10"
        logger.info(f"第 2 次生成的 SQL（修正后）: {correct_sql}")
        
        success, result = agent.executor.execute_query(correct_sql)
        assert success, "修正后的 SQL 应该执行成功"
        logger.info(f"✓ 修正后的 SQL 执行成功")
        
        # 清理
        agent.close()
        
        logger.info("✓ 测试 3 通过\n")
        return True
        
    except Exception as e:
        logger.error(f"✗ 测试 3 失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


def test_max_retry_limit():
    """
    测试最大重试次数限制
    
    验证：系统在达到最大重试次数（3次）后停止重试
    """
    logger.info("=" * 70)
    logger.info("测试 4: 最大重试次数限制")
    logger.info("=" * 70)
    
    try:
        # 加载数据库配置
        db_config = DatabaseConfig.from_file("rds_demo/db_config.txt")
        
        # 初始化 Agent
        agent = NL2SQLAgent(
            db_config=db_config,
            use_mock_llm=True
        )
        
        # 模拟一个始终失败的 SQL
        error_sql = "SELECT * FROM non_existent_table"
        logger.info(f"测试 SQL（始终失败）: {error_sql}")
        
        # 尝试执行多次
        retry_count = 0
        max_retries = 3
        
        for i in range(max_retries + 1):
            success, result = agent.executor.execute_query(error_sql)
            retry_count += 1
            
            logger.info(f"第 {retry_count} 次尝试: {'成功' if success else '失败'}")
            
            if not success:
                logger.info(f"  错误信息: {result}")
            
            # 在实际应用中，应该在达到最大重试次数后停止
            if retry_count >= max_retries:
                logger.info(f"✓ 达到最大重试次数 ({max_retries})，停止重试")
                break
        
        assert retry_count == max_retries, f"应该重试 {max_retries} 次"
        logger.info(f"✓ 重试机制工作正常")
        
        # 清理
        agent.close()
        
        logger.info("✓ 测试 4 通过\n")
        return True
        
    except Exception as e:
        logger.error(f"✗ 测试 4 失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


def test_error_logging():
    """
    测试错误日志记录
    
    验证：系统正确记录错误信息和重试过程
    """
    logger.info("=" * 70)
    logger.info("测试 5: 错误日志记录")
    logger.info("=" * 70)
    
    try:
        # 加载数据库配置
        db_config = DatabaseConfig.from_file("rds_demo/db_config.txt")
        
        # 初始化 Agent
        agent = NL2SQLAgent(
            db_config=db_config,
            use_mock_llm=True
        )
        
        # 执行一个会失败的查询
        error_sql = "SELECT * FROM invalid_table"
        logger.info(f"执行错误 SQL: {error_sql}")
        
        success, result = agent.executor.execute_query(error_sql)
        
        # 验证错误被记录
        assert not success, "查询应该失败"
        assert result is not None, "应该有错误信息"
        assert len(result) > 0, "错误信息不应为空"
        
        logger.info(f"✓ 错误信息已记录: {result}")
        
        # 清理
        agent.close()
        
        logger.info("✓ 测试 5 通过\n")
        return True
        
    except Exception as e:
        logger.error(f"✗ 测试 5 失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


def test_sql_validation():
    """
    测试 SQL 安全验证
    
    验证：系统拒绝危险的 SQL 操作
    """
    logger.info("=" * 70)
    logger.info("测试 6: SQL 安全验证")
    logger.info("=" * 70)
    
    try:
        # 加载数据库配置
        db_config = DatabaseConfig.from_file("rds_demo/db_config.txt")
        
        # 初始化 Agent
        agent = NL2SQLAgent(
            db_config=db_config,
            use_mock_llm=True
        )
        
        # 测试危险的 SQL 操作
        dangerous_sqls = [
            ("DROP TABLE customers", "DROP 操作"),
            ("DELETE FROM customers", "DELETE 操作"),
            ("UPDATE customers SET name='test'", "UPDATE 操作"),
            ("INSERT INTO customers VALUES (1, 'test')", "INSERT 操作"),
            ("SELECT * FROM customers; DROP TABLE orders;", "SQL 注入"),
        ]
        
        for sql, description in dangerous_sqls:
            logger.info(f"测试 {description}: {sql}")
            
            success, result = agent.executor.execute_query(sql)
            
            assert not success, f"{description} 应该被拒绝"
            assert "安全验证失败" in result or "不允许" in result, "应该有安全验证错误信息"
            
            logger.info(f"✓ {description} 被正确拒绝: {result}")
        
        # 清理
        agent.close()
        
        logger.info("✓ 测试 6 通过\n")
        return True
        
    except Exception as e:
        logger.error(f"✗ 测试 6 失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


def test_error_message_quality():
    """
    测试错误消息质量
    
    验证：错误消息清晰、有用，能够帮助定位问题
    """
    logger.info("=" * 70)
    logger.info("测试 7: 错误消息质量")
    logger.info("=" * 70)
    
    try:
        # 加载数据库配置
        db_config = DatabaseConfig.from_file("rds_demo/db_config.txt")
        
        # 初始化 Agent
        agent = NL2SQLAgent(
            db_config=db_config,
            use_mock_llm=True
        )
        
        # 测试不同类型的错误
        test_cases = [
            ("SELECT * FROM non_existent_table", "表不存在"),
            ("SELECT invalid_column FROM customers", "字段不存在"),
            ("SELECT * FORM customers", "语法错误"),
        ]
        
        for sql, error_type in test_cases:
            logger.info(f"测试 {error_type}: {sql}")
            
            success, result = agent.executor.execute_query(sql)
            
            assert not success, f"{error_type} 应该导致失败"
            assert result is not None, "应该有错误消息"
            assert len(result) > 10, "错误消息应该足够详细"
            
            logger.info(f"✓ {error_type} 的错误消息: {result[:100]}...")
        
        # 清理
        agent.close()
        
        logger.info("✓ 测试 7 通过\n")
        return True
        
    except Exception as e:
        logger.error(f"✗ 测试 7 失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


def main():
    """
    运行所有错误自愈测试
    """
    logger.info("\n" + "=" * 70)
    logger.info("错误自愈测试")
    logger.info("=" * 70 + "\n")
    
    # 测试列表
    tests = [
        ("SQL 语法错误自动修正", test_syntax_error_recovery),
        ("表名错误自动修正", test_table_name_error_recovery),
        ("字段名错误自动修正", test_column_name_error_recovery),
        ("最大重试次数限制", test_max_retry_limit),
        ("错误日志记录", test_error_logging),
        ("SQL 安全验证", test_sql_validation),
        ("错误消息质量", test_error_message_quality),
    ]
    
    # 运行测试
    results = []
    for test_name, test_func in tests:
        try:
            result = test_func()
            results.append((test_name, result))
        except Exception as e:
            logger.error(f"测试 '{test_name}' 发生异常: {str(e)}")
            results.append((test_name, False))
    
    # 输出测试结果摘要
    logger.info("=" * 70)
    logger.info("测试结果摘要")
    logger.info("=" * 70)
    
    passed = sum(1 for _, result in results if result)
    total = len(results)
    
    for test_name, result in results:
        status = "✓ 通过" if result else "✗ 失败"
        logger.info(f"{status}: {test_name}")
    
    logger.info("=" * 70)
    logger.info(f"总计: {passed}/{total} 测试通过")
    logger.info("=" * 70)
    
    return passed == total


if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)